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

castor_extractor/utils/salesforce/client.py

@@ -1,14 +1,38 @@
 import logging
-from typing import Iterator, Optional, Tuple
-
-from requests import Response
-
-from ...utils.client.api import APIClient
+from functools import partial
+from typing import Iterator, Optional
+
+import requests
+
+from ...utils import (
+    APIClient,
+    BearerAuth,
+    build_url,
+    fetch_all_pages,
+    handle_response,
+)
 from .constants import DEFAULT_API_VERSION, DEFAULT_PAGINATION_LIMIT
 from .credentials import SalesforceCredentials
+from .pagination import SalesforcePagination
 
 logger = logging.getLogger(__name__)
 
+SALESFORCE_TIMEOUT_S = 120
+
+
+class SalesforceAuth(BearerAuth):
+    _AUTH_ENDPOINT = "services/oauth2/token"
+
+    def __init__(self, credentials: SalesforceCredentials):
+        self._host = credentials.base_url
+        self._token_payload = credentials.token_request_payload()
+
+    def fetch_token(self) -> Optional[str]:
+        url = build_url(self._host, self._AUTH_ENDPOINT)
+        response = requests.post(url, "POST", params=self._token_payload)
+        handled_response = handle_response(response)
+        return handled_response["access_token"]
+
 
 class SalesforceBaseClient(APIClient):
     """
@@ -22,45 +46,24 @@ class SalesforceBaseClient(APIClient):
     PATH_TPL = "services/data/v{version}/{suffix}"
 
     def __init__(self, credentials: SalesforceCredentials):
-        super().__init__(host=credentials.base_url)
-        self._token = self._access_token(credentials)
-
-    def _access_token(self, credentials: SalesforceCredentials) -> str:
-        url = self.build_url(self._host, "services/oauth2/token")
-        response = self._call(
-            url, "POST", params=credentials.token_request_payload()
+        auth = SalesforceAuth(credentials)
+        super().__init__(
+            host=credentials.base_url, auth=auth, timeout=SALESFORCE_TIMEOUT_S
         )
-        return response["access_token"]
 
-    def _full_url(self, suffix: str) -> str:
+    def _endpoint(self, suffix: str) -> str:
         path = self.PATH_TPL.format(version=self.api_version, suffix=suffix)
-        return self.build_url(self._host, path)
+        return path
 
     @property
-    def query_url(self) -> str:
+    def query_endpoint(self) -> str:
         """Returns the query API url"""
-        return self._full_url("query")
+        return self._endpoint("query")
 
     @property
-    def tooling_url(self) -> str:
+    def tooling_endpoint(self) -> str:
         """Returns the tooling API url"""
-        return self._full_url("tooling/query")
-
-    @staticmethod
-    def _query_processor(response: Response) -> Tuple[dict, Optional[str]]:
-        results = response.json()
-        return results["records"], results.get("nextRecordsUrl")
-
-    def _has_reached_pagination_limit(self, page_number: int) -> bool:
-        return page_number > self.pagination_limit
-
-    def _query_first_page(self, query: str) -> Tuple[Iterator[dict], str]:
-        url = self.query_url
-        logger.info("querying page 0")
-        records, next_page_url = self._call(
-            url, params={"q": query}, processor=self._query_processor
-        )
-        return records, next_page_url
+        return self._endpoint("tooling/query")
 
     def _query_all(self, query: str) -> Iterator[dict]:
         """
@@ -68,17 +71,9 @@ class SalesforceBaseClient(APIClient):
 
         more: https://developer.salesforce.com/docs/atlas.en-us.api_rest.meta/api_rest/dome_query.htm
         """
-        records, next_page_path = self._query_first_page(query)
-        yield from records
-
-        page_count = 1
-        while next_page_path and not self._has_reached_pagination_limit(
-            page_count
-        ):
-            logger.info(f"querying page {page_count}")
-            url = self.build_url(self._host, next_page_path)
-            records, next_page_path = self._call(
-                url, processor=self._query_processor
-            )
-            yield from records
-            page_count += 1
+        request = partial(
+            self._get,
+            endpoint=self.query_endpoint,
+            params={"q": query},
+        )
+        yield from fetch_all_pages(request, SalesforcePagination)

castor_extractor/utils/salesforce/client_test.py

@@ -17,5 +17,5 @@ def test_SalesforceBaseClient__urls(mock_call):
     )
     client = SalesforceBaseClient(credentials)
 
-    assert client.query_url == "https://url/services/data/v59.0/query"
-    assert client.tooling_url == "https://url/services/data/v59.0/tooling/query"
+    assert client.query_endpoint == "services/data/v59.0/query"
+    assert client.tooling_endpoint == "services/data/v59.0/tooling/query"

castor_extractor/utils/salesforce/pagination.py

@@ -0,0 +1,33 @@
+from typing import Optional
+
+from pydantic import ConfigDict
+from pydantic.alias_generators import to_camel
+
+from ...utils import (
+    FetchNextPageBy,
+    PaginationModel,
+)
+
+LIMIT_RECORDS_PER_PAGE = 2000
+
+
+class SalesforcePagination(PaginationModel):
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )
+    fetch_by: FetchNextPageBy = FetchNextPageBy.URL
+    records: list
+    next_records_url: Optional[str] = None
+
+    def is_last(self) -> bool:
+        no_next_page = not self.next_records_url
+        page_incomplete = len(self.records) < LIMIT_RECORDS_PER_PAGE
+        return no_next_page or page_incomplete
+
+    def next_page_payload(self) -> Optional[str]:
+        return self.next_records_url
+
+    def page_results(self) -> list:
+        return self.records

castor_extractor/visualization/domo/client/client.py

@@ -153,11 +153,16 @@ class DomoClient:
 
         return all_results
 
-    def _datasources(self, page_id: str) -> RawData:
-        endpoint = self._endpoint_factory.page_content(page_id)
-        page_content = self._get_element(endpoint)
+    def _datasources(self, card_ids: List[int]) -> RawData:
+        """Yields all distinct datasources associated to the given cards"""
+        if not card_ids:
+            return empty_iterator()
+
+        endpoint = self._endpoint_factory.cards_metadata(card_ids)
+        cards_metadata = self._get_element(endpoint)
+
         processed: set[str] = set()
-        for card in page_content.get("cards", []):
+        for card in cards_metadata:
             for datasource in card["datasources"]:
                 id_ = datasource["dataSourceId"]
                 if id_ in processed:
@@ -195,7 +200,7 @@ class DomoClient:
             if not detail:
                 continue
 
-            datasources = self._datasources(page_id)
+            datasources = self._datasources(detail.get("cardIds", []))
             yield {
                 **detail,
                 "datasources": list(datasources),

castor_extractor/visualization/domo/client/credentials.py

@@ -1,6 +1,6 @@
 from typing import Dict, Optional
 
-from pydantic import Field, SecretStr
+from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from requests.auth import HTTPBasicAuth
 

castor_extractor/visualization/domo/client/endpoints.py

@@ -1,10 +1,19 @@
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import List, Optional
 
-_DOMO_PUBLIC_URL = "https://api.domo.com"
 _AUTH_URL = (
     "grant_type=client_credentials&scope=data%20dashboard%20audit%20user"
 )
+_DOMO_PUBLIC_URL = "https://api.domo.com"
+_MAX_URL_LENGTH = 2048
+
+
+class URLTooLongException(Exception):
+    """Custom error handling if the URL is too big"""
+
+    def __init__(self, url: str, max_size: int):
+        message = f"The URL is too long ({len(url)} > {max_size}) : {url}"
+        super().__init__(message)
 
 
 @dataclass
@@ -70,8 +79,11 @@ class EndpointFactory:
             is_private=True,
         )
 
-    def page_content(self, page_id: str) -> Endpoint:
-        return Endpoint(
-            base_url=f"{self.base_url}/api/content/v3/stacks/{page_id}/cards?parts=datasources",
-            is_private=True,
-        )
+    def cards_metadata(self, card_ids: List[int]) -> Endpoint:
+        urns = ",".join(map(str, card_ids))
+        url = f"{self.base_url}/api/content/v1/cards?urns={urns}&parts=datasources"
+
+        if len(url) > _MAX_URL_LENGTH:
+            raise URLTooLongException(url, _MAX_URL_LENGTH)
+
+        return Endpoint(base_url=url, is_private=True)

castor_extractor/visualization/looker/api/credentials.py

@@ -1,5 +1,5 @@
 from looker_sdk.rtl.api_settings import SettingsConfig
-from pydantic import Field, SecretStr
+from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from ..constants import DEFAULT_LOOKER_TIMEOUT_SECOND, LOOKER_ENV_PREFIX

castor_extractor/visualization/metabase/client/api/client.py

@@ -1,10 +1,16 @@
 import logging
-from typing import Dict, Iterator, List, cast
+from http import HTTPStatus
+from typing import Any, Dict, Iterator, List, Optional, cast
 
 import requests
-from requests import HTTPError, Response
-
-from .....utils import JsonType, SerializedAsset
+from requests import HTTPError
+
+from .....utils import (
+    JsonType,
+    RequestSafeMode,
+    SerializedAsset,
+    handle_response,
+)
 from ...assets import EXPORTED_FIELDS, MetabaseAsset
 from ...errors import MetabaseLoginError, SuperuserCredentialsRequired
 from ...types import IdsType
@@ -13,6 +19,17 @@ from .credentials import MetabaseApiCredentials
 
 logger = logging.getLogger(__name__)
 
+# Safe mode
+VOLUME_IGNORED = 5
+IGNORED_ERROR_CODES = (
+    HTTPStatus.BAD_REQUEST,
+    HTTPStatus.NOT_FOUND,
+)
+METABASE_SAFE_MODE = RequestSafeMode(
+    max_errors=VOLUME_IGNORED,
+    status_codes=IGNORED_ERROR_CODES,
+)
+
 URL_TEMPLATE = "{base_url}/api/{endpoint}"
 
 ROOT_KEY = "root"
@@ -31,12 +48,14 @@ class ApiClient:
     def __init__(
         self,
         credentials: MetabaseApiCredentials,
+        safe_mode: Optional[RequestSafeMode] = None,
     ):
         self.base_url = credentials.base_url
 
         self._credentials = credentials
         self._session = requests.Session()
         self._session_id = self._login()
+        self.safe_mode = safe_mode or METABASE_SAFE_MODE
         self._check_permissions()  # verify that the given user is superuser
 
     @staticmethod
@@ -57,8 +76,8 @@ class ApiClient:
         }
 
     @staticmethod
-    def _answer(response: Response):
-        answer = response.json()
+    def _answer(response: Any):
+        answer = response
         if isinstance(answer, Dict) and DATA_KEY in answer:
             # v0.41 of Metabase introduced embedded data for certain calls
             # {'data': [{ }, ...] , 'total': 15, 'limit': None, 'offset': None}"
@@ -69,7 +88,7 @@ class ApiClient:
         url = self._url(endpoint)
         headers = self._headers()
         response = self._session.get(url=url, headers=headers)
-        response.raise_for_status()  # check for errors
+        response = handle_response(response, safe_mode=self.safe_mode)
         return self._answer(response)
 
     def _check_permissions(self) -> None:
@@ -110,10 +129,6 @@ class ApiClient:
         return ids
 
     def _dashboards(self) -> Iterator[dict]:
-        """
-        GET /api/dashboard is deprecated
-        https://github.com/metabase/metabase/pull/35235
-        """
         collection_ids = self._fetch_ids(MetabaseAsset.COLLECTION)
         for _id in collection_ids:
             collection = self._call(f"collection/{_id}/items?models=dashboard")

castor_extractor/visualization/metabase/client/api/credentials.py

@@ -1,4 +1,4 @@
-from pydantic import Field, SecretStr
+from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 METABASE_API_ENV_PREFIX = "CASTOR_METABASE_API_"

castor_extractor/visualization/metabase/client/db/credentials.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from pydantic import Field, SecretStr
+from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 METABASE_DB_ENV_PREFIX = "CASTOR_METABASE_DB_"

castor_extractor/visualization/mode/client/credentials.py

@@ -1,4 +1,4 @@
-from pydantic import Field, SecretStr, field_validator
+from pydantic import Field, field_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 MODE_ENV_PREFIX = "CASTOR_MODE_ANALYTICS_"

castor_extractor/visualization/qlik/client/engine/credentials.py

@@ -1,4 +1,4 @@
-from pydantic import Field, SecretStr, field_validator
+from pydantic import Field, field_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from .....utils import validate_baseurl

castor_extractor/visualization/salesforce_reporting/client/rest.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Dict, Iterator, List, Optional
 
+from ....utils import build_url
 from ....utils.salesforce import SalesforceBaseClient
 from ..assets import SalesforceReportingAsset
 from .soql import queries
@@ -28,15 +29,15 @@ class SalesforceReportingClient(SalesforceBaseClient):
 
         if asset_type == SalesforceReportingAsset.DASHBOARDS:
             path = f"lightning/r/Dashboard/{asset['Id']}/view"
-            return self.build_url(self._host, path)
+            return build_url(self._host, path)
 
         if asset_type == SalesforceReportingAsset.FOLDERS:
             path = asset["attributes"]["url"].lstrip("/")
-            return self.build_url(self._host, path)
+            return build_url(self._host, path)
 
         if asset_type == SalesforceReportingAsset.REPORTS:
             path = f"lightning/r/Report/{asset['Id']}/view"
-            return self.build_url(self._host, path)
+            return build_url(self._host, path)
 
         return None
 

castor_extractor/visualization/sigma/client/client.py

@@ -1,150 +1,142 @@
-import logging
-from typing import Dict, Iterator, List, Optional, Tuple
-from urllib.parse import urljoin
+from functools import partial
+from http import HTTPStatus
+from typing import Callable, Dict, Iterator, List, Optional, Tuple
 
 import requests
 
+from ....utils import (
+    APIClient,
+    BearerAuth,
+    RequestSafeMode,
+    build_url,
+    fetch_all_pages,
+    handle_response,
+)
 from ..assets import SigmaAsset
 from .credentials import SigmaCredentials
-from .endpoints import EndpointFactory
-from .pagination import Pagination
-
-logger = logging.getLogger()
+from .endpoints import SigmaEndpointFactory
+from .pagination import SIGMA_API_LIMIT, SigmaPagination
 
+_CONTENT_TYPE = "application/x-www-form-urlencoded"
 
-DATA_ELEMENTS: Tuple[str, ...] = (
+_DATA_ELEMENTS: Tuple[str, ...] = (
     "input-table",
     "pivot-table",
     "table",
     "visualization",
     "viz",
 )
-_CONTENT_TYPE = "application/x-www-form-urlencoded"
 
+_AUTH_TIMEOUT_S = 60
 
-class SigmaClient:
-    """Client used for all Sigma's assets extractions"""
-
-    def __init__(self, credentials: SigmaCredentials):
-        self.host = credentials.host
-        self.client_id = credentials.client_id
-        self.api_token = credentials.api_token
-        self.grant_type = credentials.grant_type
-        self.headers: Optional[Dict[str, str]] = None
-
-    def _get_token(self) -> Dict[str, str]:
-        auth_endpoint = EndpointFactory.authentication()
-        token_api_path = urljoin(self.host, auth_endpoint)
-        token_response = requests.post(  # noqa: S113
-            token_api_path,
-            data={
-                "grant_type": self.grant_type,
-                "client_id": self.client_id,
-                "client_secret": self.api_token,
-            },
-        )
-        if token_response.status_code != requests.codes.OK:
-            raise ValueError("Couldn't fetch the token in the API")
-        return token_response.json()
-
-    def _get_headers(self, reset=False) -> Dict[str, str]:
-        """
-        If reset is True, will re-create the headers with a new authentication token
-
-        Note : From this [documentation](https://help.sigmacomputing.com/docs/api-authentication-with-curl),
-        instead of re-creating a token we could refresh it, but I don't see any benefit.
-        """
-        if reset or not self.headers:
-            headers = {"Content-Type": _CONTENT_TYPE}
-            token = self._get_token()
-            headers["Authorization"] = f"Bearer {token['access_token']}"
-            self.headers = headers
-        return self.headers
-
-    def _get(self, endpoint_url: str) -> dict:
-        url = urljoin(self.host, endpoint_url)
-        result = requests.get(url, headers=self._get_headers())  # noqa: S113
-
-        if result.status_code == requests.codes.UNAUTHORIZED:
-            logger.info("Regenerating access token")
-            result = requests.get(  # noqa: S113
-                url, headers=self._get_headers(reset=True)
-            )
+_SIGMA_HEADERS = {
+    "Content-Type": _CONTENT_TYPE,
+}
 
