lakehouse_engine.utils.rest_api

  1"""Module to handle REST API operations."""
  2
  3import time
  4from enum import Enum
  5
  6import requests
  7from requests.adapters import HTTPAdapter
  8from urllib3.util.retry import Retry
  9
 10from lakehouse_engine.utils.logging_handler import LoggingHandler
 11
 12LOG = LoggingHandler(__name__).get_logger()
 13DEFAULT_CONTENT_TYPE = "application/json"
 14
 15
 16class RestMethods(Enum):
 17    """Methods for REST API calls."""
 18
 19    POST = "POST"
 20    PUT = "PUT"
 21    ALLOWED_METHODS = ["POST", "PUT"]
 22
 23
 24class RestStatusCodes(Enum):
 25    """REST Status Code."""
 26
 27    RETRY_STATUS_CODES = [429, 500, 502, 503, 504]
 28    OK_STATUS_CODES = [200]
 29
 30
 31class RESTApiException(requests.RequestException):
 32    """Class representing any possible REST API Exception."""
 33
 34    def __init__(self, message: str) -> None:
 35        """Construct RESTApiException instances.
 36
 37        Args:
 38            message: message to display on exception event.
 39        """
 40        super().__init__(message)
 41
 42
 43def get_basic_auth(username: str, password: str) -> requests.auth.HTTPBasicAuth:
 44    """Get the basic authentication object to authenticate REST requests.
 45
 46    Args:
 47        username: username.
 48        password: password.
 49
 50    Returns:
 51        requests.auth.HTTPBasicAuth: the HTTPBasicAuth object.
 52    """
 53    return requests.auth.HTTPBasicAuth(username, password)
 54
 55
 56def get_configured_session(
 57    sleep_seconds: float = 0.2,
 58    total_retries: int = 5,
 59    backoff_factor: int = 2,
 60    retry_status_codes: list = None,
 61    allowed_methods: list = None,
 62    protocol: str = "https://",
 63) -> requests.Session:
 64    """Get a configured requests Session with exponential backoff.
 65
 66    Args:
 67        sleep_seconds: seconds to sleep before each request to avoid rate limits.
 68        total_retries: number of times to retry.
 69        backoff_factor: factor for the exponential backoff.
 70        retry_status_codes: list of status code that triggers a retry.
 71        allowed_methods: http methods that are allowed for retry.
 72        protocol: http:// or https://.
 73
 74    Returns
 75        requests.Session: the configured session.
 76    """
 77    retry_status_codes = (
 78        retry_status_codes
 79        if retry_status_codes
 80        else RestStatusCodes.RETRY_STATUS_CODES.value
 81    )
 82    allowed_methods = (
 83        allowed_methods if allowed_methods else RestMethods.ALLOWED_METHODS.value
 84    )
 85    time.sleep(sleep_seconds)
 86    session = requests.Session()
 87    retries = Retry(
 88        total=total_retries,
 89        backoff_factor=backoff_factor,
 90        status_forcelist=retry_status_codes,
 91        allowed_methods=allowed_methods,
 92    )
 93    session.mount(protocol, HTTPAdapter(max_retries=retries))
 94    return session
 95
 96
 97def execute_api_request(
 98    method: str,
 99    url: str,
100    headers: dict = None,
101    basic_auth_dict: dict = None,
102    json: dict = None,
103    files: dict = None,
104    sleep_seconds: float = 0.2,
105) -> requests.Response:
106    """Execute a REST API request.
107
108    Args:
109        method: REST method (e.g., POST or PUT).
110        url: url of the api.
111        headers: request headers.
112        basic_auth_dict: basic http authentication details
113            (e.g., {"username": "x", "password": "y"}).
114        json: json payload to send in the request.
115        files: files payload to send in the request.
116        sleep_seconds: for how many seconds to sleep to avoid error 429.
117
118    Returns:
119        response from the HTTP request.
120    """
121    basic_auth: requests.auth.HTTPBasicAuth = None
122    if basic_auth_dict:
123        basic_auth = get_basic_auth(
124            basic_auth_dict["username"], basic_auth_dict["password"]
125        )
126
127    return get_configured_session(sleep_seconds=sleep_seconds).request(
128        method=method,
129        url=url,
130        headers=headers,
131        auth=basic_auth,
132        json=json,
133        files=files,
134    )