castor-extractor 0.17.3__py3-none-any.whl → 0.18.2__py3-none-any.whl

castor_extractor/visualization/metabase/errors.py

@@ -1,8 +1,10 @@
 from typing import Any
 
+from pydantic_settings import BaseSettings
+
 
 class MetabaseLoginError(ConnectionError):
-    def __init__(self, credentials_info: dict, error_details: Any):
+    def __init__(self, credentials_info: BaseSettings, error_details: Any):
         msg = "Connection to Metabase failed!\n"
         msg += f"Credentials: {credentials_info}\n"
         msg += f"Details: {error_details}\n"
@@ -10,7 +12,7 @@ class MetabaseLoginError(ConnectionError):
 
 
 class SuperuserCredentialsRequired(AssertionError):
-    def __init__(self, credentials_info: dict, error_details: Any):
+    def __init__(self, credentials_info: BaseSettings, error_details: Any):
         msg = "Superuser credentials are required!\n"
         msg += f"Credentials: {credentials_info}\n"
         msg += f"Details: {error_details}\n"
@@ -22,7 +24,7 @@ class EncryptionSecretKeyRequired(AssertionError):
     Raised when missing the encryption secret key while it was required.
     """
 
-    def __init__(self, credentials_info: dict, error_details: Any):
+    def __init__(self, credentials_info: BaseSettings, error_details: Any):
         msg = "Encryption secret key is required!\n"
         msg += f"Credentials: {credentials_info}\n"
         msg += f"Details: {error_details}\n"

castor_extractor/visualization/metabase/extract.py

@@ -55,6 +55,6 @@ def extract_all(client: ClientMetabase, **kwargs: str) -> None:
     write_summary(
         output_directory,
         ts,
-        base_url=client.base_url(),
+        base_url=client.base_url,
         client_name=client.name(),
     )

castor_extractor/visualization/mode/__init__.py

@@ -1,3 +1,3 @@
 from .assets import ModeAnalyticsAsset
-from .client import Client
+from .client import Client, ModeCredentials as Credentials
 from .extract import extract_all, iterate_all_data

castor_extractor/visualization/mode/client/__init__.py

@@ -1 +1,2 @@
 from .client import Client
+from .credentials import ModeCredentials

castor_extractor/visualization/mode/client/client.py

@@ -2,6 +2,7 @@ import logging
 from typing import Dict, List, Optional, cast
 
 import requests
+from requests.auth import HTTPBasicAuth
 
 from ....utils import retry
 from ..assets import (
@@ -22,7 +23,7 @@ from .constants import (
     RETRY_JITTER_MS,
     RETRY_STRATEGY,
 )
-from .credentials import Credentials, CredentialsKey, get_value
+from .credentials import ModeCredentials
 
 logger = logging.getLogger(__name__)
 
@@ -40,20 +41,16 @@ class Client:
 
     def __init__(
         self,
-        **kwargs,
+        credentials: ModeCredentials,
     ):
-        self._credentials = Credentials(
-            host=get_value(CredentialsKey.HOST, kwargs).rstrip("/"),
-            workspace=get_value(CredentialsKey.WORKSPACE, kwargs),
-            token=get_value(CredentialsKey.TOKEN, kwargs),
-            secret=get_value(CredentialsKey.SECRET, kwargs),
-        )
+        self._credentials = credentials
         self._session = requests.Session()
-        if not kwargs.get("no_checks"):
-            self._check_connection()
+
+    def authenticate(self) -> HTTPBasicAuth:
+        return HTTPBasicAuth(self._credentials.token, self._credentials.secret)
 
     def _check_connection(self):
-        authentication = self._credentials.authentication()
+        authentication = self.authenticate()
         url = self._url(with_workspace=False) + "/account"
         response = self._session.get(url, auth=authentication)
         self._handle_response(response)
@@ -121,7 +118,7 @@ class Client:
         report: Optional[str] = None,
         resource_name: Optional[str] = None,
     ) -> RawData:
-        authentication = self._credentials.authentication()
+        authentication = self.authenticate()
         url = self._url(with_workspace, space, report, resource_name)
         logger.info(f"Calling {url}")
         response = self._session.get(url, auth=authentication)

castor_extractor/visualization/mode/client/client_test.py

@@ -2,20 +2,20 @@ import json
 
 from ....utils import load_file
 from ..assets import EXPORTED_FIELDS, ModeAnalyticsAsset
-from .client import Client
+from .client import Client, ModeCredentials
 
 _HOST = "https://mode.com"
 _WORKSPACE = "castor"
 
 
 def _dummy_client() -> Client:
-    return Client(  # noqa: S106
+    credentials = ModeCredentials(
         host=_HOST,
         workspace=_WORKSPACE,
         token="dummy-token",
         secret="******",
-        no_checks=True,
     )
+    return Client(credentials=credentials)
 
 
 def test__url():

castor_extractor/visualization/mode/client/credentials.py

@@ -1,57 +1,24 @@
-from enum import Enum
-from typing import Dict
+from pydantic import Field, SecretStr, field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
-from requests.auth import HTTPBasicAuth
+MODE_ENV_PREFIX = "CASTOR_MODE_ANALYTICS_"
 
-from ....utils import from_env
 
+class ModeCredentials(BaseSettings):
+    """Class holding Mode credentials attributes"""
 
-class CredentialsKey(Enum):
-    HOST = "host"
-    WORKSPACE = "workspace"
-    TOKEN = "token"  # noqa: S105
-    SECRET = "secret"  # noqa: S105
+    model_config = SettingsConfigDict(
+        env_prefix=MODE_ENV_PREFIX,
+        extra="ignore",
+        populate_by_name=True,
+    )
 
+    host: str
+    secret: str = Field(repr=False)
+    token: str
+    workspace: str
 
-CREDENTIALS_ENV: Dict[CredentialsKey, str] = {
-    CredentialsKey.HOST: "CASTOR_MODE_ANALYTICS_HOST",
-    CredentialsKey.WORKSPACE: "CASTOR_MODE_ANALYTICS_WORKSPACE",
-    CredentialsKey.TOKEN: "CASTOR_MODE_ANALYTICS_TOKEN",
-    CredentialsKey.SECRET: "CASTOR_MODE_ANALYTICS_SECRET",
-}
-
-
-def get_value(key: CredentialsKey, kwargs: dict) -> str:
-    """
-    Returns the value of the given key:
-    - from kwargs in priority
-    - from ENV if not provided (raises an error if not found in ENV)
-    """
-    env_key = CREDENTIALS_ENV[key]
-    return str(kwargs.get(key.value) or from_env(env_key))
-
-
-class Credentials:
-    """ValueObject for the credentials"""
-
-    def __init__(self, host: str, workspace: str, token: str, secret: str):
-        self.host = host
-        self.workspace = workspace
-        self.token = token
-        self.secret = secret
-
-    def to_dict(self, hide: bool = False) -> Dict[str, str]:
-        safe = (
-            CredentialsKey.HOST,
-            CredentialsKey.WORKSPACE,
-            CredentialsKey.TOKEN,
-        )
-        unsafe = (CredentialsKey.SECRET,)
-
-        def val(k: CredentialsKey, v: str) -> str:
-            return "*" + v[-4:] if hide and k in unsafe else v
-
-        return {a.value: val(a, getattr(self, a.value)) for a in safe + unsafe}
-
-    def authentication(self) -> HTTPBasicAuth:
-        return HTTPBasicAuth(self.token, self.secret)
+    @field_validator("host", mode="before")
+    @classmethod
+    def _check_base_url(cls, host: str) -> str:
+        return host.rstrip("/")

castor_extractor/visualization/mode/extract.py

@@ -11,7 +11,7 @@ from ...utils import (
     write_summary,
 )
 from .assets import ModeAnalyticsAsset as Asset
-from .client import Client
+from .client import Client, ModeCredentials
 
 logger = logging.getLogger(__name__)
 
@@ -37,9 +37,11 @@ def iterate_all_data(
     yield Asset.MEMBER, deep_serialize(members)
 
 
-def extract_all(client: Client, **kwargs: str) -> None:
+def extract_all(credentials: ModeCredentials, output_directory: str) -> None:
     """Extract Data From Mode Analytics and store it locally in files under the output_directory"""
-    output_directory = kwargs.get("output_directory") or from_env(OUTPUT_DIR)
+    output_directory = output_directory or from_env(OUTPUT_DIR)
+
+    client = Client(credentials)
     ts = current_timestamp()
 
     for key, data in iterate_all_data(client):

castor_extractor/visualization/powerbi/__init__.py

@@ -1,3 +1,3 @@
 from .assets import PowerBiAsset
-from .client import Client, Credentials
+from .client import Client, PowerbiCredentials, Urls
 from .extract import extract_all

castor_extractor/visualization/powerbi/client/__init__.py

@@ -1,2 +1,3 @@
-from .credentials import Credentials
+from .constants import Urls
+from .credentials import PowerbiCredentials
 from .rest import Client

castor_extractor/visualization/powerbi/client/credentials.py

@@ -1,20 +1,28 @@
-from dataclasses import field
 from typing import List, Optional
 
-from pydantic.dataclasses import dataclass
+from pydantic import Field, field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from .constants import Urls
 
+POWERBI_ENV_PREFIX = "CASTOR_POWERBI_"
 
-@dataclass
-class Credentials:
+
+class PowerbiCredentials(BaseSettings):
     """Class to handle PowerBI rest API permissions"""
 
+    model_config = SettingsConfigDict(
+        env_prefix=POWERBI_ENV_PREFIX,
+        extra="ignore",
+        populate_by_name=True,
+    )
+
     client_id: str
     tenant_id: str
-    secret: str = field(metadata={"sensitive": True})
-    scopes: Optional[List[str]] = None
+    secret: str = Field(repr=False)
+    scopes: List[str] = [Urls.DEFAULT_SCOPE]
 
-    def __post_init__(self):
-        if self.scopes is None:
-            self.scopes = [Urls.DEFAULT_SCOPE]
+    @field_validator("scopes", mode="before")
+    @classmethod
+    def _check_scopes(cls, scopes: Optional[List[str]]) -> List[str]:
+        return scopes if scopes is not None else [Urls.DEFAULT_SCOPE]

castor_extractor/visualization/powerbi/client/credentials_test.py

@@ -1,5 +1,5 @@
 from .constants import Urls
-from .credentials import Credentials
+from .credentials import PowerbiCredentials
 
 
 def test_credentials():
@@ -8,15 +8,23 @@ def test_credentials():
     secret = "🤫"
 
     # no scopes provided
-    credentials = Credentials(
+    credentials = PowerbiCredentials(
         tenant_id=tenant_id,
         client_id=client_id,
         secret=secret,
     )
     assert credentials.scopes == [Urls.DEFAULT_SCOPE]
 
+    credentials = PowerbiCredentials(
+        tenant_id=tenant_id,
+        client_id=client_id,
+        secret=secret,
+        scopes=None,
+    )
+    assert credentials.scopes == [Urls.DEFAULT_SCOPE]
+
     # empty scopes
-    credentials = Credentials(
+    credentials = PowerbiCredentials(
         tenant_id=tenant_id,
         client_id=client_id,
         secret=secret,
@@ -26,7 +34,7 @@ def test_credentials():
 
     # with scopes
     scopes = ["foo"]
-    credentials = Credentials(
+    credentials = PowerbiCredentials(
         tenant_id=tenant_id,
         client_id=client_id,
         secret=secret,

castor_extractor/visualization/powerbi/client/rest.py

@@ -18,7 +18,7 @@ from .constants import (
     QueryParams,
     Urls,
 )
-from .credentials import Credentials
+from .credentials import PowerbiCredentials
 from .utils import batch_size_is_valid_or_assert, datetime_is_recent_or_assert
 
 logger = logging.getLogger(__name__)
@@ -62,7 +62,7 @@ class Client:
     https://learn.microsoft.com/en-us/rest/api/power-bi/admin
     """
 
