castor-extractor 0.21.7__py3-none-any.whl → 0.22.0__py3-none-any.whl

castor_extractor/visualization/tableau/assets.py

@@ -1,5 +1,4 @@
 from enum import Enum
-from typing import Set
 
 from ...types import ExternalAsset, classproperty
 
@@ -24,7 +23,7 @@ class TableauAsset(ExternalAsset):
     WORKBOOK_TO_DATASOURCE = "workbooks_to_datasource"
 
     @classproperty
-    def optional(cls) -> Set["TableauAsset"]:
+    def optional(cls) -> set["TableauAsset"]:
         return {
             TableauAsset.DASHBOARD,
             TableauAsset.DASHBOARD_SHEET,

castor_extractor/visualization/tableau/client/client.py

@@ -1,5 +1,4 @@
 import logging
-from typing import List
 
 import tableauserverclient as TSC  # type: ignore
 
@@ -43,7 +42,7 @@ class ApiClient:
         self._page_size = PAGE_SIZE
         self._server.version = TABLEAU_SERVER_VERSION
         self._safe_mode = bool(kwargs.get("safe_mode"))
-        self.errors: List[str] = []
+        self.errors: list[str] = []
 
     @staticmethod
     def name() -> str:

castor_extractor/visualization/tableau/client/client_utils.py

@@ -1,4 +1,5 @@
-from typing import Dict, Iterator, Optional
+from collections.abc import Iterator
+from typing import Optional
 
 from ....utils import SerializedAsset
 from ..assets import TableauAsset
@@ -69,6 +70,6 @@ def query_scroll(
             break
 
 
-def extract_asset(asset: Dict, asset_type: TableauAsset) -> Dict:
+def extract_asset(asset: dict, asset_type: TableauAsset) -> dict:
     """Agnostic function extracting dedicated attributes with define asset"""
     return {key: getattr(asset, key) for key in TSC_FIELDS[asset_type]}

castor_extractor/visualization/tableau/client/credentials.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Dict, Optional
+from typing import Optional
 
 from ....utils import from_env
 
@@ -20,7 +20,7 @@ class CredentialsKey(Enum):
     TABLEAU_SERVER_URL = "server_url"
 
 
-CREDENTIALS_ENV: Dict[CredentialsKey, str] = {
+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",
@@ -89,7 +89,7 @@ class CredentialsApi:
             CredentialsKey.TABLEAU_TOKEN: token,
         }
 
-    def to_dict(self, hide: bool = False) -> Dict[str, str]:
+    def to_dict(self, hide: bool = False) -> dict[str, str]:
         safe = (
             CredentialsKey.TABLEAU_USER,
             CredentialsKey.TABLEAU_SITE_ID,

castor_extractor/visualization/tableau/client/safe_mode.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Dict, List
 
 import tableauserverclient as TSC  # type: ignore
 
@@ -48,7 +47,7 @@ def safe_mode_fetch_usage(client) -> SerializedAsset:
     Returns computed usages when page number is not found
     Log errors if ServerResponseError is return
     """
-    list_usages: List[Dict] = []
+    list_usages: list[dict] = []
     page_number: int = 0
 
     while True:

castor_extractor/visualization/tableau/extract.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Iterable, Tuple
+from collections.abc import Iterable
 
 from ...utils import (
     OUTPUT_DIR,
@@ -19,7 +19,7 @@ logger = logging.getLogger(__name__)
 
 def iterate_all_data(
     client: Client,
-) -> Iterable[Tuple[TableauAsset, list]]:
+) -> Iterable[tuple[TableauAsset, list]]:
     """Iterate over the extracted Data from Tableau"""
 
     logger.info("Extracting USER from Tableau API")

castor_extractor/visualization/tableau/gql_fields.py

@@ -1,6 +1,6 @@
 # Fields which will be use for Tableau GraphQL API
 from enum import Enum
-from typing import Dict, List, Union
+from typing import Union
 
 from .assets import TableauAsset, TableauGraphqlAsset
 
@@ -189,9 +189,9 @@ class GQLQueryFields(Enum):
     """
 
 
-QueryInfo = List[Dict[str, Union[GQLQueryFields, TableauGraphqlAsset]]]
+QueryInfo = list[dict[str, Union[GQLQueryFields, TableauGraphqlAsset]]]
 
-QUERY_FIELDS: Dict[TableauAsset, QueryInfo] = {
+QUERY_FIELDS: dict[TableauAsset, QueryInfo] = {
     TableauAsset.CUSTOM_SQL_TABLE: [
         {
             FIELDS: GQLQueryFields.CUSTOM_SQL_TABLE,

castor_extractor/visualization/tableau/tsc_fields.py

@@ -1,10 +1,9 @@
 # TSC for TableauServerClient: basic REST API to extracting core objects
-from typing import Dict, Set
 
 from .assets import TableauAsset
 
 # TSC fields extracted per assets
-TSC_FIELDS: Dict[TableauAsset, Set[str]] = {
+TSC_FIELDS: dict[TableauAsset, set[str]] = {
     TableauAsset.PROJECT: {
         "id",
         "name",

castor_extractor/visualization/tableau/types.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Tuple, Union
+from typing import Union
 
 from tableauserverclient import ServerResponseError  # type: ignore
 from typing_extensions import Literal
@@ -6,6 +6,6 @@ from typing_extensions import Literal
 from .errors import TableauErrorCode
 
 PageReturn = Union[
-    Tuple[List[Dict], Literal[None]],
-    Tuple[Literal[None], Union[TableauErrorCode, ServerResponseError]],
+    tuple[list[dict], Literal[None]],
+    tuple[Literal[None], Union[TableauErrorCode, ServerResponseError]],
 ]

castor_extractor/visualization/tableau_revamp/client/client.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Optional
 
 import tableauserverclient as TSC  # type: ignore
 
@@ -121,12 +122,16 @@ class TableauRevampClient:
         credentials: TableauRevampCredentials,
         timeout_sec: int = DEFAULT_TIMEOUT_SECONDS,
         with_pulse: bool = False,
+        override_page_size: Optional[int] = None,
     ):
         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_metadata = TableauClientMetadataApi(
+            server=self._server,
+            override_page_size=override_page_size,
+        )
         self._client_rest = TableauClientRestApi(server=self._server)
         self._client_tsc = TableauClientTSC(server=self._server)
 

castor_extractor/visualization/tableau_revamp/client/client_metadata_api.py

@@ -1,16 +1,17 @@
-from typing import Dict, Iterator, Optional
+from collections.abc import Iterator
+from typing import Optional
 
 import tableauserverclient as TSC  # type: ignore
 
-from ....utils import SerializedAsset
+from ....utils import SerializedAsset, retry
 from ..assets import TableauRevampAsset
 from ..constants import DEFAULT_PAGE_SIZE
-from .errors import TableauApiError
+from .errors import TableauApiError, TableauApiTimeout
 from .gql_queries import FIELDS_QUERIES, GQL_QUERIES, QUERY_TEMPLATE
 
 # increase the value when extraction is too slow
 # decrease the value when timeouts arise
-_CUSTOM_PAGE_SIZE: Dict[TableauRevampAsset, int] = {
+_CUSTOM_PAGE_SIZE: dict[TableauRevampAsset, int] = {
     # for some clients, extraction of columns tend to hit the node limit
     # https://community.tableau.com/s/question/0D54T00000YuK60SAF/metadata-query-nodelimitexceeded-error
     # the workaround is to reduce pagination
@@ -20,21 +21,58 @@ _CUSTOM_PAGE_SIZE: Dict[TableauRevampAsset, int] = {
     TableauRevampAsset.TABLE: 50,
 }
 
+_TIMEOUT_MESSAGE = (
+    "Execution canceled because timeout of 30000 millis was reached"
+)
+
+_RETRY_BASE_MS = 10_000
+_RETRY_COUNT = 4
+
+
+def _check_errors(answer: dict) -> None:
+    """
+    handle errors in graphql response:
+    - return None when there's no errors in the answer
+    - TableauApiTimeout if any of the errors is a timeout
+    - TableauApiError (generic) otherwise
+    """
+    if "errors" not in answer:
+        return
+
+    errors = answer["errors"]
+
+    for error in errors:
+        if error.get("message") == _TIMEOUT_MESSAGE:
+            # we need specific handling for timeout issues (retry strategy)
+            raise TableauApiTimeout(errors)
+
+    raise TableauApiError(answer["errors"])
+
 
 def gql_query_scroll(
     server,
     query: str,
     resource: str,
 ) -> Iterator[SerializedAsset]:
-    """Iterate over GQL query results, handling pagination and cursor"""
+    """
+    Iterate over GQL query results, handling pagination and cursor
 
+    We have a retry strategy when timeout issues arise.
+    It's a known issue on Tableau side, still waiting for their fix:
+    https://issues.salesforce.com/issue/a028c00000zKahoAAC/undefined
+    """
+
+    @retry(
+        exceptions=(TableauApiTimeout,),
+        max_retries=_RETRY_COUNT,
+        base_ms=_RETRY_BASE_MS,
+    )
     def _call(cursor: Optional[str]) -> dict:
         # If cursor is defined it must be quoted else use null token
         token = "null" if cursor is None else f'"{cursor}"'
         query_ = query.replace("AFTER_TOKEN_SIGNAL", token)
         answer = server.metadata.query(query_)
-        if "errors" in answer:
-            raise TableauApiError(answer["errors"])
+        _check_errors(answer)
         return answer["data"][f"{resource}Connection"]
 
     cursor = None
@@ -58,8 +96,10 @@ class TableauClientMetadataApi:
     def __init__(
         self,
         server: TSC.Server,
+        override_page_size: Optional[int] = None,
     ):
         self._server = server
+        self._forced_page_size = override_page_size
 
     def _call(
         self,
@@ -75,9 +115,16 @@ class TableauClientMetadataApi:
         result_pages = gql_query_scroll(self._server, query, resource)
         return [asset for page in result_pages for asset in page]
 
+    def _page_size(self, asset: TableauRevampAsset) -> int:
+        return (
+            self._forced_page_size
+            or _CUSTOM_PAGE_SIZE.get(asset)
+            or DEFAULT_PAGE_SIZE
+        )
+
     def _fetch_fields(self) -> SerializedAsset:
         result: SerializedAsset = []
-        page_size = _CUSTOM_PAGE_SIZE[TableauRevampAsset.FIELD]
+        page_size = self._page_size(TableauRevampAsset.FIELD)
         for resource, fields in FIELDS_QUERIES:
             current = self._call(resource, fields, page_size)
             result.extend(current)
@@ -90,6 +137,6 @@ class TableauClientMetadataApi:
         if asset == TableauRevampAsset.FIELD:
             return self._fetch_fields()
 
-        page_size = _CUSTOM_PAGE_SIZE.get(asset) or DEFAULT_PAGE_SIZE
+        page_size = self._page_size(asset)
         resource, fields = GQL_QUERIES[asset]
         return self._call(resource, fields, page_size)

castor_extractor/visualization/tableau_revamp/client/client_rest_api.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict, Optional
+from typing import Optional
 
 import requests
 import tableauserverclient as TSC  # type: ignore
@@ -40,7 +40,7 @@ class TableauClientRestApi:
         return self._server.http_options["timeout"]
 
     @property
-    def headers(self) -> Dict[str, str]:
+    def headers(self) -> dict[str, str]:
         return {"x-tableau-auth": self._server.auth_token}
 
     def _get_site_name(self) -> str:
@@ -52,7 +52,7 @@ class TableauClientRestApi:
         self,
         url: str,
         page_token: Optional[str] = None,
-    ) -> Dict:
+    ) -> dict:
         if page_token:
             url += f"?page_token={page_token}"
 

castor_extractor/visualization/tableau_revamp/client/client_tsc.py

@@ -1,4 +1,5 @@
-from typing import Any, Dict, Iterable, Iterator
+from collections.abc import Iterable, Iterator
+from typing import Any
 
 import tableauserverclient as TSC  # type: ignore
 
@@ -30,7 +31,7 @@ class TableauClientTSC:
         self,
         data: Iterable,
         asset: TableauRevampAsset,
-    ) -> Iterator[Dict]:
+    ) -> Iterator[dict]:
         keys = REST_FIELDS[asset]
 
         for row in data:

castor_extractor/visualization/tableau_revamp/client/errors.py

@@ -1,3 +1,8 @@
 class TableauApiError(ValueError):
     def __init__(self, error: str):
         super().__init__(f"Tableau API returned the following error: {error}")
+
+
+class TableauApiTimeout(ValueError):
+    def __init__(self, error: str):
+        super().__init__(f"Tableau API returned a timeout error: {error}")

castor_extractor/visualization/tableau_revamp/client/gql_queries.py

@@ -1,5 +1,3 @@
-from typing import Dict, Tuple
-
 from ..assets import TableauRevampAsset
 
 QUERY_TEMPLATE = """
@@ -130,7 +128,7 @@ workbook { id }
 """
 
 
-GQL_QUERIES: Dict[TableauRevampAsset, Tuple[str, str]] = {
+GQL_QUERIES: dict[TableauRevampAsset, tuple[str, str]] = {
     TableauRevampAsset.COLUMN: ("columns", _COLUMNS_QUERY),
     TableauRevampAsset.DASHBOARD: ("dashboards", _DASHBOARDS_QUERY),
     TableauRevampAsset.DATASOURCE: ("datasources", _DATASOURCES_QUERY),

castor_extractor/visualization/tableau_revamp/client/rest_fields.py

@@ -1,9 +1,7 @@
-from typing import Dict, Set
-
 from ..assets import TableauRevampAsset
 
 # list of fields to pick in REST API or TSC responses
-REST_FIELDS: Dict[TableauRevampAsset, Set[str]] = {
+REST_FIELDS: dict[TableauRevampAsset, set[str]] = {
     TableauRevampAsset.DATASOURCE: {
         "id",
         "project_id",

castor_extractor/visualization/tableau_revamp/extract.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Iterable, Tuple
+from collections.abc import Iterable
 
 from ...utils import (
     OUTPUT_DIR,
@@ -18,7 +18,7 @@ logger = logging.getLogger(__name__)
 
 def iterate_all_data(
     client: TableauRevampClient,
-) -> Iterable[Tuple[TableauRevampAsset, list]]:
+) -> Iterable[tuple[TableauRevampAsset, list]]:
     """Iterate over the extracted Data from Tableau"""
 
     for asset in TableauRevampAsset:

castor_extractor/visualization/thoughtspot/client/client.py

@@ -1,4 +1,5 @@
-from typing import Dict, Iterator, Optional
+from collections.abc import Iterator
+from typing import Optional
 
 import requests
 
@@ -35,7 +36,7 @@ THOUGHTSPOT_SAFE_MODE = RequestSafeMode()
 
 
 class ThoughtspotBearerAuth(BearerAuth):
-    def __init__(self, host: str, token_payload: Dict[str, str]):
+    def __init__(self, host: str, token_payload: dict[str, str]):
         auth_endpoint = ThoughtspotEndpointFactory.authentication()
         self.authentication_url = build_url(host, auth_endpoint)
         self.token_payload = token_payload

castor_extractor/visualization/thoughtspot/client/utils.py

@@ -1,6 +1,6 @@
 import csv
+from collections.abc import Iterator
 from io import StringIO
-from typing import Iterator
 
 
 def usage_liveboard_reader(usage_liveboard_csv: str) -> Iterator[dict]:

castor_extractor/visualization/thoughtspot/extract.py

@@ -1,5 +1,6 @@
 import logging
-from typing import Iterable, Iterator, Tuple, Union
+from collections.abc import Iterable, Iterator
+from typing import Union
 
 from ...utils import (
     OUTPUT_DIR,
@@ -21,7 +22,7 @@ logger = logging.getLogger(__name__)
 
 def iterate_all_data(
     client: ThoughtspotClient,
-) -> Iterable[Tuple[ThoughtspotAsset, Union[list, Iterator, dict]]]:
+) -> Iterable[tuple[ThoughtspotAsset, Union[list, Iterator, dict]]]:
     """Iterate over the extracted data from Thoughtspot"""
 
     for asset in ThoughtspotAsset:

castor_extractor/warehouse/abstract/asset.py

@@ -1,5 +1,4 @@
 from enum import Enum
-from typing import Dict, List, Set, Tuple
 
 from ...types import ExternalAsset, classproperty
 
@@ -26,7 +25,7 @@ class WarehouseAsset(ExternalAsset):
     VIEW_DDL = "view_ddl"
 
     @classproperty
-    def optional(cls) -> Set["WarehouseAsset"]:
+    def optional(cls) -> set["WarehouseAsset"]:
         return {
             WarehouseAsset.ADDITIONAL_COLUMN_LINEAGE,
             WarehouseAsset.ADDITIONAL_TABLE_LINEAGE,
@@ -50,7 +49,7 @@ class WarehouseAssetGroup(Enum):
 
 
 # tuple of supported assets for each group (depends on the technology)
-SupportedAssets = Dict[WarehouseAssetGroup, Tuple[WarehouseAsset, ...]]
+SupportedAssets = dict[WarehouseAssetGroup, tuple[WarehouseAsset, ...]]
 
 # shared by all technologies
 CATALOG_ASSETS = (
@@ -80,13 +79,13 @@ NON_EXTRACTABLE_ASSETS = {WarehouseAssetGroup.EXTERNAL_LINEAGE}
 
 def extractable_asset_groups(
     supported_assets: SupportedAssets,
-) -> List[Tuple[WarehouseAsset, ...]]:
+) -> list[tuple[WarehouseAsset, ...]]:
     """
     helper function to differentiate
     extractable assets vs supported (ingest-able) assets
     """
     groups = set(supported_assets).difference(NON_EXTRACTABLE_ASSETS)
-    extractable: Set[Tuple[WarehouseAsset, ...]] = {
+    extractable: set[tuple[WarehouseAsset, ...]] = {
         supported_assets[group] for group in groups
     }
     return list(extractable)

castor_extractor/warehouse/abstract/extract.py

@@ -1,6 +1,7 @@
 import logging
+from collections.abc import Iterator
 from itertools import chain
-from typing import Callable, Iterator, List, Optional, Tuple
+from typing import Callable, Optional
 
 from ...utils import (
     OUTPUT_DIR,
@@ -16,7 +17,7 @@ from .query import AbstractQueryBuilder, ExtractionQuery
 logger = logging.getLogger(__name__)
 
 
-def common_args(kwargs: dict) -> Tuple[str, bool]:
+def common_args(kwargs: dict) -> tuple[str, bool]:
     """Args used by all technologies"""
     output_directory = kwargs.get("output_directory") or from_env(OUTPUT_DIR)
     skip_existing = kwargs.get("skip_existing") or False
@@ -39,7 +40,7 @@ class SQLExtractionProcessor:
         self._safe_mode = safe_mode
 
     @staticmethod
-    def _unique(data: Iterator[dict]) -> List[dict]:
+    def _unique(data: Iterator[dict]) -> list[dict]:
         """
         Remove duplicate in the given data.
         Remark: this method implies loading all data in memory: it breaks the streaming pipeline !

castor_extractor/warehouse/abstract/query.py

@@ -1,7 +1,7 @@
 import inspect
 import os
 from abc import ABC, abstractmethod
-from typing import List, Optional, Tuple
+from typing import Optional
 
 from .asset import WarehouseAsset
 from .time_filter import TimeFilter
@@ -37,7 +37,7 @@ class AbstractQueryBuilder(ABC):
     def __init__(
         self,
         time_filter: Optional[TimeFilter],
-        duplicated: Optional[Tuple[WarehouseAsset, ...]] = None,
+        duplicated: Optional[tuple[WarehouseAsset, ...]] = None,
     ):
         self._time_filter = time_filter or TimeFilter.default()
         self._duplicated = duplicated
@@ -55,7 +55,7 @@ class AbstractQueryBuilder(ABC):
         """read from a file located in queries directory"""
         root = os.path.dirname(inspect.getfile(self.__class__))
         path = os.path.join(root, QUERIES_DIR, filename)
-        with open(path, "r") as f:
+        with open(path) as f:
             return f.read()
 
     def load_statement(self, asset: WarehouseAsset) -> str:
@@ -75,7 +75,7 @@ class AbstractQueryBuilder(ABC):
         return ExtractionQuery(statement, params)
 
     @abstractmethod
-    def build(self, asset: WarehouseAsset) -> List[ExtractionQuery]:
+    def build(self, asset: WarehouseAsset) -> list[ExtractionQuery]:
         """
         Build the Query allowing extraction of the given asset
         - Most of the time, returns a single query

castor_extractor/warehouse/bigquery/client.py

@@ -1,6 +1,6 @@
 import itertools
 import logging
-from typing import List, Optional, Set
+from typing import Optional
 
 from google.api_core.exceptions import Forbidden  # type: ignore
 from google.cloud.bigquery import Client as GoogleCloudClient  # type: ignore
@@ -27,9 +27,9 @@ class BigQueryClient(SqlalchemyClient):
     def __init__(
         self,
         credentials: dict,
-        db_allowed: Optional[Set[str]] = None,
-        db_blocked: Optional[Set[str]] = None,
-        dataset_blocked: Optional[Set[str]] = None,
+        db_allowed: Optional[set[str]] = None,
+        db_blocked: Optional[set[str]] = None,
+        dataset_blocked: Optional[set[str]] = None,
     ):
         super().__init__(credentials)
         self._db_allowed = db_allowed
@@ -37,8 +37,8 @@ class BigQueryClient(SqlalchemyClient):
         self._dataset_blocked = dataset_blocked
         self.credentials = self._credentials()
         self.client = self._client()
-        self._projects: List[str] | None = None
-        self._datasets: List[Dataset] | None = None
+        self._projects: list[str] | None = None
+        self._datasets: list[Dataset] | None = None
 
     @staticmethod
     def name() -> str:
@@ -78,7 +78,7 @@ class BigQueryClient(SqlalchemyClient):
             credentials=self.credentials,
         )
 
-    def _list_datasets(self) -> List[Dataset]:
+    def _list_datasets(self) -> list[Dataset]:
         """
         Returns datasets available for the given GCP client
         Cache the result in self._datasets to reduce number of API calls
@@ -98,7 +98,7 @@ class BigQueryClient(SqlalchemyClient):
         base_ms=_RETRY_BASE_MS,
         log_exc_info=True,
     )
-    def get_projects(self) -> List[str]:
+    def get_projects(self) -> list[str]:
         """
         Returns distinct project_id available for the given GCP client
         Cache the result in self._projects to reduce number of API calls.

castor_extractor/warehouse/bigquery/extract.py

@@ -38,7 +38,7 @@ def _credentials(params: dict) -> dict:
     """extract GCP credentials"""
     path = params.get("credentials") or from_env(BIGQUERY_CREDENTIALS)
     logger.info(f"Credentials fetched from {path}")
-    with open(path, "r") as file:
+    with open(path) as file:
         return cast(dict, json.load(file))
 
 

castor_extractor/warehouse/bigquery/query.py

@@ -1,5 +1,5 @@
 import logging
-from typing import List, Optional
+from typing import Optional
 
 from ..abstract import (
     AbstractQueryBuilder,
@@ -109,7 +109,7 @@ class BigQueryQueryBuilder(AbstractQueryBuilder):
             else self._regions
         )
 
-    def build(self, asset: WarehouseAsset) -> List[ExtractionQuery]:
+    def build(self, asset: WarehouseAsset) -> list[ExtractionQuery]:
         """
         It would be easier to stitch data directly in the query statement (UNION ALL).
         Unfortunately, querying INFORMATION_SCHEMA on multiple regions

castor_extractor/warehouse/bigquery/types.py

@@ -1,4 +1,2 @@
-from typing import Set, Tuple
-
-SetString = Set[str]
-SetTwoString = Set[Tuple[str, str]]
+SetString = set[str]
+SetTwoString = set[tuple[str, str]]