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

CHANGELOG.md

@@ -1,14 +1,41 @@
+
 # Changelog
 
+## 0.19.6 - 2024-09-03
+
+* Metabase: Adding error handler on API call
+
+## 0.19.5 - 2024-09-02
+
+* Databricks/Salesforce: Remove deprecated client dependencies
+
+## 0.19.4 - 2024-08-29
+
+* Tableau Pulse: extract Metrics and Subscriptions
+
+## 0.19.3 - 2024-08-27
+
+- Sigma: Add SafeMode to SigmaClient
+
+## 0.19.2 - 2024-08-23
+
+- Reworked APIClient to unify client's behaviours
+
+Impacted Technologies: Sigma, Soda, Notion
+
+## 0.19.1 - 2024-08-23
+
+* Domo: extract datasources via cards metadata endpoint
+
 ## 0.19.0 - 2024-08-21
 
-* Breaking change Looker CLI: 
+* Breaking change Looker CLI:
 
 `-u` and `--username` changed to `-c` `--client-id`
 
 `-p` and `--password` changed to `-s` `--client-secret`
 
-* Breaking change Metabase CLI: 
+* Breaking change Metabase CLI:
 
 `-u` and `--username` changed to `-u` `--user`
 

castor_extractor/file_checker/templates/generic_warehouse.py

@@ -1,7 +1,7 @@
 from typing import Set
 
 from ..column import ColumnChecker
-from ..constants import COLUMN_TYPES, TABLE_TYPES
+from ..constants import TABLE_TYPES
 from ..enums import DataType
 from ..file import FileTemplate
 

castor_extractor/knowledge/notion/client/client.py

@@ -1,24 +1,13 @@
-import logging
+from functools import partial
 from http import HTTPStatus
-from typing import Dict, Iterator
+from typing import Dict, Iterator, Optional
 
-from pydantic import ValidationError
-
-from ....utils import RequestSafeMode, empty_iterator
-from ....utils.client.api import APIClient, HttpMethod
+from ....utils import APIClient, BearerAuth, RequestSafeMode, fetch_all_pages
 from ..assets import NotionAsset
-from .constants import (
-    CASTOR_NOTION_USER_AGENT,
-    NOTION_BASE_URL,
-    NOTION_TIMEOUT_MS,
-)
+from .constants import CASTOR_NOTION_USER_AGENT, NOTION_BASE_URL, NOTION_VERSION
 from .credentials import NotionCredentials
-from .endpoints import EndpointFactory
-from .pagination import PaginatedResponse
-
-logger = logging.getLogger("__name__")
-
-NOTION_VERSION = "2021-08-16"
+from .endpoints import NotionEndpointFactory
+from .pagination import NotionPagination
 
 VOLUME_IGNORED = 10
 IGNORED_ERROR_CODES = (HTTPStatus.BAD_GATEWAY,)
@@ -26,83 +15,64 @@ NOTION_SAFE_MODE = RequestSafeMode(
     max_errors=VOLUME_IGNORED,
     status_codes=IGNORED_ERROR_CODES,
 )
+NOTION_BASE_HEADERS = {
+    "Notion-Version": NOTION_VERSION,
+    "User-Agent": CASTOR_NOTION_USER_AGENT,
+}
 
 
 def _search_filter(asset: str) -> Dict[str, Dict[str, str]]:
     return {"filter": {"value": asset, "property": "object"}}
 
 
-class NotionClient(APIClient):
-    """Client fetching data from Notion"""
+class NotionAuth(BearerAuth):
+    def __init__(self, token: str):
+        self.token = token
 
+    def fetch_token(self):
+        return self.token
+
+
+class NotionClient(APIClient):
     def __init__(
         self,
         credentials: NotionCredentials,
-        timeout: int = NOTION_TIMEOUT_MS,
-        base_url: str = NOTION_BASE_URL,
-        user_agent: str = CASTOR_NOTION_USER_AGENT,
-        safe_mode: RequestSafeMode = NOTION_SAFE_MODE,
-    ) -> None:
-        base_headers = {
-            "Notion-Version": NOTION_VERSION,
-            "User-Agent": user_agent,
-        }
-
+        safe_mode: Optional[RequestSafeMode] = None,
+    ):
+        auth = NotionAuth(token=credentials.token)
         super().__init__(
-            host=base_url,
-            token=credentials.token,
-            headers=base_headers,
-            timeout=timeout,
-            safe_mode=safe_mode,
+            host=NOTION_BASE_URL,
+            auth=auth,
+            headers=NOTION_BASE_HEADERS,
+            safe_mode=safe_mode or NOTION_SAFE_MODE,
         )
 