-        try:
-            return result.json()
-        except Exception as e:
-            logger.warning(
-                f"Couldn't deserialize result from url {url}."
-                f" with status code {result.status_code} and"
-                f" exception {type(e)}"
-            )
-            return dict()
-
-    def _get_with_pagination(self, endpoint_url: str) -> Iterator[dict]:
-        pagination = Pagination(next_page="0")
-
-        while pagination.next_page is not None:
-            paginated_url = pagination.generate_url(endpoint_url)
-            response = self._get(paginated_url)
-            pagination = Pagination(
-                next_page=response.get("nextPage"),
-                entries=response.get("entries"),
-                total=response.get("total"),
-            )
-            yield from pagination.entries
+_VOLUME_IGNORED = 10_000
+_IGNORED_ERROR_CODES = (
+    HTTPStatus.BAD_REQUEST,
+    HTTPStatus.BAD_GATEWAY,
+    HTTPStatus.INTERNAL_SERVER_ERROR,
+    HTTPStatus.CONFLICT,
+    HTTPStatus.NOT_FOUND,
+)
+SIGMA_SAFE_MODE = RequestSafeMode(
+    max_errors=_VOLUME_IGNORED,
+    status_codes=_IGNORED_ERROR_CODES,
+)
+
+
+class SigmaBearerAuth(BearerAuth):
+    def __init__(self, host: str, token_payload: Dict[str, str]):
+        auth_endpoint = SigmaEndpointFactory.authentication()
+        self.authentication_url = build_url(host, auth_endpoint)
+        self.token_payload = token_payload
+
+    def fetch_token(self):
+        token_api_path = self.authentication_url
+        token_response = requests.post(
+            token_api_path, data=self.token_payload, timeout=_AUTH_TIMEOUT_S
+        )
+        return handle_response(token_response)["access_token"]
 
-    def _per_workbook_get_pages(self, workbook_id: str) -> Iterator[dict]:
-        endpoint = EndpointFactory.pages(workbook_id)
-        yield from self._get_with_pagination(endpoint)
 
-    def _per_page_get_elements(
+class SigmaClient(APIClient):
+    def __init__(
         self,
-        workbook_id: str,
-        page_id: str,
-    ) -> Iterator[dict]:
-        endpoint = EndpointFactory.elements(workbook_id, page_id)
-        yield from self._get_with_pagination(endpoint)
+        credentials: SigmaCredentials,
+        safe_mode: Optional[RequestSafeMode] = None,
+    ):
+        auth = SigmaBearerAuth(
+            host=credentials.host,
+            token_payload=credentials.token_payload,
+        )
+        super().__init__(
+            host=credentials.host,
+            auth=auth,
+            headers=_SIGMA_HEADERS,
+            safe_mode=safe_mode or SIGMA_SAFE_MODE,
+        )
+
+    def _get_paginated(self, endpoint: str) -> Callable:
+        return partial(
+            self._get, endpoint=endpoint, params={"limit": SIGMA_API_LIMIT}
+        )
 
     def _get_all_datasets(self) -> Iterator[dict]:
-        endpoint = EndpointFactory.datasets()
-        yield from self._get_with_pagination(endpoint)
+        request = self._get_paginated(endpoint=SigmaEndpointFactory.datasets())
+        yield from fetch_all_pages(request, SigmaPagination)
 
     def _get_all_files(self) -> Iterator[dict]:
