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

castor_extractor/visualization/sigma/client/credentials.py

@@ -1,4 +1,6 @@
-from pydantic import Field, SecretStr
+from typing import Dict
+
+from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 CASTOR_ENV_PREFIX = "CASTOR_SIGMA_"
@@ -17,3 +19,11 @@ class SigmaCredentials(BaseSettings):
     client_id: str
     host: str
     grant_type: str = "client_credentials"
+
+    @property
+    def token_payload(self) -> Dict[str, str]:
+        return {
+            "grant_type": self.grant_type,
+            "client_id": self.client_id,
+            "client_secret": self.api_token,
+        }

castor_extractor/visualization/sigma/client/endpoints.py

@@ -1,4 +1,4 @@
-class EndpointFactory:
+class SigmaEndpointFactory:
     """Wrapper class around all endpoints we're using"""
 
     DATASETS = "datasets"

castor_extractor/visualization/sigma/client/pagination.py

@@ -1,24 +1,28 @@
-import logging
 from typing import Optional
 
-logger = logging.getLogger(__name__)
+from pydantic import ConfigDict, Field
+from pydantic.alias_generators import to_camel
 
+from ....utils import PaginationModel
 
-class Pagination:
-    """This is a wrapper around Sigma's pagination system"""
+SIGMA_API_LIMIT = 500  # default number of records per page
 
-    def __init__(
-        self,
-        next_page: Optional[str],
-        entries: Optional[list] = None,
-        total: Optional[int] = 0,
-    ):
-        self.next_page = next_page
-        self.entries = entries or []
-        self.total = total
 
-    def generate_url(self, endpoint_url: str) -> str:
-        """Generates the paginated url based on the targeted endpoint"""
-        pagination = f"?page={self.next_page}"
-        paginated_url = f"{endpoint_url}{pagination}"
-        return paginated_url
+class SigmaPagination(PaginationModel):
+    next_page: Optional[str] = "0"
+    entries: list = Field(default_factory=list)
+
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )
+
+    def is_last(self) -> bool:
+        return self.next_page is None
+
+    def next_page_payload(self) -> dict:
+        return {"page": self.next_page}
+
+    def page_results(self) -> list:
+        return self.entries

castor_extractor/visualization/tableau/tests/unit/rest_api/auth_test.py

@@ -3,7 +3,6 @@ import unittest
 from unittest import mock
 
 import requests_mock  # type: ignore
-import tableauserverclient as TSC  # type: ignore
 
 from ....client import ApiClient
 from ..utils import KEYS

castor_extractor/visualization/tableau/tests/unit/rest_api/credentials_test.py

@@ -1,6 +1,3 @@
-import requests_mock  # type: ignore
-import tableauserverclient as TSC  # type: ignore
-
 from ....client.credentials import CredentialsApi
 
 

castor_extractor/visualization/tableau_revamp/assets.py

@@ -10,9 +10,20 @@ class TableauRevampAsset(ExternalAsset):
     DASHBOARD = "dashboards"
     DATASOURCE = "datasources"
     FIELD = "fields"
+    METRIC = "metrics"
+    METRIC_DEFINITION = "metrics_definitions"
     PROJECT = "projects"
     SHEET = "sheets"
+    SUBSCRIPTION = "subscriptions"
     TABLE = "tables"
     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

@@ -1,59 +1,39 @@
 import logging
-from typing import Dict, Iterator, List, Optional
 
 import tableauserverclient as TSC  # type: ignore
-from tableauserverclient import Pager
 
 from ....utils import SerializedAsset
-from ..assets import TableauRevampAsset
-from ..constants import (
-    DEFAULT_PAGE_SIZE,
-    DEFAULT_TIMEOUT_SECONDS,
-    TABLEAU_SERVER_VERSION,
-)
+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
+from .client_tsc import TableauClientTSC
 from .credentials import TableauRevampCredentials
-from .errors import TableauApiError
-from .gql_queries import FIELDS_QUERIES, GQL_QUERIES, QUERY_TEMPLATE
-from .tsc_fields import TSC_FIELDS
 
 logger = logging.getLogger(__name__)
 
