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

castor_extractor/uploader/upload.py

@@ -32,7 +32,6 @@ def _path_and_url(
     file_type: FileType,
     file_path: str,
 ) -> Tuple[str, str]:
-
     now = datetime.utcnow()
     timestamp = int(now.timestamp())
     filename = ntpath.basename(file_path)

castor_extractor/utils/__init__.py

@@ -33,8 +33,10 @@ from .time import (
     current_datetime,
     current_timestamp,
     date_after,
+    format_date,
     past_date,
     timestamp_ms,
+    yesterday,
 )
 from .type import Callback, Getter, JsonType, SerializedAsset
 from .validation import validate_baseurl

castor_extractor/utils/argument_parser_test.py

@@ -4,7 +4,6 @@ from .argument_parser import parse_filled_arguments
 
 
 class MockArgumentParser:
-
     def __init__(self):
         self.attributes = {}
 

castor_extractor/utils/client/api.py

@@ -3,6 +3,9 @@ from typing import Any, Callable, Dict, Literal, Optional
 
 import requests
 
+from ...types import ExternalAsset
+from ..safe_request import RequestSafeMode, handle_response
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_TIMEOUT_S = 30
@@ -11,16 +14,31 @@ DEFAULT_TIMEOUT_S = 30
 HttpMethod = Literal["GET", "OPTIONS", "HEAD", "POST", "PUT", "PATCH", "DELETE"]
 
 
+def _authentication_header(token: Optional[str] = None) -> Dict[str, str]:
+    if token:
+        return {"Authorization": f"Bearer {token}"}
+    return dict()
+
+
 class APIClient:
     """
     API client
     - authentication via access token
     """
 
-    def __init__(self, host: str, token: Optional[str] = None):
+    def __init__(
+        self,
+        host: str,
+        token: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout: int = DEFAULT_TIMEOUT_S,
+        safe_mode: RequestSafeMode = RequestSafeMode(),
+    ):
         self._host = host
         self._token = token or ""
-        self._timeout = DEFAULT_TIMEOUT_S
+        self._timeout = timeout
+        self._base_headers = headers or {}
+        self.safe_mode = safe_mode
 
     @staticmethod
     def build_url(host: str, path: str):
@@ -28,10 +46,10 @@ class APIClient:
             host = "https://" + host
         return f"{host.strip('/')}/{path}"
 
-    def _headers(self) -> Dict[str, str]:
-        if self._token:
-            return {"Authorization": f"Bearer {self._token}"}
-        return dict()
+    @property
+    def _headers(self):
+        """Returns specified headers and authentication headers altogether"""
+        return {**self._base_headers, **_authentication_header(self._token)}
 
     def _call(
         self,
@@ -43,20 +61,20 @@ class APIClient:
         processor: Optional[Callable] = None,
     ) -> Any:
         logger.debug(f"Calling {method} on {url}")
-        result = requests.request(
+        response = requests.request(
             method,
             url,
-            headers=self._headers(),
+            headers=self._headers,
             params=params,
             json=data,
             timeout=self._timeout,
         )
-        result.raise_for_status()
+        response_payload = handle_response(response, self.safe_mode)
 
         if processor:
-            return processor(result)
+            return processor(response)
 
-        return result.json()
+        return response_payload
 
     def get(
         self,

castor_extractor/utils/client/api_test.py

@@ -1,4 +1,4 @@
-from .api import APIClient
+from .api import APIClient, _authentication_header
 
 
 def test_APIClient_build_url():
@@ -14,3 +14,11 @@ def test_APIClient_build_url():
 
     host_with_trailing_slash = "https://3.14.azuredatabricks.net/"
     assert expected == APIClient.build_url(host_with_trailing_slash, path)
+
+
+def test__authentication_header():
+    expected_headers = {"Authorization": "Bearer my_token"}
+    assert _authentication_header("my_token") == expected_headers
+
+    assert _authentication_header("") == {}
+    assert _authentication_header() == {}

castor_extractor/utils/object_test.py

@@ -19,7 +19,7 @@ class _User:
 
 
 def test_deep_serialize__None():
-    assert deep_serialize(None) == None
+    assert deep_serialize(None) is None
 
 
 def test_deep_serialize__str():

castor_extractor/utils/pager/pager.py

@@ -48,7 +48,7 @@ class AbstractPager(Generic[T]):
         return should_stop
 
 
-class Pager(AbstractPager):
+class Pager(AbstractPager[T]):
     def __init__(
         self,
         callback: Callable[[int, int], Sequence[T]],

castor_extractor/utils/pager/pager_on_id.py

@@ -1,11 +1,16 @@
-from typing import Callable, Iterator, Optional, Sequence, TypeVar
+from typing import Callable, Iterator, Optional, Protocol, Sequence, TypeVar
 from uuid import UUID
 
 from .pager import DEFAULT_PER_PAGE, AbstractPager, PagerStopStrategy
 
 _DEFAULT_MIN_UUID = UUID("00000000-0000-0000-0000-000000000000")
 
-T = TypeVar("T")
+
+class IndexableObject(Protocol):
+    def __getitem__(self, key: str) -> UUID: ...
+
+
+Indexable = TypeVar("Indexable", bound=IndexableObject)
 
 
 class PagerOnIdLogger:
@@ -16,10 +21,10 @@ class PagerOnIdLogger:
         pass
 
 
-class PagerOnId(AbstractPager):
+class PagerOnId(AbstractPager[Indexable]):
     def __init__(
         self,
-        callback: Callable[[UUID, int], Sequence[T]],
+        callback: Callable[[UUID, int], Sequence[Indexable]],
         *,
         logger: Optional[PagerOnIdLogger] = None,
         stop_strategy: PagerStopStrategy = PagerStopStrategy.EMPTY_PAGE,
@@ -29,13 +34,13 @@ class PagerOnId(AbstractPager):
         self._stop_strategy = stop_strategy
 
     @staticmethod
-    def _max_id(items: Sequence) -> UUID:
+    def _max_id(items: Sequence[Indexable]) -> UUID:
         return max(item["id"] for item in items)
 
     def iterator(
         self,
         per_page: int = DEFAULT_PER_PAGE,
-    ) -> Iterator[Sequence[T]]:
+    ) -> Iterator[Sequence[Indexable]]:
         """Yields data provided by the callback as a list using the greatest UUID as a reference point"""
         greater_than_id = _DEFAULT_MIN_UUID
         stop_on_empty_page = self._stop_strategy == PagerStopStrategy.EMPTY_PAGE

castor_extractor/utils/safe_request.py

@@ -1,5 +1,5 @@
 import logging
-from typing import List, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 from requests import HTTPError, Response
 
@@ -39,12 +39,14 @@ class RequestSafeMode:
 
 
 def handle_response(
-    response: Response, safe_mode: RequestSafeMode
-) -> ResponseJson:
+    response: Response,
+    safe_mode: Optional[RequestSafeMode] = None,
+) -> Any:
     """
     Util to handle a HTTP Response based on the response status code and the
     safe mode used
     """
+    safe_mode = safe_mode if safe_mode else RequestSafeMode()
     try:
         response.raise_for_status()
     except HTTPError as e:

castor_extractor/utils/safe_request_test.py

@@ -15,10 +15,8 @@ def mock_response(status_code: int):
 
 
 def test_http_error_with_no_safe_mode():
-    safe_params = RequestSafeMode()  # Caught
-
     with pytest.raises(HTTPError):
-        handle_response(mock_response(HTTPStatus.FORBIDDEN), safe_params)
+        handle_response(mock_response(HTTPStatus.FORBIDDEN))
 
 
 def test_http_error_with_no_status_code():

castor_extractor/utils/string_test.py

@@ -75,4 +75,4 @@ def test_decode_when_bytes():
     assert not decode_when_bytes(None)
     assert decode_when_bytes(1) == 1
     assert decode_when_bytes(1.34) == 1.34
-    assert decode_when_bytes(True) == True
+    assert decode_when_bytes(True)

castor_extractor/utils/time.py

@@ -1,4 +1,7 @@
 from datetime import date, datetime, timedelta, timezone
+from typing import Union
+
+ISO_FORMAT = "%Y-%m-%dT%H:%M:%S"
 
 
 def current_datetime() -> datetime:
@@ -54,3 +57,11 @@ def at_midnight(date_: date) -> datetime:
 def date_after(day: date, future_days: int) -> date:
     """returns the date `future_days` after `day`"""
     return day + timedelta(future_days)
+
+
+def format_date(timestamp: Union[datetime, date]) -> str:
+    return timestamp.strftime(ISO_FORMAT)
+
+
+def yesterday() -> date:
+    return current_date() - timedelta(days=1)

castor_extractor/visualization/domo/client/client.py

@@ -1,13 +1,12 @@
 import logging
 from datetime import datetime, timedelta
 from http import HTTPStatus
-from typing import Iterator, List, Optional, Set
+from typing import Any, Iterator, List, Optional, Set
 
 import requests
 
 from ....utils import (
     RequestSafeMode,
-    ResponseJson,
     at_midnight,
     current_date,
     empty_iterator,
@@ -108,7 +107,7 @@ class DomoClient:
         endpoint: Endpoint,
         params: Optional[dict] = None,
         asset_id: Optional[str] = None,
-    ) -> ResponseJson:
+    ) -> Any:
         params = params if params else {}
         is_private = endpoint.is_private
         headers = self._private_headers if is_private else self._bearer_auth()

castor_extractor/visualization/looker/api/client.py

@@ -22,6 +22,7 @@ from looker_sdk.sdk.api40.models import (
 from looker_sdk.sdk.constants import sdk_version
 
 from ....utils import Pager, PagerLogger, SafeMode, past_date, safe_mode
+from ..assets import LookerAsset
 from ..constants import DEFAULT_LOOKER_PAGE_SIZE
 from ..fields import format_fields
 from .constants import (
@@ -291,3 +292,37 @@ class ApiClient:
         self._on_api_call()
 
         return user_attributes
+
+    def fetch(
+        self,
+        asset: LookerAsset,
+        *,
+        folder_id: Optional[str] = None,
+        explore_names: Optional[Iterator[Tuple[str, str]]] = None,
+    ) -> list:
+        if asset == LookerAsset.USERS:
+            return self.users()
+        if asset == LookerAsset.CONNECTIONS:
+            return self.connections()
+        if asset == LookerAsset.LOOKS:
+            return self.looks(folder_id=folder_id)
+        if asset == LookerAsset.DASHBOARDS:
+            return self.dashboards(folder_id=folder_id)
+        if asset == LookerAsset.CONTENT_VIEWS:
+            return self.content_views()
+        if asset == LookerAsset.EXPLORES:
+            assert explore_names is not None
+            return list(self.explores(explore_names=explore_names))
+        if asset == LookerAsset.FOLDERS:
+            return self.folders()
+        if asset == LookerAsset.GROUPS_HIERARCHY:
+            return self.groups_hierarchy()
+        if asset == LookerAsset.GROUPS_ROLES:
+            return self.groups_roles()
+        if asset == LookerAsset.LOOKML_MODELS:
+            return self.lookml_models()
+        if asset == LookerAsset.PROJECTS:
+            return self.projects()
+        if asset == LookerAsset.USERS_ATTRIBUTES:
+            return self.users_attributes()
+        raise ValueError(f"Asset {asset.value} is not supported")

castor_extractor/visualization/looker/api/extraction_parameters.py

@@ -1,6 +1,7 @@
 from pydantic import Field, field_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
+from ....utils import OUTPUT_DIR
 from ..constants import (
     DEFAULT_LOOKER_PAGE_SIZE,
     DEFAULT_LOOKER_THREAD_POOL_SIZE,
@@ -24,7 +25,7 @@ class ExtractionParameters(BaseSettings):
 
     is_safe_mode: bool = False
     log_to_stdout: bool
-    output_directory: str
+    output: str = Field(validation_alias=OUTPUT_DIR)
     search_per_folder: bool
     page_size: int = Field(default=DEFAULT_LOOKER_PAGE_SIZE)
     thread_pool_size: int = Field(default=DEFAULT_LOOKER_THREAD_POOL_SIZE)

castor_extractor/visualization/looker/extract.py

@@ -39,7 +39,7 @@ def _safe_mode(
 ) -> Optional[SafeMode]:
     if extraction_parameters.is_safe_mode:
         return None
-    add_logging_file_handler(extraction_parameters.output_directory)
+    add_logging_file_handler(extraction_parameters.output)
     return SafeMode((Exception,), float("inf"))
 
 
@@ -127,7 +127,7 @@ def extract_all(**kwargs) -> None:
     output_directory
     """
     extraction_parameters = ExtractionParameters(**kwargs)
-    output_directory = extraction_parameters.output_directory
+    output_directory = extraction_parameters.output
 
     credentials = LookerCredentials(**kwargs)
 

castor_extractor/visualization/metabase/assets.py

@@ -1,3 +1,5 @@
+from typing import Dict, Tuple
+
 from ...types import ExternalAsset
 
 
@@ -13,7 +15,7 @@ class MetabaseAsset(ExternalAsset):
     DASHBOARD_CARDS = "dashboard_cards"
 
 
-EXPORTED_FIELDS = {
+EXPORTED_FIELDS: Dict[MetabaseAsset, Tuple[str, ...]] = {
     MetabaseAsset.COLLECTION: (
         "id",
         "name",

castor_extractor/visualization/metabase/extract.py

@@ -24,19 +24,31 @@ def iterate_all_data(
     """Iterate over the extracted Data From metabase"""
 
     yield MetabaseAsset.USER, deep_serialize(client.fetch(MetabaseAsset.USER))
-    yield MetabaseAsset.COLLECTION, deep_serialize(
-        client.fetch(MetabaseAsset.COLLECTION),
+    yield (
+        MetabaseAsset.COLLECTION,
+        deep_serialize(
+            client.fetch(MetabaseAsset.COLLECTION),
+        ),
     )
-    yield MetabaseAsset.DATABASE, deep_serialize(
-        client.fetch(MetabaseAsset.DATABASE),
+    yield (
+        MetabaseAsset.DATABASE,
+        deep_serialize(
+            client.fetch(MetabaseAsset.DATABASE),
+        ),
     )
     yield MetabaseAsset.TABLE, deep_serialize(client.fetch(MetabaseAsset.TABLE))
     yield MetabaseAsset.CARD, deep_serialize(client.fetch(MetabaseAsset.CARD))
-    yield MetabaseAsset.DASHBOARD, deep_serialize(
-        client.fetch(MetabaseAsset.DASHBOARD),
+    yield (
+        MetabaseAsset.DASHBOARD,
+        deep_serialize(
+            client.fetch(MetabaseAsset.DASHBOARD),
+        ),
     )
-    yield MetabaseAsset.DASHBOARD_CARDS, deep_serialize(
-        client.fetch(MetabaseAsset.DASHBOARD_CARDS),
+    yield (
+        MetabaseAsset.DASHBOARD_CARDS,
+        deep_serialize(
+            client.fetch(MetabaseAsset.DASHBOARD_CARDS),
+        ),
     )
 
 

castor_extractor/visualization/mode/client/client.py

@@ -220,6 +220,6 @@ class Client:
             # generic calls
             # example: https://modeanalytics.com/api/{workspace}/spaces
             # example: https://modeanalytics.com/api/{workspace}/data_sources
-            result = self._call(resource_name=asset.value)
+            result = self._call(resource_name=str(asset.value))
         logger.info(f"{len(result)} rows extracted")
         return self._post_processing(asset, result)

castor_extractor/visualization/powerbi/client/constants.py

@@ -1,5 +1,5 @@
 """
-    File regrouping all constants used in PowerBi client
+File regrouping all constants used in PowerBi client
 """
 
 DEFAULT_TIMEOUT_IN_SECS = 30

castor_extractor/visualization/powerbi/client/rest.py

@@ -1,12 +1,12 @@
 import logging
-from datetime import date, datetime, timedelta
+from datetime import date, datetime
 from time import sleep
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
 
 import msal  # type: ignore
 import requests
 
-from ....utils import at_midnight, current_date
+from ....utils import at_midnight, format_date, yesterday
 from ..assets import PowerBiAsset
 from .constants import (
     DEFAULT_TIMEOUT_IN_SECS,
@@ -23,19 +23,9 @@ from .utils import batch_size_is_valid_or_assert, datetime_is_recent_or_assert
 
 logger = logging.getLogger(__name__)
 
-_FORMAT = "%Y-%m-%dT%H:%M:%S"
-
-
-def _format(timestamp: datetime) -> str:
-    return timestamp.strftime(_FORMAT)
-
-
-def _yesterday() -> date:
-    return current_date() - timedelta(days=1)
-
 
 def _time_filter(day: Optional[date]) -> Tuple[datetime, datetime]:
-    target_day = day or _yesterday()
+    target_day = day or yesterday()
     start = at_midnight(target_day)
     end = datetime.combine(target_day, datetime.max.time())
     return start, end
@@ -51,8 +41,8 @@ def _url(
     url = Urls.ACTIVITY_EVENTS
     start, end = _time_filter(day)
     url += "?$filter=Activity eq 'viewreport'"
-    url += f"&startDateTime='{_format(start)}'"
-    url += f"&endDateTime='{_format(end)}'"
+    url += f"&startDateTime='{format_date(start)}'"
+    url += f"&endDateTime='{format_date(end)}'"
     return url
 
 

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

@@ -3,7 +3,12 @@ import logging
 from .....utils import SafeMode, safe_mode
 from .constants import MEASURES_SESSION_PARAMS, JsonRpcMethod
 from .credentials import QlikCredentials
-from .error import AccessDeniedError, AppSizeExceededError
+from .error import (
+    AccessDeniedError,
+    AppSizeExceededError,
+    PersistenceReadFailedError,
+    QlikResponseKeyError,
+)
 from .json_rpc import JsonRpcClient
 from .websocket import open_websocket
 
@@ -11,8 +16,29 @@ logger = logging.getLogger(__name__)
 
 
 def _handle(response: dict) -> int:
-    """Returns the object Handle from the response payload"""
-    return response["result"]["qReturn"]["qHandle"]
+    """
+    Returns the object Handle from the response payload, or raises an error
+    if one of the keys can't be found
+    """
+    try:
+        return response["result"]["qReturn"]["qHandle"]
+    except KeyError:
+        raise QlikResponseKeyError(
+            f"Could not fetch handle from response {response}"
+        )
+
+
+def _measure(response: dict) -> list:
+    """
+    Returns the measure from the response payload, or raises a custom error
+    if one of the keys can't be found
+    """
+    try:
+        return response["result"]["qLayout"]["qMeasureList"]["qItems"]
+    except KeyError:
+        raise QlikResponseKeyError(
+            f"Could not fetch measure from response {response}"
+        )
 
 
 def _list_measures(client: JsonRpcClient, app_id: str) -> list:
@@ -41,7 +67,7 @@ def _list_measures(client: JsonRpcClient, app_id: str) -> list:
         method=JsonRpcMethod.GET_LAYOUT,
         handle=session_handle,
     )
-    return response["result"]["qLayout"]["qMeasureList"]["qItems"]
+    return _measure(response)
 
 
 class EngineApiClient:
@@ -53,7 +79,12 @@ class EngineApiClient:
     def __init__(self, credentials: QlikCredentials):
         self.credentials = credentials
         self._safe_mode = SafeMode(
-            exceptions=(AccessDeniedError, AppSizeExceededError),
+            exceptions=(
+                AccessDeniedError,
+                AppSizeExceededError,
+                PersistenceReadFailedError,
+                QlikResponseKeyError,
+            ),
             max_errors=float("inf"),
         )
 

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

@@ -6,6 +6,7 @@ DEFAULT_HANDLE = -1
 
 ACCESS_DENIED_ERROR_CODE = 5
 APP_SIZE_EXCEEDED_ERROR_CODE = 1015
+PERSISTENCE_READ_FAILED_ERROR_CODE = 9001
 
 MEASURES_SESSION_PARAMS = (
     {

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

@@ -1,6 +1,10 @@
 from typing import Dict, Type
 
-from .constants import ACCESS_DENIED_ERROR_CODE, APP_SIZE_EXCEEDED_ERROR_CODE
+from .constants import (
+    ACCESS_DENIED_ERROR_CODE,
+    APP_SIZE_EXCEEDED_ERROR_CODE,
+    PERSISTENCE_READ_FAILED_ERROR_CODE,
+)
 
 
 class JsonRpcError(Exception):
@@ -32,9 +36,22 @@ class AppSizeExceededError(JsonRpcError):
     ...
 
 
+class PersistenceReadFailedError(JsonRpcError):
+    """
+    Error class to be raised when JSON-RPC error is persistence read failed
+    """
+
+
+class QlikResponseKeyError(Exception):
+    """
+    Error class to be raised when an expected Key is missing from the JSON-RPC response
+    """
+
+
 ERROR_CODE_MAPPING: Dict[int, Type[JsonRpcError]] = {
     ACCESS_DENIED_ERROR_CODE: AccessDeniedError,
     APP_SIZE_EXCEEDED_ERROR_CODE: AppSizeExceededError,
+    PERSISTENCE_READ_FAILED_ERROR_CODE: PersistenceReadFailedError,
 }
 
 

castor_extractor/visualization/salesforce_reporting/client/soql.py

@@ -1,6 +1,8 @@
+from typing import Dict
+
 from ..assets import SalesforceReportingAsset
 
-queries = {
+queries: Dict[SalesforceReportingAsset, str] = {
     SalesforceReportingAsset.DASHBOARDS: """
         SELECT
             CreatedBy.Id,