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

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

@@ -1,4 +1,4 @@
-from typing import Iterable, List, Optional, Tuple, Union
+from typing import Iterable, List, Tuple, Union
 
 from ...utils import (
     OUTPUT_DIR,
@@ -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(
@@ -24,24 +24,13 @@ def iterate_all_data(
         yield asset, deep_serialize(data)
 
 
-def extract_all(
-    tenant_id: str,
-    client_id: str,
-    secret: str,
-    scopes: Optional[List[str]] = None,
-    output_directory: Optional[str] = None,
-) -> None:
+def extract_all(**kwargs) -> None:
     """
     Extract data from PowerBI REST API
     Store the output files locally under the given output_directory
     """
-    _output_directory = output_directory or from_env(OUTPUT_DIR)
-    creds = Credentials(
-        tenant_id=tenant_id,
-        client_id=client_id,
-        secret=secret,
-        scopes=scopes,
-    )
+    _output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
+    creds = PowerbiCredentials(**kwargs)
     client = Client(creds)
     ts = current_timestamp()
 

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__)
 
@@ -49,23 +47,18 @@ def iterate_all_data(
 
 
 def extract_all(
-    base_url: Optional[str] = None,
-    api_key: Optional[str] = None,
-    output_directory: Optional[str] = None,
-    except_http_error_statuses: Optional[List[int]] = None,
+    except_http_error_statuses: Optional[List[int]] = None, **kwargs
 ) -> None:
     """
     Extract data from Qlik REST API
     Store the output files locally under the given output_directory
     """
 
-    _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)
+    credentials = QlikCredentials(**kwargs)
+    _output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
 
     client = QlikClient(
-        server_url=_base_url,
-        api_key=_api_key,
+        credentials=credentials,
         except_http_error_statuses=except_http_error_statuses,
     )
 
@@ -75,4 +68,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/salesforce_reporting/extract.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Iterable, Optional, Tuple, Union
+from typing import Iterable, Tuple, Union
 
 from ...utils import (
     OUTPUT_DIR,
@@ -28,29 +28,15 @@ def iterate_all_data(
         yield asset.name.lower(), deep_serialize(data)
 
 
-def extract_all(
-    username: str,
-    password: str,
-    client_id: str,
-    client_secret: str,
-    security_token: str,
-    base_url: str,
-    output_directory: Optional[str] = None,
-) -> None:
+def extract_all(**kwargs) -> None:
     """
     Extract data from Salesforce REST API
     Store the output files locally under the given output_directory
     """
-    _output_directory = output_directory or from_env(OUTPUT_DIR)
-    creds = SalesforceCredentials(
-        username=username,
-        password=password,
-        client_id=client_id,
-        client_secret=client_secret,
-        security_token=security_token,
-        base_url=base_url,
-    )
-    client = SalesforceReportingClient(credentials=creds)
+    _output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
+    credentials = SalesforceCredentials(**kwargs)
+
+    client = SalesforceReportingClient(credentials=credentials)
     ts = current_timestamp()
 
     for key, data in iterate_all_data(client):

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,7 +5,7 @@ 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
 
@@ -29,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]:
@@ -37,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:

castor_extractor/visualization/sigma/client/credentials.py

@@ -1,35 +1,19 @@
-from dataclasses import dataclass
-from enum import Enum
+from pydantic import Field, SecretStr
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
+CASTOR_ENV_PREFIX = "CASTOR_SIGMA_"
 
-class CredentialsKey(Enum):
-    """Value enum object for the credentials"""
 
-    CLIENT_SECRET = "client_secret"  # noqa: S105
-    CLIENT_ID = "client_id"
-    HOST = "host"
-    GRANT_TYPE = "grant_type"
-    API_TOKEN = "api_token"  # noqa: S105
-
-
-CLIENT_ALLOWED_KEYS = (
-    CredentialsKey.HOST.value,
-    CredentialsKey.CLIENT_ID.value,
-    CredentialsKey.API_TOKEN.value,
-)
-
-
-@dataclass
-class SigmaCredentials:
+class SigmaCredentials(BaseSettings):
     """Class to handle Sigma rest API permissions"""
 
-    api_token: str
+    model_config = SettingsConfigDict(
+        env_prefix=CASTOR_ENV_PREFIX,
+        extra="ignore",
+        populate_by_name=True,
+    )
+
+    api_token: str = Field(repr=False)
     client_id: str
     host: str
-
-    @classmethod
-    def from_secret(cls, secret: dict) -> "SigmaCredentials":
-        credentials = {
-            k: v for k, v in secret.items() if k in CLIENT_ALLOWED_KEYS
-        }
-        return cls(**credentials)
+    grant_type: str = "client_credentials"

castor_extractor/visualization/sigma/extract.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Iterable, Iterator, Optional, Tuple, Union
+from typing import Iterable, Iterator, Tuple, Union
 
 from ...utils import (
     OUTPUT_DIR,
@@ -12,7 +12,6 @@ from ...utils import (
 )
 from .assets import SigmaAsset
 from .client import SigmaClient, SigmaCredentials
-from .constants import API_TOKEN, CLIENT_ID, HOST
 
 logger = logging.getLogger(__name__)
 
@@ -51,27 +50,15 @@ def iterate_all_data(
     yield SigmaAsset.LINEAGES, list(deep_serialize(lineages))
 
 
-def extract_all(
-    host: Optional[str] = None,
-    client_id: Optional[str] = None,
-    api_token: Optional[str] = None,
-    output_directory: Optional[str] = None,
-) -> None:
+def extract_all(**kwargs) -> None:
     """
     Extract data from Sigma API
     Store the output files locally under the given output_directory
     """
 
-    _output_directory = output_directory or from_env(OUTPUT_DIR)
-    _client_id = client_id or from_env(CLIENT_ID)
-    _host = host or from_env(HOST)
-    _api_token = api_token or from_env(API_TOKEN)
+    _output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
 
-    credentials = SigmaCredentials(
-        host=_host,
-        client_id=_client_id,
-        api_token=_api_token,
-    )
+    credentials = SigmaCredentials(**kwargs)
     client = SigmaClient(credentials=credentials)
 
     ts = current_timestamp()
@@ -80,4 +67,4 @@ def extract_all(
         filename = get_output_filename(key.name.lower(), _output_directory, ts)
         write_json(filename, data)
 
-    write_summary(_output_directory, ts, host=_host)
+    write_summary(_output_directory, ts, host=credentials.host)

castor_extractor/visualization/tableau_revamp/client/credentials.py

@@ -1,104 +1,57 @@
-from enum import Enum
-from typing import Dict, Literal, Optional, overload
+from typing import Optional
 
-from ....utils import from_env
-
-_AUTH_ERROR_MSG = "Need either user and password or token_name and token"
+from pydantic import field_validator, model_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
 # To specify the default site on Tableau Server, you can use an empty string
 # https://tableau.github.io/server-client-python/docs/api-ref#authentication
 _DEFAULT_SERVER_SITE_ID = ""
 
-
-class CredentialsKey(Enum):
-    """Value enum object for the credentials"""
-
-    TABLEAU_USER = "user"
-    TABLEAU_PASSWORD = "password"  # noqa: S105
-    TABLEAU_TOKEN_NAME = "token_name"  # noqa: S105
-    TABLEAU_TOKEN = "token"  # noqa: S105
-    TABLEAU_SITE_ID = "site_id"
-    TABLEAU_SERVER_URL = "server_url"
-
-
-CREDENTIALS_ENV: Dict[CredentialsKey, str] = {
-    CredentialsKey.TABLEAU_USER: "CASTOR_TABLEAU_USER",
-    CredentialsKey.TABLEAU_PASSWORD: "CASTOR_TABLEAU_PASSWORD",
-    CredentialsKey.TABLEAU_TOKEN_NAME: "CASTOR_TABLEAU_TOKEN_NAME",
-    CredentialsKey.TABLEAU_TOKEN: "CASTOR_TABLEAU_TOKEN",
-    CredentialsKey.TABLEAU_SITE_ID: "CASTOR_TABLEAU_SITE_ID",
-    CredentialsKey.TABLEAU_SERVER_URL: "CASTOR_TABLEAU_SERVER_URL",
-}
-
-
-@overload
-def get_value(key: CredentialsKey, kwargs: dict) -> Optional[str]: ...
+# In Castor APP, site_id is mandatory: users can't let this field empty
+# In that case, we encourage users to write "Default" instead
+_DEFAULT_SITE_ID_USER_INPUT = "default"
 
 
-@overload
-def get_value(
-    key: CredentialsKey, kwargs: dict, optional: Literal[True]
-) -> Optional[str]: ...
+TABLEAU_ENV_PREFIX = "CASTOR_TABLEAU_"
 
 
-@overload
-def get_value(
-    key: CredentialsKey, kwargs: dict, optional: Literal[False]
-) -> str: ...
-
-
-def get_value(
-    key: CredentialsKey,
-    kwargs: dict,
-    optional: bool = True,
-) -> Optional[str]:
+class TableauRevampCredentials(BaseSettings):
     """
-    Returns the value of the given key:
-    - from kwargs in priority
-    - from ENV otherwise
-    Raises an error if not found (unless optional)
+    Tableau's credentials to connect to both APIs (REST and GRAPHQL)
     """
 
-    if key.value in kwargs:
-        return kwargs[key.value]
-
-    env_key = CREDENTIALS_ENV[key]
-    return from_env(env_key, optional)
+    model_config = SettingsConfigDict(
+        env_prefix=TABLEAU_ENV_PREFIX,
+        extra="ignore",
+        populate_by_name=True,
+    )
 
+    server_url: str
+    site_id: str = ""
 
-class TableauRevampCredentials:
-    """
-    Tableau's credentials to connect to REST API
-    """
-
-    def __init__(
-        self,
-        *,
-        server_url: str,
-        site_id: Optional[str],
-        user: Optional[str],
-        password: Optional[str],
-        token_name: Optional[str],
-        token: Optional[str],
-    ):
-        self.user = user
-        self.site_id = site_id or _DEFAULT_SERVER_SITE_ID
-        self.server_url = server_url
-        self.password = password
-        self.token_name = token_name
-        self.token = token
+    password: Optional[str] = None
+    token: Optional[str] = None
+    token_name: Optional[str] = None
+    user: Optional[str] = None
 
+    @field_validator("site_id", mode="before")
     @classmethod
-    def from_env(cls, kwargs: dict) -> "TableauRevampCredentials":
-        return TableauRevampCredentials(
-            server_url=get_value(
-                CredentialsKey.TABLEAU_SERVER_URL,
-                kwargs,
-                optional=False,
-            ),
-            site_id=get_value(CredentialsKey.TABLEAU_SITE_ID, kwargs),
-            user=get_value(CredentialsKey.TABLEAU_USER, kwargs),
-            password=get_value(CredentialsKey.TABLEAU_PASSWORD, kwargs),
-            token_name=get_value(CredentialsKey.TABLEAU_TOKEN_NAME, kwargs),
-            token=get_value(CredentialsKey.TABLEAU_TOKEN, kwargs),
-        )
+    def _check_site_id(cls, site_id: Optional[str]) -> str:
+        if not site_id or site_id.lower() == _DEFAULT_SITE_ID_USER_INPUT:
+            return _DEFAULT_SERVER_SITE_ID
+        return site_id
+
+    @model_validator(mode="after")
+    def _check_user_xor_pat_login(self) -> "TableauRevampCredentials":
+        """
+        Checks that credentials are correctly input, it means either:
+        - User and password are filled
+        - Token and Token name are filled
+        """
+        user_login = self.password and self.user
+        pat_login = self.token_name and self.token
+        if not user_login and not pat_login:
+            raise ValueError("Either token or user identification is required")
+        if user_login and pat_login:
+            raise ValueError("Can't have both token and user identification")
+        return self

castor_extractor/warehouse/databricks/client.py

@@ -28,6 +28,7 @@ from .utils import build_path, tag_label
 
 logger = logging.getLogger(__name__)
 
+_DATABRICKS_CLIENT_TIMEOUT = 60
 _MAX_NUMBER_OF_LINEAGE_ERRORS = 1000
 _MAX_THREADS = 10
 _RETRY_ATTEMPTS = 3
@@ -98,6 +99,8 @@ class DatabricksClient(APIClient):
         self._db_blocked = db_blocked
         self._has_table_tags = has_table_tags
         self._has_column_tags = has_column_tags
+
+        self._timeout = _DATABRICKS_CLIENT_TIMEOUT
         self.formatter = DatabricksFormatter()
 
     def execute_sql(

castor_extractor/warehouse/redshift/queries/column.sql

@@ -28,7 +28,6 @@ information_tables AS (
         i.table_id || '.' || c.column_name AS column_id,
         c.data_type,
         c.ordinal_position,
-        c.column_default,
         c.is_nullable,
         c.character_maximum_length,
         c.character_octet_length,
@@ -59,7 +58,6 @@ raw_tables AS (
         a.attname AS column_name,
         c.oid::TEXT || '.' || a.attname AS column_id,
         a.attnum AS ordinal_position,
-        ad.adsrc AS column_default,
         CASE
             WHEN t.typname = 'bpchar' THEN 'char'
             ELSE t.typname
@@ -90,7 +88,6 @@ tables AS (
         COALESCE(i.data_type, r.data_type) AS data_type,
         COALESCE(i.ordinal_position, r.ordinal_position) AS ordinal_position,
         COALESCE(i.is_nullable, r.is_nullable) AS is_nullable,
-        COALESCE(i.column_default, r.column_default) AS column_default,
         i.character_maximum_length::INT AS character_maximum_length,
         i.character_octet_length::INT AS character_octet_length,
         i.numeric_precision::INT AS numeric_precision,
@@ -117,7 +114,6 @@ views_late_binding AS (
         c.data_type,
         c.ordinal_position,
         'YES' AS is_nullable,
-        NULL::TEXT AS column_default,
         NULL::INT AS character_maximum_length,
         NULL::INT AS character_octet_length,
         NULL::INT AS numeric_precision,
@@ -162,7 +158,6 @@ external_columns AS (
         c.external_type AS data_type,
         MIN(c.columnnum) AS ordinal_position,
         CASE c.is_nullable WHEN 'false' THEN 'NO' ELSE 'YES' END AS is_nullable,
-        NULL AS column_default,
         NULL AS character_maximum_length,
         NULL AS character_octet_length,
         NULL AS numeric_precision,