-# these assets must be extracted via TableauServerClient
+# these assets must be extracted via TableauServerClient (TSC)
 _TSC_ASSETS = (
-    # only users who published content can be extracted from MetadataAPI
-    TableauRevampAsset.USER,
     # projects are not available in Metadata API
     TableauRevampAsset.PROJECT,
     # view count are not available in Metadata API
     TableauRevampAsset.USAGE,
+    # only users who published content can be extracted from MetadataAPI
+    TableauRevampAsset.USER,
 )
 
-# increase the value when extraction is too slow
-# decrease the value when timeouts arise
-_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
-    TableauRevampAsset.COLUMN: 50,
-    # fields are light but volumes are bigger
-    TableauRevampAsset.FIELD: 1000,
-    TableauRevampAsset.TABLE: 50,
-}
-
-
-def _pick_fields(
-    data: Pager,
-    asset: TableauRevampAsset,
-) -> SerializedAsset:
-    fields = TSC_FIELDS[asset]
-
-    def _pick(row: dict):
-        return {field: getattr(row, field) for field in fields}
+# these assets must be extracted via the REST API
+_REST_API_ASSETS = (
+    # Tableau Pulse assets are only available in REST API
+    TableauRevampAsset.METRIC,
+    TableauRevampAsset.METRIC_DEFINITION,
+    TableauRevampAsset.SUBSCRIPTION,
+)
 