-    def __init__(self, credentials: Credentials):
+    def __init__(self, credentials: PowerbiCredentials):
         self.creds = credentials
         client_app = f"{Urls.CLIENT_APP_BASE}{self.creds.tenant_id}"
         self.app = msal.ConfidentialClientApplication(

castor_extractor/visualization/powerbi/client/rest_test.py

@@ -5,7 +5,7 @@ import pytest
 from requests import HTTPError
 
 from .constants import GET, POST, Assertions, Keys, QueryParams, Urls
-from .credentials import Credentials
+from .credentials import PowerbiCredentials
 from .rest import Client, msal
 
 FAKE_TENANT_ID = "IamFake"
@@ -14,7 +14,7 @@ FAKE_SECRET = "MeThree"
 
 
 def _client() -> Client:
-    creds = Credentials(
+    creds = PowerbiCredentials(
         tenant_id=FAKE_TENANT_ID,
         client_id=FAKE_CLIENT_ID,
         secret=FAKE_SECRET,

castor_extractor/visualization/powerbi/extract.py

@@ -10,7 +10,7 @@ from ...utils import (
     write_summary,
 )
 from .assets import METADATA_ASSETS, PowerBiAsset
-from .client import Client, Credentials
+from .client import Client, PowerbiCredentials
 
 
 def iterate_all_data(
@@ -36,7 +36,7 @@ def extract_all(
     Store the output files locally under the given output_directory
     """
     _output_directory = output_directory or from_env(OUTPUT_DIR)
-    creds = Credentials(
+    creds = PowerbiCredentials(
         tenant_id=tenant_id,
         client_id=client_id,
         secret=secret,

castor_extractor/visualization/qlik/__init__.py

@@ -1,3 +1,7 @@
 from .assets import QlikAsset
-from .client import APP_EXTERNAL_ID_KEY as QLIK_APP_EXTERNAL_ID_KEY, QlikClient
+from .client import (
+    APP_EXTERNAL_ID_KEY as QLIK_APP_EXTERNAL_ID_KEY,
+    QlikClient,
+    QlikCredentials,
+)
 from .extract import extract_all

castor_extractor/visualization/qlik/client/__init__.py

@@ -1,2 +1,3 @@
 from .constants import APP_EXTERNAL_ID_KEY
+from .engine import QlikCredentials
 from .master import QlikMasterClient as QlikClient

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

@@ -1 +1,2 @@
 from .client import EngineApiClient
+from .credentials import QlikCredentials

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

@@ -2,6 +2,7 @@ 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 .json_rpc import JsonRpcClient
 from .websocket import open_websocket
@@ -49,10 +50,8 @@ class EngineApiClient:
     get measures using JsonRpcClient and websocket connection.
     """
 
-    def __init__(self, server_url: str, api_key: str):
-        self.server_url = server_url
-        self.api_key = api_key
-
+    def __init__(self, credentials: QlikCredentials):
+        self.credentials = credentials
         self._safe_mode = SafeMode(
             exceptions=(AccessDeniedError, AppSizeExceededError),
             max_errors=float("inf"),
@@ -70,8 +69,8 @@ class EngineApiClient:
 
         with open_websocket(
             app_id=app_id,
-            server_url=self.server_url,
-            api_key=self.api_key,
+            server_url=self.credentials.base_url,
+            api_key=self.credentials.api_key,
         ) as websocket:
             json_rpc_client = JsonRpcClient(websocket=websocket)
             return _call(json_rpc_client, app_id)

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

@@ -0,0 +1,26 @@
+from pydantic import Field, SecretStr, field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from .....utils import validate_baseurl
+
+QLIK_ENV_PREFIX = "CASTOR_QLIK_"
+
+
+class QlikCredentials(BaseSettings):
+    """
+    Qlik's credentials to connect to the API
+    """
+
+    model_config = SettingsConfigDict(
+        env_prefix=QLIK_ENV_PREFIX,
+        extra="ignore",
+        populate_by_name=True,
+    )
+
+    api_key: str = Field(repr=False)
+    base_url: str
+
+    @field_validator("base_url", mode="before")
+    @classmethod
+    def _check_base_url(cls, base_url: str) -> str:
+        return validate_baseurl(base_url)

castor_extractor/visualization/qlik/client/master.py

@@ -4,7 +4,7 @@ from tqdm import tqdm  # type: ignore
 
 from ..assets import QlikAsset
 from .constants import APP_EXTERNAL_ID_KEY, SCOPED_ASSETS
-from .engine import EngineApiClient
+from .engine import EngineApiClient, QlikCredentials
 from .rest import RestApiClient
 
 ListedData = List[dict]
@@ -53,25 +53,19 @@ class QlikMasterClient:
 
     def __init__(
         self,
-        server_url: str,
-        api_key: str,
+        credentials: QlikCredentials,
         except_http_error_statuses: Optional[List[int]] = None,
         display_progress: bool = True,
     ):
-        self._server_url = server_url
-        self._api_key = api_key
+        self._server_url = credentials.base_url
         self.display_progress = display_progress
 
         self.rest_api_client = RestApiClient(
-            server_url=self._server_url,
-            api_key=self._api_key,
+            credentials=credentials,
             except_http_error_statuses=except_http_error_statuses,
         )
 
-        self.engine_api_client = EngineApiClient(
-            server_url=self._server_url,
-            api_key=self._api_key,
-        )
+        self.engine_api_client = EngineApiClient(credentials=credentials)
 
     def _fetch_lineage(self, apps: ListedData) -> ListedData:
         callback = self.rest_api_client.data_lineage

castor_extractor/visualization/qlik/client/rest.py

@@ -16,6 +16,7 @@ from .constants import (
     RETRY_COUNTS,
     RETRY_STATUSES,
 )
+from .engine import QlikCredentials
 
 logger = logging.getLogger(__name__)
 
@@ -60,12 +61,11 @@ class RestApiClient:
 
     def __init__(
         self,
-        server_url: str,
-        api_key: str,
+        credentials: QlikCredentials,
         except_http_error_statuses: Optional[List[int]] = None,
     ):
-        self._server_url = server_url
-        self._api_key = api_key
+        self._server_url = credentials.base_url
+        self._api_key = credentials.api_key
         self._session = _session()
         self._except_http_error_statuses = except_http_error_statuses or []
         self._authenticate()

castor_extractor/visualization/qlik/client/rest_test.py

@@ -1,6 +1,7 @@
 from typing import Optional
 from unittest.mock import call, patch
 
+from .engine import QlikCredentials
 from .rest import RestApiClient
 
 
@@ -22,8 +23,11 @@ def _check_called_once(
 def test_rest_api_client_pager():
     dummy_server_url = "https://clic.kom"
     dummy_api_key = "i-am-the-key-dont-let-others-know-about"
-
-    client = RestApiClient(server_url=dummy_server_url, api_key=dummy_api_key)
+    credentials = QlikCredentials(
+        base_url=dummy_server_url,
+        api_key=dummy_api_key,
+    )
+    client = RestApiClient(credentials=credentials)
 
     first_page_url = "https://clic.kom/assets"
 

castor_extractor/visualization/qlik/extract.py

@@ -7,13 +7,11 @@ from ...utils import (
     deep_serialize,
     from_env,
     get_output_filename,
-    validate_baseurl,
     write_json,
     write_summary,
 )
 from .assets import QlikAsset
-from .client import QlikClient
-from .constants import API_KEY, BASE_URL
+from .client import QlikClient, QlikCredentials
 
 logger = logging.getLogger(__name__)
 
@@ -59,13 +57,11 @@ def extract_all(
     Store the output files locally under the given output_directory
     """
 
+    credentials = QlikCredentials(base_url=base_url, api_key=api_key)
     _output_directory = output_directory or from_env(OUTPUT_DIR)
-    _base_url = validate_baseurl(base_url or from_env(BASE_URL))
-    _api_key = api_key or from_env(API_KEY)
 
     client = QlikClient(
-        server_url=_base_url,
-        api_key=_api_key,
+        credentials=credentials,
         except_http_error_statuses=except_http_error_statuses,
     )
 
@@ -75,4 +71,4 @@ def extract_all(
         filename = get_output_filename(key.name.lower(), _output_directory, ts)
         write_json(filename, data)
 
-    write_summary(_output_directory, ts, base_url=_base_url)
+    write_summary(_output_directory, ts, base_url=credentials.base_url)

castor_extractor/visualization/sigma/__init__.py

@@ -1,3 +1,3 @@
 from .assets import SigmaAsset
-from .client import CredentialsKey, SigmaClient, SigmaCredentials
+from .client import SigmaClient, SigmaCredentials
 from .extract import extract_all

castor_extractor/visualization/sigma/client/__init__.py

@@ -1,2 +1,2 @@
 from .client import SigmaClient
-from .credentials import CredentialsKey, SigmaCredentials
+from .credentials import SigmaCredentials

castor_extractor/visualization/sigma/client/client.py

@@ -5,14 +5,20 @@ from urllib.parse import urljoin
 import requests
 
 from ..assets import SigmaAsset
-from .credentials import CredentialsKey, SigmaCredentials
+from .credentials import SigmaCredentials
 from .endpoints import EndpointFactory
 from .pagination import Pagination
 
 logger = logging.getLogger()
 
 
-DATA_ELEMENTS: Tuple[str, str] = ("table", "visualization")
+DATA_ELEMENTS: Tuple[str, ...] = (
+    "input-table",
+    "pivot-table",
+    "table",
+    "visualization",
+    "viz",
+)
 _CONTENT_TYPE = "application/x-www-form-urlencoded"
 
 
@@ -23,6 +29,7 @@ class SigmaClient:
         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]:
@@ -31,9 +38,9 @@ class SigmaClient:
         token_response = requests.post(  # noqa: S113
             token_api_path,
             data={
-                CredentialsKey.GRANT_TYPE.value: "client_credentials",
-                CredentialsKey.CLIENT_ID.value: self.client_id,
-                CredentialsKey.CLIENT_SECRET.value: self.api_token,
+                "grant_type": self.grant_type,
+                "client_id": self.client_id,
+                "client_secret": self.api_token,
             },
         )
         if token_response.status_code != requests.codes.OK: