castor-extractor 0.19.0__py3-none-any.whl → 0.19.6__py3-none-any.whl

castor_extractor/utils/client/api/auth_test.py

@@ -0,0 +1,49 @@
+from typing import Dict, Optional
+
+from .auth import BasicAuth, BearerAuth, CustomAuth
+
+
+class _MockRequest:
+    def __init__(self):
+        self.headers = {}
+
+
+class _CustomAuth(CustomAuth):
+    def _authentication_header(self) -> Dict[str, str]:
+        return {"custom-token": "token"}
+
+
+class _BearAuth(BearerAuth):
+    def fetch_token(self) -> Optional[str]:
+        return "token"
+
+
+def test_BasicAuth():
+    prepared_request = _MockRequest()
+    auth = BasicAuth(username="simple", password="basic")
+    auth.__call__(prepared_request)
+    assert prepared_request.headers == {
+        "Authorization": "Basic c2ltcGxlOmJhc2lj"
+    }
+
+
+def test_BearerAuth():
+    prepared_request = _MockRequest()
+    auth = _BearAuth()
+    auth.__call__(prepared_request)
+    assert prepared_request.headers == {"Authorization": "Bearer token"}
+
+    auth._token = "expired_token"
+    auth.__call__(prepared_request)
+    assert prepared_request.headers == {"Authorization": "Bearer expired_token"}
+
+    auth.refresh_token()
+    auth.__call__(prepared_request)
+    assert prepared_request.headers == {"Authorization": "Bearer token"}
+
+
+def test_CustomAuth():
+    prepared_request = _MockRequest()
+    auth = _CustomAuth()
+    auth.__call__(prepared_request)
+    assert prepared_request.headers == {"custom-token": "token"}

castor_extractor/utils/client/api/client.py

@@ -0,0 +1,153 @@
+import logging
+from http import HTTPStatus
+from typing import Dict, Literal, Optional, Tuple
+
+import requests
+from requests import Response
+
+from ...retry import retry_request
+from .auth import Auth
+from .safe_request import RequestSafeMode, handle_response
+from .utils import build_url
+
+logger = logging.getLogger(__name__)
+
+Headers = Optional[Dict[str, str]]
+
+# https://requests.readthedocs.io/en/latest/api/#requests.request
+HttpMethod = Literal["GET", "OPTIONS", "HEAD", "POST", "PUT", "PATCH", "DELETE"]
+
+DEFAULT_TIMEOUT = 60
+RETRY_ON_EXPIRED_TOKEN = 1
+
+
+def _generate_payloads(
+    method: HttpMethod,
+    params: Optional[dict],
+    data: Optional[dict],
+    pagination_params: Optional[dict],
+) -> Tuple[dict, dict]:
+    _pagination_params = pagination_params or {}
+    params = params or {}
+    data = data or {}
+
+    if method == "GET":
+        params = {**params, **_pagination_params}
+    elif method == "POST":
+        data = {**data, **_pagination_params}
+    else:
+        raise ValueError(f"Method {method} is not yet supported")
+    return data, params
+
+
+class APIClient:
+    """
+    Interface to easily query REST-API with GET and POST requests
+
+    Args:
+        auth: auth class to enable logging to the API
+        host: base url of the API
+        headers: common headers to all calls that will be made
+        timeout: read timeout for each request
+        safe_mode: ignore certain exceptions based on status codes
+
+    Note:
+        If the auth implements a refreshing mechanism (refresh_token)
+        the token is automatically refreshed once upon receiving the
+        401: UNAUTHORIZED status code
+    """
+
+    def __init__(
+        self,
+        auth: Auth,
+        host: Optional[str] = None,
+        headers: Headers = None,
+        timeout: int = DEFAULT_TIMEOUT,
+        safe_mode: RequestSafeMode = RequestSafeMode(),
+    ):
+        self.base_headers = headers or {}
+        self._host = host
+        self._timeout = timeout
+        self._auth = auth
+        self._safe_mode = safe_mode
+
+    def _call(
+        self,
+        method: HttpMethod,
+        endpoint: str,
+        *,
+        headers: Headers = None,
+        params: Optional[dict] = None,
+        data: Optional[dict] = None,
+        pagination_params: Optional[dict] = None,
+    ) -> Response:
+        headers = headers or {}
+
+        data, params = _generate_payloads(
+            method=method,
+            params=params,
+            data=data,
+            pagination_params=pagination_params,
+        )
+
+        url = build_url(self._host, endpoint)
+
+        return requests.request(
+            method=method,
+            url=url,
+            auth=self._auth,
+            headers={**self.base_headers, **headers},
+            params=params,
+            json=data,
+            timeout=self._timeout,
+        )
+
+    @retry_request(
+        status_codes=(HTTPStatus.UNAUTHORIZED,),
+        max_retries=RETRY_ON_EXPIRED_TOKEN,
+    )
+    def _get(
+        self,
+        endpoint: str,
+        *,
+        headers: Headers = None,
+        params: Optional[dict] = None,
+        data: Optional[dict] = None,
+        pagination_params: Optional[dict] = None,
+    ):
+        response = self._call(
+            method="GET",
+            endpoint=endpoint,
+            params=params,
+            data=data,
+            pagination_params=pagination_params,
+            headers=headers,
+        )
+        if response.status_code == HTTPStatus.UNAUTHORIZED:
+            self._auth.refresh_token()
+
+        return handle_response(response, safe_mode=self._safe_mode)
+
+    @retry_request(
+        status_codes=(HTTPStatus.UNAUTHORIZED,),
+        max_retries=RETRY_ON_EXPIRED_TOKEN,
+    )
+    def _post(
+        self,
+        endpoint: str,
+        *,
+        headers: Headers = None,
+        data: Optional[dict] = None,
+        pagination_params: Optional[dict] = None,
+    ):
+        response = self._call(
+            method="POST",
+            endpoint=endpoint,
+            data=data,
+            pagination_params=pagination_params,
+            headers=headers,
+        )
+        if response.status_code == HTTPStatus.UNAUTHORIZED:
+            self._auth.refresh_token()
+
+        return handle_response(response, safe_mode=self._safe_mode)

castor_extractor/utils/client/api/client_test.py

@@ -0,0 +1,47 @@
+from unittest.mock import patch
+
+from requests import PreparedRequest, Request, Session
+
+from .auth import BasicAuth
+from .client import APIClient
+
+
+class MockSession:
+    def __init__(self):
+        self.request_data = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    def request(self, **kwargs) -> PreparedRequest:
+        kwargs.pop("timeout")
+        request = Request(**kwargs)
+        prepared_request = Session().prepare_request(request=request)
+        return prepared_request
+
+
+@patch("requests.sessions.Session", MockSession)
+def test__get():
+    auth = BasicAuth(username="user_id", password="secret")
+    client = APIClient(
+        auth=auth,
+        host="https://example.api.com/v1/",
+        headers={"content-type": "test"},
+        timeout=9,
+    )
+
+    prepared_request = client._call("GET", "endpoint")
+
+    # test method
+    assert prepared_request.method == "GET"
+
+    # test headers
+    assert prepared_request.headers["content-type"] == "test"
+    assert (
+        prepared_request.headers["Authorization"]
+        == "Basic dXNlcl9pZDpzZWNyZXQ="
+    )
+    assert prepared_request.url == "https://example.api.com/v1/endpoint"

castor_extractor/utils/client/api/pagination.py

@@ -0,0 +1,83 @@
+import logging
+from abc import abstractmethod
+from enum import Enum
+from functools import partial
+from time import sleep
+from typing import Callable, Iterator, Optional, Type, Union
+
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class FetchNextPageBy(Enum):
+    """
+    Enum to pick which APIClient._call() argument we want
+    to use for calling the next page in the pagination.
+    Supported arguments are :
+    - params (PAYLOAD)
+    - endpoint (URL)
+    """
+
+    PAYLOAD = "pagination_params"
+    URL = "endpoint"
+
+
+class PaginationModel(BaseModel):
+    """
+    Base abstract class defining a pagination model
+
+    By implementing the 3 abstract methods below, enables
+    to fetch all elements of a Paginated API by using the
+    `fetch_all_pages` method
+    """
+
+    fetch_by: FetchNextPageBy = FetchNextPageBy.PAYLOAD
+    current_page_payload: Optional["dict"] = None
+
+    @abstractmethod
+    def is_last(self) -> bool:
+        """Stopping condition for the pagination"""
+        pass
+
+    @abstractmethod
+    def next_page_payload(self) -> Optional[Union[dict, str]]:
+        """Payload enabling to generate the request for the next page"""
+        pass
+
+    @abstractmethod
+    def page_results(self) -> list:
+        """List of results of the current page"""
+        pass
+
+    def next_page_parameters(self) -> dict:
+        return {self.fetch_by.value: self.next_page_payload()}
+
+
+def fetch_all_pages(
+    request: Callable,
+    pagination_model: Type[PaginationModel],
+    rate_limit: Optional[int] = None,
+) -> Iterator:
+    """
+    Method to return all results of a Paginated API based on the
+    pagination model and the first request call
+    """
+    page_number = 1
+    response_payload = request()
+    paginated_response = pagination_model(**response_payload)
+    while not paginated_response.is_last():
+        logger.info(f"Fetching page number {page_number}")
+        yield from paginated_response.page_results()
+        next_page_parameters = paginated_response.next_page_parameters()
+        new_request = partial(request, **next_page_parameters)
+        if rate_limit:
+            sleep(rate_limit)
+        paginated_response = pagination_model(
+            current_page_payload=next_page_parameters, **new_request()
+        )
+        page_number += 1
+
+    # send last page's results
+    logger.info(f"Fetching page number {page_number}")
+    yield from paginated_response.page_results()

