castor-extractor 0.19.4__py3-none-any.whl → 0.19.7__py3-none-any.whl

CHANGELOG.md

@@ -1,5 +1,18 @@
+
 # Changelog
 
+## 0.19.7 - 2024-09-05
+
+* Metabase: Handle compatibility with older version
+
+## 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

castor_extractor/quality/soda/client/pagination.py

@@ -15,7 +15,7 @@ class SodaCloudPagination(PaginationModel):
 
     def next_page_payload(self) -> dict:
         current_page = (
-            self.current_page_payload["page"]
+            self.current_page_payload[self.fetch_by.value]["page"]
             if self.current_page_payload
             else _CLOUD_FIRST_PAGE
         )

castor_extractor/utils/__init__.py

@@ -7,6 +7,7 @@ from .client import (
     BearerAuth,
     CustomAuth,
     ExtractionQuery,
+    FetchNextPageBy,
     PaginationModel,
     PostgresClient,
     RequestSafeMode,

castor_extractor/utils/client/__init__.py

@@ -5,6 +5,7 @@ from .api import (
     BasicAuth,
     BearerAuth,
     CustomAuth,
+    FetchNextPageBy,
     PaginationModel,
     RequestSafeMode,
     ResponseJson,
@@ -12,7 +13,6 @@ from .api import (
     fetch_all_pages,
     handle_response,
 )
-from .api_deprecated import APIClientDeprecated
 from .postgres import PostgresClient
 from .query import ExtractionQuery
 from .uri import uri_encode

castor_extractor/utils/client/api/__init__.py

@@ -1,5 +1,5 @@
 from .auth import Auth, BasicAuth, BearerAuth, CustomAuth
 from .client import APIClient
-from .pagination import PaginationModel, fetch_all_pages
+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/client.py

@@ -1,6 +1,6 @@
 import logging
 from http import HTTPStatus
-from typing import Dict, Literal, Optional
+from typing import Dict, Literal, Optional, Tuple
 
 import requests
 from requests import Response
@@ -21,6 +21,25 @@ 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
@@ -59,22 +78,27 @@ class APIClient:
         *,
         headers: Headers = None,
         params: Optional[dict] = None,
+        data: Optional[dict] = None,
         pagination_params: Optional[dict] = None,
     ) -> Response:
         headers = headers or {}
-        _pagination_params = pagination_params or {}
-        params = params or {}
+
+        data, params = _generate_payloads(
+            method=method,
+            params=params,
+            data=data,
+            pagination_params=pagination_params,
+        )
 
         url = build_url(self._host, endpoint)
-        combined_params = {**params, **_pagination_params}
 
         return requests.request(
             method=method,
             url=url,
             auth=self._auth,
             headers={**self.base_headers, **headers},
-            params=combined_params if method == "GET" else None,
-            json=combined_params if method == "POST" else None,
+            params=params,
+            json=data,
             timeout=self._timeout,
         )
 
@@ -88,12 +112,14 @@ class APIClient:
         *,
         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,
         )
@@ -117,7 +143,7 @@ class APIClient:
         response = self._call(
             method="POST",
             endpoint=endpoint,
-            params=data,
+            data=data,
             pagination_params=pagination_params,
             headers=headers,
         )

castor_extractor/utils/client/api/pagination.py

@@ -1,14 +1,28 @@
 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
+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
@@ -18,6 +32,7 @@ class PaginationModel(BaseModel):
     `fetch_all_pages` method
     """
 
+    fetch_by: FetchNextPageBy = FetchNextPageBy.PAYLOAD
     current_page_payload: Optional["dict"] = None
 
     @abstractmethod
@@ -26,7 +41,7 @@ class PaginationModel(BaseModel):
         pass
 
     @abstractmethod
-    def next_page_payload(self) -> dict:
+    def next_page_payload(self) -> Optional[Union[dict, str]]:
         """Payload enabling to generate the request for the next page"""
         pass
 
@@ -35,6 +50,9 @@ class PaginationModel(BaseModel):
         """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,
@@ -51,13 +69,12 @@ def fetch_all_pages(
     while not paginated_response.is_last():
         logger.info(f"Fetching page number {page_number}")
         yield from paginated_response.page_results()
-        next_payload = paginated_response.next_page_payload()
-        new_request = partial(request, pagination_params=next_payload)
+        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_payload, **new_request()
+            current_page_payload=next_page_parameters, **new_request()
         )
         page_number += 1
 

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/salesforce/client.py

@@ -1,16 +1,40 @@
 import logging
-from typing import Iterator, Optional, Tuple
-
-from requests import Response
-
-from ...utils.client import APIClientDeprecated
+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(APIClientDeprecated):
+
+class SalesforceBaseClient(APIClient):
     """
     Salesforce API client.
     https://developer.salesforce.com/docs/atlas.en-us.api_rest.meta/api_rest/intro_rest.htm
@@ -22,45 +46,24 @@ class SalesforceBaseClient(APIClientDeprecated):
     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(APIClientDeprecated):
 
         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/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,16 +129,16 @@ 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")
             if not collection:
                 continue
             for dashboard in cast(SerializedAsset, collection):
+                if dashboard.get("model") != "dashboard":
+                    # This is to maintain compatibility with older versions
+                    # where ?models=dashboard has no effects
+                    continue
                 dashboard_id = dashboard.get("id")
                 if not dashboard_id:
                     continue

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

@@ -33,12 +33,13 @@ _SIGMA_HEADERS = {
     "Content-Type": _CONTENT_TYPE,
 }
 
-_VOLUME_IGNORED = 200
+_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,

castor_extractor/visualization/tableau_revamp/assets.py

@@ -19,3 +19,11 @@ class TableauRevampAsset(ExternalAsset):
     USAGE = "usage"
     USER = "users"
     WORKBOOK = "workbooks"
+
+
+# assets that are only available for clients using Tableau Pulse
+TABLEAU_PULSE_ASSETS = (
+    TableauRevampAsset.METRIC,
+    TableauRevampAsset.METRIC_DEFINITION,
+    TableauRevampAsset.SUBSCRIPTION,
+)

castor_extractor/visualization/tableau_revamp/client/client.py

@@ -3,7 +3,7 @@ import logging
 import tableauserverclient as TSC  # type: ignore
 
 from ....utils import SerializedAsset
-from ..assets import TableauRevampAsset
+from ..assets import TABLEAU_PULSE_ASSETS, TableauRevampAsset
 from ..constants import DEFAULT_TIMEOUT_SECONDS
 from .client_metadata_api import TableauClientMetadataApi
 from .client_rest_api import TableauClientRestApi
@@ -111,9 +111,11 @@ class TableauRevampClient:
         self,
         credentials: TableauRevampCredentials,
         timeout_sec: int = DEFAULT_TIMEOUT_SECONDS,
+        with_pulse: bool = False,
     ):
         self._credentials = credentials
         self._server = _server(credentials.server_url, timeout_sec)
+        self._with_pulse = with_pulse
 
         self._client_metadata = TableauClientMetadataApi(server=self._server)
         self._client_rest = TableauClientRestApi(server=self._server)
@@ -192,6 +194,9 @@ class TableauRevampClient:
         """
         Extract the given Tableau Asset
         """
+        if asset in TABLEAU_PULSE_ASSETS and not self._with_pulse:
+            logger.info(f"Skipping asset {asset} - Tableau Pulse de-activated")
+            return []
 
         logger.info(f"Extracting {asset.name}...")