-    return [_pick(row) for row in data]
+logging.getLogger("tableau.endpoint").setLevel(logging.WARNING)
 
 
-def _enrich_datasources_with_tsc(
+def _merge_datasources(
     datasources: SerializedAsset,
     tsc_datasources: SerializedAsset,
 ) -> SerializedAsset:
@@ -71,14 +51,19 @@ def _enrich_datasources_with_tsc(
             # embedded datasources cannot be accessed via URL => no webpage_url
             continue
         luid = datasource["luid"]
-        tsc_datasource = mapping[luid]
+        tsc_datasource = mapping.get(luid)
+        if not tsc_datasource:
+            # it happens that a datasource is in Metadata API but not in TSC
+            datasource["projectLuid"] = None
+            datasource["webpageUrl"] = None
+            continue
         datasource["projectLuid"] = tsc_datasource["project_id"]
         datasource["webpageUrl"] = tsc_datasource["webpage_url"]
 
     return datasources
 
 
-def _enrich_workbooks_with_tsc(
+def _merge_workbooks(
     workbooks: SerializedAsset,
     tsc_workbooks: SerializedAsset,
 ) -> SerializedAsset:
@@ -104,61 +89,44 @@ def _enrich_workbooks_with_tsc(
     return workbooks
 
 
-def gql_query_scroll(
-    server,
-    query: str,
-    resource: str,
-) -> Iterator[SerializedAsset]:
-    """Iterate over GQL query results, handling pagination and cursor"""
-
-    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"])
-        return answer["data"][f"{resource}Connection"]
-
-    cursor = None
-    while True:
-        payload = _call(cursor)
-        yield payload["nodes"]
-
-        page_info = payload["pageInfo"]
-        if page_info["hasNextPage"]:
-            cursor = page_info["endCursor"]
-        else:
-            break
+def _server(
+    server_url: str,
+    timeout_sec: int,
+) -> TSC.Server:
+    options = {"verify": True, "timeout": timeout_sec}
+    server = TSC.Server(server_url, use_server_version=True)
+    server.add_http_options(options)
+    return server
 
 
 class TableauRevampClient:
     """
     Connect to Tableau's API and extract assets.
 
-    Relies on TableauServerClient overlay:
+    Relies on TableauServerClient (TSC) overlay for authentication
     https://tableau.github.io/server-client-python/docs/
-    - for connection
-    - to extract Users (Metadata
-
-    Calls the MetadataAPI, using graphQL
-    https://help.tableau.com/current/api/metadata_api/en-us/reference/index.html
     """
 
     def __init__(
         self,
         credentials: TableauRevampCredentials,
         timeout_sec: int = DEFAULT_TIMEOUT_SECONDS,
+        with_pulse: bool = False,
     ):
         self._credentials = credentials
-        self._server = TSC.Server(self._credentials.server_url)
-        options = {"verify": True, "timeout": timeout_sec}
-        self._server.add_http_options(options)
-        self._server.version = TABLEAU_SERVER_VERSION
-        self.errors: List[str] = []
-
-    @staticmethod
-    def name() -> str:
+        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)
+        self._client_tsc = TableauClientTSC(server=self._server)
+
+    @property
+    def base_url(self) -> str:
+        return self._credentials.server_url
+
+    @property
+    def name(self) -> str:
         return "Tableau/API"
 
     def _user_password_login(self) -> None:
@@ -203,98 +171,50 @@ class TableauRevampClient:
             "Invalid credentials: either user/password or PAT must be provided",
         )
 
-    def base_url(self) -> str:
-        return self._credentials.server_url
-
-    def _fetch_from_tsc(
-        self,
-        asset: TableauRevampAsset,
-    ) -> SerializedAsset:
-        if asset == TableauRevampAsset.DATASOURCE:
-            data = TSC.Pager(self._server.datasources)
-
-        elif asset == TableauRevampAsset.PROJECT:
-            data = TSC.Pager(self._server.projects)
-
-        elif asset == TableauRevampAsset.USAGE:
-            data = TSC.Pager(self._server.views, usage=True)
-
-        elif asset == TableauRevampAsset.USER:
-            data = TSC.Pager(self._server.users)
-
-        elif asset == TableauRevampAsset.WORKBOOK:
-            data = TSC.Pager(self._server.workbooks)
-
-        else:
-            raise AssertionError(f"Fetching from TSC not supported for {asset}")
-
-        return _pick_fields(data, asset)
-
-    def _run_graphql_query(
-        self,
-        resource: str,
-        fields: str,
-        page_size: int = DEFAULT_PAGE_SIZE,
-    ) -> SerializedAsset:
-        query = QUERY_TEMPLATE.format(
-            resource=resource,
-            fields=fields,
-            page_size=page_size,
-        )
-        result_pages = gql_query_scroll(self._server, query, resource)
-        return [asset for page in result_pages for asset in page]
-
-    def _fetch_fields(self) -> SerializedAsset:
-        result: SerializedAsset = []
-        page_size = _CUSTOM_PAGE_SIZE[TableauRevampAsset.FIELD]
-        for resource, fields in FIELDS_QUERIES:
-            current = self._run_graphql_query(resource, fields, page_size)
-            result.extend(current)
-        return result
-
-    def _fetch_from_metadata_api(
-        self,
-        asset: TableauRevampAsset,
-    ) -> SerializedAsset:
-        if asset == TableauRevampAsset.FIELD:
-            return self._fetch_fields()
-
-        page_size = _CUSTOM_PAGE_SIZE.get(asset) or DEFAULT_PAGE_SIZE
-        resource, fields = GQL_QUERIES[asset]
-        return self._run_graphql_query(resource, fields, page_size)
-
     def _fetch_datasources(self) -> SerializedAsset:
         asset = TableauRevampAsset.DATASOURCE
 
-        datasources = self._fetch_from_metadata_api(asset)
-        datasource_projects = self._fetch_from_tsc(asset)
+        datasources = self._client_metadata.fetch(asset)
+        tsc_datasources = self._client_tsc.fetch(asset)
 
-        return _enrich_datasources_with_tsc(datasources, datasource_projects)
+        return _merge_datasources(datasources, tsc_datasources)
 
     def _fetch_workbooks(self) -> SerializedAsset:
         asset = TableauRevampAsset.WORKBOOK
 
-        workbooks = self._fetch_from_metadata_api(asset)
-        workbook_projects = self._fetch_from_tsc(asset)
+        workbooks = self._client_metadata.fetch(asset)
+        workbook_projects = self._client_tsc.fetch(asset)
 
-        return _enrich_workbooks_with_tsc(workbooks, workbook_projects)
+        return _merge_workbooks(workbooks, workbook_projects)
 
-    def fetch(self, asset: TableauRevampAsset) -> SerializedAsset:
+    def fetch(
+        self,
+        asset: TableauRevampAsset,
+    ) -> SerializedAsset:
         """
         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}...")
 
         if asset == TableauRevampAsset.DATASOURCE:
-            # both APIs are required to extract datasources
+            # two APIs are required to extract datasources
             return self._fetch_datasources()
 
         if asset == TableauRevampAsset.WORKBOOK:
-            # both APIs are required to extract workbooks
+            # two APIs are required to extract workbooks
             return self._fetch_workbooks()
 
         if asset in _TSC_ASSETS:
             # some assets can only be extracted via TSC
-            return self._fetch_from_tsc(asset)
+            return self._client_tsc.fetch(asset)
+
+        if asset in _REST_API_ASSETS:
+            # some assets can only be extracted via REST API
+            return self._client_rest.fetch(asset)
 
-        # extract most assets via Metadata API
-        return self._fetch_from_metadata_api(asset)
+        # other assets can be extracted via Metadata API
+        return self._client_metadata.fetch(asset)

castor_extractor/visualization/tableau_revamp/client/client_metadata_api.py

@@ -0,0 +1,95 @@
+from typing import Dict, Iterator, Optional
+
+import tableauserverclient as TSC  # type: ignore
+
+from ....utils import SerializedAsset
+from ..assets import TableauRevampAsset
+from ..constants import DEFAULT_PAGE_SIZE
+from .errors import TableauApiError
+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] = {
+    # 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
+    TableauRevampAsset.COLUMN: 50,
+    # fields are light but volumes are bigger
+    TableauRevampAsset.FIELD: 1000,
+    TableauRevampAsset.TABLE: 50,
+}
+
+
+def gql_query_scroll(
+    server,
+    query: str,
+    resource: str,
+) -> Iterator[SerializedAsset]:
+    """Iterate over GQL query results, handling pagination and cursor"""
+
+    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"])
+        return answer["data"][f"{resource}Connection"]
+
+    cursor = None
+    while True:
+        payload = _call(cursor)
+        yield payload["nodes"]
+
+        page_info = payload["pageInfo"]
+        if page_info["hasNextPage"]:
+            cursor = page_info["endCursor"]
+        else:
+            break
+
+
+class TableauClientMetadataApi:
+    """
+    Calls the MetadataAPI, using graphQL
+    https://help.tableau.com/current/api/metadata_api/en-us/reference/index.html
+    """
+
+    def __init__(
+        self,
+        server: TSC.Server,
+    ):
+        self._server = server
+
+    def _call(
+        self,
+        resource: str,
+        fields: str,
+        page_size: int = DEFAULT_PAGE_SIZE,
+    ) -> SerializedAsset:
+        query = QUERY_TEMPLATE.format(
+            resource=resource,
+            fields=fields,
+            page_size=page_size,
+        )
+        result_pages = gql_query_scroll(self._server, query, resource)
+        return [asset for page in result_pages for asset in page]
+
+    def _fetch_fields(self) -> SerializedAsset:
+        result: SerializedAsset = []
+        page_size = _CUSTOM_PAGE_SIZE[TableauRevampAsset.FIELD]
+        for resource, fields in FIELDS_QUERIES:
+            current = self._call(resource, fields, page_size)
+            result.extend(current)
+        return result
+
+    def fetch(
+        self,
+        asset: TableauRevampAsset,
+    ) -> SerializedAsset:
+        if asset == TableauRevampAsset.FIELD:
+            return self._fetch_fields()
+
+        page_size = _CUSTOM_PAGE_SIZE.get(asset) or DEFAULT_PAGE_SIZE
+        resource, fields = GQL_QUERIES[asset]
+        return self._call(resource, fields, page_size)

castor_extractor/visualization/tableau_revamp/client/client_rest_api.py

@@ -0,0 +1,128 @@
+import logging
+from typing import Dict, Optional
+
+import requests
+import tableauserverclient as TSC  # type: ignore
+
+from ....utils import SerializedAsset, deduplicate
+from ..assets import TableauRevampAsset
+from .rest_fields import REST_FIELDS
+
+logger = logging.getLogger(__name__)
+
+_PULSE_API = "api/-/pulse"
+
+_METRICS_URL = "{base}/pulse/site/{site}/metrics/{definition_id}"
+
+
+def _pick(
+    data: SerializedAsset,
+    asset: TableauRevampAsset,
+) -> SerializedAsset:
+    keys = REST_FIELDS[asset]
+    return [{key: row[key] for key in keys} for row in data]
+
+
+class TableauClientRestApi:
+    """
+    Extract Tableau Assets using REST API
+    https://help.tableau.com/current/api/rest_api/en-us/REST/rest_api_ref.htm
+    """
+
+    def __init__(
+        self,
+        server: TSC.Server,
+    ):
+        self._server = server
+
+    @property
+    def timeout_sec(self) -> int:
+        return self._server.http_options["timeout"]
+
+    @property
+    def headers(self) -> Dict[str, str]:
+        return {"x-tableau-auth": self._server.auth_token}
+
+    def _get_site_name(self) -> str:
+        site_id = self._server.site_id
+        site = self._server.sites.get_by_id(site_id)
+        return site.content_url
+
+    def _get(
+        self,
+        url: str,
+        page_token: Optional[str] = None,
+    ) -> Dict:
+        if page_token:
+            url += f"?page_token={page_token}"
+
+        logger.debug(f"Calling REST API: {url}")
+        response = requests.get(
+            url,
+            headers=self.headers,
+            timeout=self.timeout_sec,
+        )
+        response.raise_for_status()
+        return response.json()
+
+    def _call(
+        self,
+        path: str,
+        target: str,
+    ) -> SerializedAsset:
+        base = self._server.server_address.strip("/")
+        url = f"{base}/{path}/{target}"
+
+        next_page_token = None
+        data = []
+
+        while True:
+            response = self._get(url, next_page_token)
+            data += response[target]
+            next_page_token = response.get("next_page_token")
+            if not next_page_token:
+                break
+
+        return data
+
+    def _compute_metric_url(self, data: SerializedAsset) -> None:
+        site = self._get_site_name()
+        base_url = self._server.server_address.strip("/")
+        for row in data:
+            row["metadata"]["url"] = _METRICS_URL.format(
+                base=base_url,
+                site=site,
+                definition_id=row["metadata"]["id"],
+            )
+
+    def _fetch_metrics(self, definitions: SerializedAsset) -> SerializedAsset:
+        metrics = []
+        for definition in definitions:
+            definition_id = definition["metadata"]["id"]
+            path = f"{_PULSE_API}/definitions/{definition_id}"
+            metrics += self._call(path=path, target="metrics")
+
+        # for some reason, the REST API sometimes send the same metric twice
+        return deduplicate("id", metrics)
+
+    def fetch(
+        self,
+        asset: TableauRevampAsset,
+    ) -> SerializedAsset:
+        if asset == TableauRevampAsset.SUBSCRIPTION:
+            # https://help.tableau.com/current/api/rest_api/en-us/REST/rest_api_ref_pulse.htm#PulseSubscriptionService_ListSubscriptions
+            data = self._call(path=_PULSE_API, target="subscriptions")
+
+        elif asset == TableauRevampAsset.METRIC_DEFINITION:
+            # https://help.tableau.com/current/api/rest_api/en-us/REST/rest_api_ref_pulse.htm#MetricQueryService_ListDefinitions
+            data = self._call(path=_PULSE_API, target="definitions")
+            self._compute_metric_url(data)
+
+        elif asset == TableauRevampAsset.METRIC:
+            # https://help.tableau.com/current/api/rest_api/en-us/REST/rest_api_ref_pulse.htm#MetricQueryService_ListMetrics
+            definitions = self._call(path=_PULSE_API, target="definitions")
+            data = self._fetch_metrics(definitions)
+        else:
+            raise AssertionError(f"Unsupported asset {asset} for REST API")
+
+        return _pick(data, asset)