castor_extractor/utils/client/api/pagination_test.py

@@ -0,0 +1,51 @@
+from typing import Optional
+
+from pydantic import ConfigDict, Field
+from pydantic.alias_generators import to_camel
+
+from .pagination import PaginationModel, fetch_all_pages
+
+
+class _TestPagination(PaginationModel):
+    next_page: Optional[str] = None
+    entries: list = Field(default_factory=list)
+
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )
+
+    def is_last(self) -> bool:
+        return self.next_page is None
+
+    def next_page_payload(self) -> dict:
+        return {"page": self.next_page}
+
+    def page_results(self) -> list:
+        return self.entries
+
+
+def _request(pagination_params: Optional[dict] = None):
+    if not pagination_params:
+        return {
+            "nextPage": "next_page_id",
+            "entries": [1, 2, 3, 4, 5],
+        }
+    if pagination_params.get("page") == "next_page_id":
+        return {
+            "nextPage": "next_page_id_2",
+            "entries": [6, 7, 8, 9, 10],
+        }
+    if pagination_params.get("page") == "next_page_id_2":
+        return {
+            "nextPage": None,
+            "entries": [11],
+        }
+
+    raise AssertionError(f"call has unexpected parameters: {pagination_params}")
+
+
+def test__TestPagination():
+    all_results = fetch_all_pages(_request, _TestPagination)
+    assert list(all_results) == [i for i in range(1, 12)]

castor_extractor/utils/{safe_request_test.py → client/api/safe_request_test.py}

@@ -4,7 +4,10 @@ from http import HTTPStatus
 import pytest
 from requests import HTTPError, Response
 
-from .safe_request import RequestSafeMode, handle_response
+from .safe_request import (
+    RequestSafeMode,
+    handle_response,
+)
 
 
 def mock_response(status_code: int):

castor_extractor/utils/client/api/utils.py

@@ -0,0 +1,9 @@
+from typing import Optional
+
+
+def build_url(host: Optional[str], endpoint: str):
+    if not host:
+        return endpoint
+    if not host.startswith("https://"):
+        host = "https://" + host
+    return f"{host.strip('/')}/{endpoint}"

castor_extractor/utils/client/api/utils_test.py

@@ -0,0 +1,16 @@
+from .utils import build_url
+
+
+def test_APIClient_build_url():
+    expected = "https://3.14.azuredatabricks.net/api/2.1/unity-catalog/tables"
+
+    path = "api/2.1/unity-catalog/tables"
+
+    host = "3.14.azuredatabricks.net"
+    assert expected == build_url(host, path)
+
+    host_with_http = "https://3.14.azuredatabricks.net"
+    assert expected == build_url(host_with_http, path)
+
+    host_with_trailing_slash = "https://3.14.azuredatabricks.net/"
+    assert expected == build_url(host_with_trailing_slash, path)

castor_extractor/utils/collection.py

@@ -1,12 +1,23 @@
-from typing import Any, DefaultDict, Dict, List, Sequence
+from collections import defaultdict
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Sequence,
+    Set,
+    TypeVar,
+)
 
 from .object import getproperty
 from .type import Getter
 
+T = TypeVar("T")
+
 
 def group_by(identifier: Getter, elements: Sequence) -> Dict[Any, List]:
     """Groups the elements by the given key"""
-    groups: Dict[Any, List] = DefaultDict(list)
+    groups: Dict[Any, List] = defaultdict(list)
     for element in elements:
         key = getproperty(element, identifier)
         groups[key].append(element)
@@ -52,3 +63,24 @@ def empty_iterator():
     Remark: missing return type is on purpose, it breaks the typing
     """
     return iter([])
+
+
+def deduplicate(
+    identifier: Getter,
+    elements: Iterable[T],
+) -> List[T]:
+    """
+    Remove duplicates in the given elements, using the specified identifier
+    Only the first occurrence is kept.
+    """
+    deduplicated: List[T] = []
+    processed: Set[Any] = set()
+
+    for element in elements:
+        key = getproperty(element, identifier)
+        if key in processed:
+            continue
+        processed.add(key)
+        deduplicated.append(element)
+
+    return deduplicated

castor_extractor/utils/collection_test.py

@@ -1,6 +1,4 @@
-import pytest
-
-from .collection import mapping_from_rows
+from .collection import deduplicate, mapping_from_rows
 
 
 def test__mapping_from_rows__basic_mapping():
@@ -58,3 +56,19 @@ def test__mapping_from_rows__multiple_valid_rows():
     result = mapping_from_rows(rows, "id", "name")
     expected = {1: "Alice", 2: "Bob", 3: "Charlie"}
     assert result == expected
+
+
+def test_deduplicate():
+    e1 = {"id": "1", "name": "element_1"}
+    e2 = {"id": "2", "name": "element_2"}
+    e3 = {"id": "3", "name": "element_3"}
+
+    elements = [
+        e1,
+        e2,
+        e3,
+        {"id": "3", "name": "duplicate"},
+        {"id": "3", "name": "duplicate"},
+        {"id": "2", "name": "duplicate"},
+    ]
+    assert deduplicate("id", elements) == [e1, e2, e3]

castor_extractor/utils/pager/__init__.py

@@ -1,3 +1,2 @@
 from .pager import AbstractPager, Pager, PagerLogger, PagerStopStrategy
 from .pager_on_id import PagerOnId, PagerOnIdLogger
-from .pager_on_token import PagerOnToken

castor_extractor/utils/retry.py

@@ -6,6 +6,7 @@ from typing import Any, Callable, Sequence, Tuple, Type, Union
 
 from pydantic import BaseModel, PositiveInt, PrivateAttr
 from pydantic.fields import Field
+from requests import HTTPError
 
 logger = logging.getLogger(__name__)
 
@@ -115,3 +116,46 @@ def retry(
         return _func
 
     return _wrapper
+
+
+def retry_request(
+    status_codes: Sequence[int],
+    max_retries: int = 1,
+    base_ms: int = 10,
+    jitter_ms: int = 20,
+    strategy: RetryStrategy = DEFAULT_STRATEGY,
+    log_exc_info: bool = False,
+) -> Callable:
+    """retry decorator"""
+
+    exceptions_ = tuple(e for e in status_codes)
+
+    def _wrapper(callable: Callable) -> Callable:
+        def _try(*args, **kwargs) -> WrapperReturnType:
+            try:
+                return None, callable(*args, **kwargs)
+            except HTTPError as err:
+                status_code = err.response.status_code
+                if status_code not in exceptions_:
+                    raise err
+                logger.warning(f"Exception within {callable.__name__}")
+                return err, None
+
+        def _func(*args, **kwargs) -> Any:
+            retry = Retry(
+                max_retries=max_retries,
+                base_ms=base_ms,
+                jitter_ms=jitter_ms,
+                strategy=strategy,
+            )
+            while True:
+                err, result = _try(*args, **kwargs)
+                if err is None:
+                    return result
+                if retry.check(err, log_exc_info):
+                    continue
+                raise err
+
+        return _func
+
+    return _wrapper

castor_extractor/utils/retry_test.py

@@ -1,11 +1,15 @@
+from http import HTTPStatus
 from statistics import variance
 from time import time
 from typing import List
+from unittest.mock import patch
 
 import pytest
+import requests
 from pydantic.error_wrappers import ValidationError
+from requests import HTTPError, Response
 
-from .retry import MS_IN_SEC, Retry, RetryStrategy
+from .retry import MS_IN_SEC, Retry, RetryStrategy, retry_request
 
 
 def test_retry_field_validations():
@@ -75,3 +79,24 @@ def test_retry_strategy__check():
     assert retry.check(error) is False
     delta_ms = int((after - before) * MS_IN_SEC)
     assert _within(delta_ms, 315, 345)
+
+
+@patch("requests.get")
+def test_retry_request(mocked_get):
+    def error_response():
+        response = Response()
+        response.status_code = HTTPStatus.UNAUTHORIZED
+        return response
+
+    mocked_get.return_value = error_response()
+
+    @retry_request(status_codes=(HTTPStatus.UNAUTHORIZED,), max_retries=3)
+    def get():
+        response = requests.get("hello")
+        response.raise_for_status()
+        return response.json()
+
+    with pytest.raises(HTTPError):
+        get()
+
+    assert mocked_get.call_count == 4  # 1 call + 3 retries