-        endpoint = EndpointFactory.files()
-        yield from self._get_with_pagination(endpoint)
+        request = self._get_paginated(endpoint=SigmaEndpointFactory.files())
+        yield from fetch_all_pages(request, SigmaPagination)
 
     def _get_all_members(self) -> Iterator[dict]:
-        endpoint = EndpointFactory.members()
-        yield from self._get(endpoint)
+        request = self._get_paginated(endpoint=SigmaEndpointFactory.members())
+        yield from fetch_all_pages(request, SigmaPagination)
 
     def _get_all_workbooks(self) -> Iterator[dict]:
-        endpoint = EndpointFactory.workbooks()
-        yield from self._get_with_pagination(endpoint)
+        request = self._get_paginated(endpoint=SigmaEndpointFactory.workbooks())
+        yield from fetch_all_pages(request, SigmaPagination)
+
+    def _get_elements_per_page(
+        self, page: dict, workbook_id: str
+    ) -> Iterator[dict]:
+        page_id = page["pageId"]
+        request = self._get_paginated(
+            SigmaEndpointFactory.elements(workbook_id, page_id)
+        )
+        elements = fetch_all_pages(request, SigmaPagination)
+        for element in elements:
+            if element.get("type") not in _DATA_ELEMENTS:
+                continue
+            yield {
+                **element,
+                "workbook_id": workbook_id,
+                "page_id": page_id,
+            }
 
     def _get_all_elements(self, workbooks: List[dict]) -> Iterator[dict]:
         for workbook in workbooks:
             workbook_id = workbook["workbookId"]
-            pages = self._per_workbook_get_pages(workbook_id)
+
+            request = self._get_paginated(
+                SigmaEndpointFactory.pages(workbook_id)
+            )
+            pages = fetch_all_pages(request, SigmaPagination)
 
             for page in pages:
-                page_id = page["pageId"]
-                elements = self._per_page_get_elements(workbook_id, page_id)
-                for element in elements:
-                    if element.get("type") not in DATA_ELEMENTS:
-                        continue
-                    yield {
-                        **element,
-                        "workbook_id": workbook_id,
-                        "page_id": page_id,
-                    }
+                yield from self._get_elements_per_page(
+                    page=page, workbook_id=workbook_id
+                )
 
     def _get_all_lineages(self, elements: List[dict]) -> Iterator[dict]:
         for element in elements:
             workbook_id = element["workbook_id"]
             element_id = element["elementId"]
-            endpoint = EndpointFactory.lineage(workbook_id, element_id)
-            lineage = self._get(endpoint)
+            lineage = self._get(
+                endpoint=SigmaEndpointFactory.lineage(workbook_id, element_id)
+            )
             yield {
                 **lineage,
                 "workbook_id": workbook_id,
@@ -154,8 +146,11 @@ class SigmaClient:
     def _get_all_queries(self, workbooks: List[dict]) -> Iterator[dict]:
         for workbook in workbooks:
             workbook_id = workbook["workbookId"]
-            endpoint = EndpointFactory.queries(workbook_id)
-            queries = self._get_with_pagination(endpoint)
+            request = self._get_paginated(
+                SigmaEndpointFactory.queries(workbook_id)
+            )
+            queries = fetch_all_pages(request, SigmaPagination)
+
             for query in queries:
                 yield {**query, "workbook_id": workbook_id}