-    def _request(
-        self, method: HttpMethod, endpoint: str, params: dict
-    ) -> Iterator[dict]:
-        """
-        API call to Notion:
-        - If result is paginated, yield all pages
-        - If not, yield only the response payload
-        """
-        built_url = self.build_url(self._host, endpoint)
-        response_payload = self._call(
-            method=method,
-            url=built_url,
-            params=params if params and method == "GET" else None,
-            data=params if params and method == "POST" else None,
-        )
-        try:
-            paginated_response = PaginatedResponse(**response_payload)
-            yield from paginated_response.results
-
-            if not paginated_response.has_more:
-                return empty_iterator()
-
-            yield from self._request(
-                method=method,
-                endpoint=endpoint,
-                params=paginated_response.start_cursor,
-            )
-
-        except ValidationError:
-            yield response_payload
+    def users(self) -> Iterator[dict]:
+        request = partial(self._get, endpoint=NotionEndpointFactory.users())
+        yield from fetch_all_pages(request, NotionPagination)
 
     def _page_listing(self) -> Iterator[dict]:
-        return self._request(
-            method="POST",
-            endpoint=EndpointFactory.search(),
-            params=_search_filter("page"),
+        request = partial(
+            self._post,
+            endpoint=NotionEndpointFactory.search(),
+            data=_search_filter("page"),
         )
+        yield from fetch_all_pages(request, NotionPagination)
 
     def _blocks(self, block_id: str) -> Iterator[dict]:
-        return self._request("GET", EndpointFactory.blocks(block_id), {})
+        request = partial(
+            self._get,
+            endpoint=NotionEndpointFactory.blocks(block_id),
+        )
+        yield from fetch_all_pages(request, NotionPagination)
 
     def databases(self) -> Iterator[dict]:
-        return self._request(
-            method="POST",
-            endpoint=EndpointFactory.search(),
-            params=_search_filter("database"),
+        request = partial(
+            self._post,
+            endpoint=NotionEndpointFactory.search(),
+            data=_search_filter("database"),
         )
+        yield from fetch_all_pages(request, NotionPagination)
 
     def recursive_blocks(self, block_id: str) -> Iterator[dict]:
         """Fetch recursively all children blocks of a given block or page"""
@@ -116,9 +86,7 @@ class NotionClient(APIClient):
 
     def pages(self) -> Iterator[dict]:
         """Fetch all pages with its whole content"""
-        pages = list(self._page_listing())
-        logger.info(f"Extracting {len(pages)} pages ...")
-        for page in pages:
+        for page in self._page_listing():
             if page.get("object") == "database":
                 # Notion Search API filter for page doesn't work
                 continue
@@ -126,10 +94,6 @@ class NotionClient(APIClient):
             page["child_blocks"] = content
             yield page
 
-    def users(self) -> Iterator[dict]:
-        """Fetch all users"""
-        return self._request("GET", EndpointFactory.users(), {})
-
     def fetch(self, asset: NotionAsset) -> Iterator[dict]:
         """Returns the needed metadata for the queried asset"""
         if asset == NotionAsset.PAGES:

castor_extractor/knowledge/notion/client/client_test.py

@@ -1,7 +1,10 @@
+from functools import partial
 from unittest.mock import patch
 
+from ....utils import fetch_all_pages
 from .client import NotionClient
 from .credentials import NotionCredentials
+from .pagination import NotionPagination
 
 MOCK_PAGINATED_RESPONSE = {
     "results": [{"result_id": 1}],
@@ -16,7 +19,7 @@ MOCK_PAGINATED_RESPONSE_2 = {
 MOCK_RESPONSE = {"result_id": 3}
 
 
-@patch.object(NotionClient, "_call")
+@patch.object(NotionClient, "_get")
 def test_NotionClient__request(mock_call):
     mock_call.side_effect = [
         MOCK_PAGINATED_RESPONSE,
@@ -25,11 +28,13 @@ def test_NotionClient__request(mock_call):
     ]
 
     client = NotionClient(NotionCredentials(token="MockToken"))
-    response = list(client._request("GET", "fake_endpoint", {}))
+    get_call = partial(client._get, endpoint="FAKE")
+
+    response = list(fetch_all_pages(get_call, NotionPagination))
     assert response == [{"result_id": 1}, {"result_id": 2}]
 
-    response = list(client._request("GET", "fake_endpoint", {}))
-    assert response == [{"result_id": 3}]
+    response = get_call()
+    assert response == {"result_id": 3}
 
 
 MOCK_BLOCK = [{"object": "block", "id": "1", "has_children": True}]

castor_extractor/knowledge/notion/client/constants.py

@@ -1,3 +1,4 @@
 NOTION_BASE_URL = "https://api.notion.com/v1/"
 CASTOR_NOTION_USER_AGENT = "castor-extractor"
 NOTION_TIMEOUT_MS = 300
+NOTION_VERSION = "2021-08-16"

castor_extractor/knowledge/notion/client/endpoints.py

@@ -1,4 +1,4 @@
-class EndpointFactory:
+class NotionEndpointFactory:
     """Wrapper class around all endpoints we're using"""
 
     BLOCKS = "blocks"

castor_extractor/knowledge/notion/client/pagination.py

@@ -1,16 +1,20 @@
 from typing import Optional
 
-from pydantic.dataclasses import dataclass
+from ....utils import PaginationModel
 
 
-@dataclass
-class PaginatedResponse:
+class NotionPagination(PaginationModel):
     """Class to handle paginated results"""
 
     results: list
     next_cursor: Optional[str]
     has_more: bool
 
-    @property
-    def start_cursor(self) -> dict:
+    def is_last(self) -> bool:
+        return not self.has_more
+
+    def next_page_payload(self) -> dict:
         return {"start_cursor": self.next_cursor}
+
+    def page_results(self) -> list:
+        return self.results

castor_extractor/quality/soda/assets.py

@@ -4,5 +4,5 @@ from ...types import ExternalAsset
 class SodaAsset(ExternalAsset):
     """Soda assets"""
 
-    CHECK_RESULTS = "check_results"
+    CHECKS = "checks"
     DATASETS = "datasets"

castor_extractor/quality/soda/client/client.py

@@ -1,99 +1,46 @@
-from time import sleep
-from typing import Any, Iterator, Optional
-
-import requests
-
-from ....utils import format_date, yesterday
-from ....utils.client.api import DEFAULT_TIMEOUT_S, HttpMethod
+from functools import partial
+from typing import Iterator
+
+from ....utils import (
+    APIClient,
+    BasicAuth,
+    fetch_all_pages,
+)
 from ..assets import SodaAsset
 from .credentials import SodaCredentials
 from .endpoints import SodaEndpointFactory
+from .pagination import SodaCloudPagination
 
-_REPORTING_PAGE_SIZE = 400
-_CLOUD_PAGE_SIZE = 100
+_CLOUD_API = "https://cloud.soda.io/api/v1/"
 _REQUESTS_PER_MINUTE = 10
 _SECONDS_PER_MINUTE = 60
-_RATE_LIMIT_S = (_SECONDS_PER_MINUTE // _REQUESTS_PER_MINUTE) + 1
+_RATE_LIMIT_MS = (_SECONDS_PER_MINUTE // _REQUESTS_PER_MINUTE) + 1
 
+HEADERS = {"Content-Type": "application/json"}
 
-class SodaClient:
-    def __init__(
-        self,
-        credentials: SodaCredentials,
-    ):
-        self._timeout = DEFAULT_TIMEOUT_S
-        self._reporting_headers = credentials.reporting_headers
-        self._auth = (credentials.api_key, credentials.secret)
 
-    def _call(
-        self,
-        url: str,
-        headers: dict,
-        method: HttpMethod = "GET",
-        *,
-        params: Optional[dict] = None,
-        data: Optional[dict] = None,
-        auth: Optional[tuple] = None,
-    ) -> Any:
-        response = requests.request(
-            method,
-            url,
-            headers=headers,
-            params=params,
-            json=data,
-            timeout=self._timeout,
-            auth=auth,
+class SodaClient(APIClient):
+    def __init__(self, credentials: SodaCredentials):
+        cloud_auth = BasicAuth(
+            username=credentials.api_key, password=credentials.secret
         )
-        response.raise_for_status()
-
-        return response.json()
+        super().__init__(host=_CLOUD_API, auth=cloud_auth, headers=HEADERS)
 
-    def _get_results_paginated(self, url: str, additional: dict) -> Iterator:
-        page_number = 1
-        next_page = True
-        while next_page:
-            json_data = {
-                **{"page": page_number, "size": _REPORTING_PAGE_SIZE},
-                **additional,
-            }
-            _check_results_page = self._call(
-                url=url,
-                method="POST",
-                data=json_data,
-                headers=self._reporting_headers,
-            )
-            yield from _check_results_page["data"]
-
-            next_page = len(_check_results_page["data"]) == _REPORTING_PAGE_SIZE
-            page_number += 1
-
-    def _datasets(self) -> Iterator[dict]:
-        url = SodaEndpointFactory.datasets()
-        next_page = True
-        page_number = 0
-        while next_page:
-            data = self._call(
-                url=url,
-                method="GET",
-                headers={},
-                params={"size": _CLOUD_PAGE_SIZE, "page": page_number},
-                auth=self._auth,
-            )
-            yield from data["content"]
-            next_page = not data["last"]
-            page_number += 1
-            sleep(_RATE_LIMIT_S)
+    def datasets(self) -> Iterator[dict]:
+        request = partial(self._get, endpoint=SodaEndpointFactory.datasets())
+        yield from fetch_all_pages(
+            request, SodaCloudPagination, rate_limit=_RATE_LIMIT_MS
+        )
 
-    def _check_results(self) -> Iterator:
-        url = SodaEndpointFactory.check_results()
-        _date = format_date(timestamp=yesterday())
-        return self._get_results_paginated(
-            url, additional={"from_datetime": _date}
+    def checks(self) -> Iterator[dict]:
+        request = partial(self._get, endpoint=SodaEndpointFactory.checks())
+        yield from fetch_all_pages(
+            request, SodaCloudPagination, rate_limit=_RATE_LIMIT_MS
         )
 
-    def fetch(self, asset: SodaAsset) -> Iterator[dict]:
+    def fetch(self, asset: SodaAsset) -> Iterator:
         if asset == SodaAsset.DATASETS:
-            yield from self._datasets()
-        if asset == SodaAsset.CHECK_RESULTS:
-            yield from self._check_results()
+            return self.datasets()
+        if asset == SodaAsset.CHECKS:
+            return self.checks()
         raise ValueError(f"The asset {asset}, is not supported")

castor_extractor/quality/soda/client/credentials.py

@@ -1,8 +1,5 @@
-from typing import Dict
-
 from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
-from requests.auth import HTTPBasicAuth
 
 SODA_ENV_PREFIX = "CASTOR_SODA_"
 
@@ -18,11 +15,3 @@ class SodaCredentials(BaseSettings):
 
     api_key: str = Field(repr=False)
     secret: str = Field(repr=False)
-
-    @property
-    def reporting_headers(self) -> Dict[str, str]:
-        return {
-            "Content-Type": "application/json",
-            "API_KEY_ID": self.api_key,
-            "API_KEY_SECRET": self.secret,
-        }

castor_extractor/quality/soda/client/endpoints.py

@@ -1,13 +1,10 @@
 class SodaEndpointFactory:
     """Wrapper class around all endpoints we're using"""
 
-    CLOUD_API = "https://cloud.soda.io/api/v1"
-    REPORTING_API = "https://reporting.cloud.soda.io/v1"
-
     @classmethod
     def datasets(cls) -> str:
-        return f"{cls.CLOUD_API}/datasets"
+        return "datasets"
 
     @classmethod
-    def check_results(cls) -> str:
-        return f"{cls.REPORTING_API}/quality/check_results"
+    def checks(cls) -> str:
+        return "checks"

castor_extractor/quality/soda/client/pagination.py

@@ -0,0 +1,25 @@
+from typing import List
+
+from ....utils import PaginationModel
+
+_CLOUD_PAGE_SIZE = 100
+_CLOUD_FIRST_PAGE = 0
+
+
+class SodaCloudPagination(PaginationModel):
+    content: List[dict]
+    last: bool
+
+    def is_last(self) -> bool:
+        return self.last
+
+    def next_page_payload(self) -> dict:
+        current_page = (
+            self.current_page_payload[self.fetch_by.value]["page"]
+            if self.current_page_payload
+            else _CLOUD_FIRST_PAGE
+        )
+        return {"page": current_page + 1, "size": _CLOUD_PAGE_SIZE}
+
+    def page_results(self) -> list:
+        return self.content

castor_extractor/utils/__init__.py

@@ -1,12 +1,24 @@
 from .argument_parser import parse_filled_arguments
 from .client import (
     AbstractSourceClient,
+    APIClient,
+    Auth,
+    BasicAuth,
+    BearerAuth,
+    CustomAuth,
     ExtractionQuery,
+    FetchNextPageBy,
+    PaginationModel,
     PostgresClient,
+    RequestSafeMode,
+    ResponseJson,
     SqlalchemyClient,
+    build_url,
+    fetch_all_pages,
+    handle_response,
     uri_encode,
 )
-from .collection import empty_iterator, group_by, mapping_from_rows
+from .collection import deduplicate, empty_iterator, group_by, mapping_from_rows
 from .constants import OUTPUT_DIR
 from .deprecate import deprecate_python
 from .env import from_env
@@ -24,7 +36,6 @@ from .pager import (
 )
 from .retry import RetryStrategy, retry
 from .safe import SafeMode, safe_mode
-from .safe_request import RequestSafeMode, ResponseJson, handle_response
 from .store import AbstractStorage, LocalStorage
 from .string import decode_when_bytes, string_to_tuple
 from .time import (

castor_extractor/utils/client/__init__.py

@@ -1,4 +1,18 @@
 from .abstract import AbstractSourceClient, SqlalchemyClient
+from .api import (
+    APIClient,
+    Auth,
+    BasicAuth,
+    BearerAuth,
+    CustomAuth,
+    FetchNextPageBy,
+    PaginationModel,
+    RequestSafeMode,
+    ResponseJson,
+    build_url,
+    fetch_all_pages,
+    handle_response,
+)
 from .postgres import PostgresClient
 from .query import ExtractionQuery
 from .uri import uri_encode

castor_extractor/utils/client/api/__init__.py

@@ -0,0 +1,5 @@
+from .auth import Auth, BasicAuth, BearerAuth, CustomAuth
+from .client import APIClient
+from .pagination import FetchNextPageBy, PaginationModel, fetch_all_pages
+from .safe_request import RequestSafeMode, ResponseJson, handle_response
+from .utils import build_url

castor_extractor/utils/client/api/auth.py

@@ -0,0 +1,76 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, Union
+
+from requests.auth import AuthBase, HTTPBasicAuth
+
+logger = logging.getLogger(__name__)
+
+
+class BasicAuth(HTTPBasicAuth):
+    """
+    Authentication for API using basic auth method
+
+    - Instantiate with username and password
+    - pass the Auth class to the APIClient
+    """
+
+    def refresh_token(self):
+        pass
+
+
+class CustomAuth(AuthBase, ABC):
+    """
+    Authentication for API using custom auth method
+
+    You need to:
+     - implement the `_authentication_header()` method
+     - pass the Auth class to the APIClient
+    """
+
+    def refresh_token(self):
+        """Method to refresh the token if token expires"""
+        pass
+
+    @abstractmethod
+    def _authentication_header(self) -> Dict[str, str]:
+        pass
+
+    def __call__(self, r):
+        r.headers = {**r.headers, **self._authentication_header()}
+        return r
+
+
+class BearerAuth:
+    """
+    Authentication for API using Bearer tokens
+
+    You need to:
+    - implement the `fetch_token()` method
+    - pass the Auth class to the APIClient
+    """
+
+    _token: Optional[str] = None
+    authentication_key = "Authorization"
+
+    @abstractmethod
+    def fetch_token(self) -> Optional[str]:
+        """Method that should return the bearer token"""
+        pass
+
+    def refresh_token(self):
+        """Method to refresh the token if token expires"""
+        self._fetch_token(force_refresh=True)
+
+    def _fetch_token(self, force_refresh: bool = False) -> Optional[str]:
+        if not self._token or force_refresh:
+            logger.info("Refreshing authentication token...")
+            self._token = self.fetch_token()
+        return f"Bearer {self._token}"
+
+    def __call__(self, r):
+        r.headers[self.authentication_key] = self._fetch_token()
+        return r
+
+
+Auth = Union[BasicAuth, CustomAuth, BearerAuth]