castor-extractor 0.16.5__py3-none-any.whl → 0.16.9__py3-none-any.whl

castor_extractor/visualization/tableau_revamp/client/client.py

@@ -0,0 +1,297 @@
+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 .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
+_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,
+)
+
+# speed up extraction: fields and columns are smaller but volumes are bigger
+_CUSTOM_PAGE_SIZE: Dict[TableauRevampAsset, int] = {
+    TableauRevampAsset.FIELD: 1000,
+    TableauRevampAsset.COLUMN: 1000,
+}
+
+
+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}
+
+    return [_pick(row) for row in data]
+
+
+def _enrich_datasources_with_tsc(
+    datasources: SerializedAsset,
+    tsc_datasources: SerializedAsset,
+) -> SerializedAsset:
+    """
+    Enrich datasources with fields coming from TableauServerClient:
+    - project_luid
+    - webpage_url
+    """
+
+    mapping = {row["id"]: row for row in tsc_datasources}
+
+    for datasource in datasources:
+        if datasource["__typename"] != "PublishedDatasource":
+            # embedded datasources are bound to workbooks => no project
+            # embedded datasources cannot be accessed via URL => no webpage_url
+            continue
+        luid = datasource["luid"]
+        tsc_datasource = mapping[luid]
+        datasource["projectLuid"] = tsc_datasource["project_id"]
+        datasource["webpageUrl"] = tsc_datasource["webpage_url"]
+
+    return datasources
+
+
+def _enrich_workbooks_with_tsc(
+    workbooks: SerializedAsset,
+    tsc_workbooks: SerializedAsset,
+) -> SerializedAsset:
+    """
+    Enrich workbooks with fields coming from TableauServerClient:
+    - project_luid
+    """
+
+    mapping = {row["id"]: row for row in tsc_workbooks}
+
+    for workbook in workbooks:
+        luid = workbook["luid"]
+        tsc_workbook = mapping.get(luid)
+        if not tsc_workbook:
+            # it happens that a workbook is in Metadata API but not in TSC
+            # in this case, we push the workbook with default project
+            logger.warning(f"Workbook {luid} was not found in TSC")
+            workbook["projectLuid"] = None
+            continue
+
+        workbook["projectLuid"] = tsc_workbook["project_id"]
+
+    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
+
+
+class TableauRevampClient:
+    """
+    Connect to Tableau's API and extract assets.
+
+    Relies on TableauServerClient overlay:
+    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,
+    ):
+        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:
+        return "Tableau/API"
+
+    def _user_password_login(self) -> None:
+        """Login into Tableau using user and password"""
+        self._server.auth.sign_in(
+            TSC.TableauAuth(
+                self._credentials.user,
+                self._credentials.password,
+                site_id=self._credentials.site_id,
+            ),
+        )
+
+    def _pat_login(self) -> None:
+        """Login into Tableau using personal authentication token"""
+        self._server.auth.sign_in(
+            TSC.PersonalAccessTokenAuth(
+                self._credentials.token_name,
+                self._credentials.token,
+                site_id=self._credentials.site_id,
+            ),
+        )
+
+    def login(self) -> None:
+        """
+        Depending on the given credentials, logs-in using either:
+        - user/password
+        - token_name/value (Personal Access Token)
+        https://help.tableau.com/current/api/rest_api/en-us/REST/rest_api_concepts_auth.htm
+
+        Raises an error if none can be found
+        """
+
+        if self._credentials.user and self._credentials.password:
+            logger.info("Logging in using user and password authentication")
+            return self._user_password_login()
+
+        if self._credentials.token_name and self._credentials.token:
+            logger.info("Logging in using token authentication")
+            return self._pat_login()
+
+        raise ValueError(
+            "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)
+
+        return _enrich_datasources_with_tsc(datasources, datasource_projects)
+
+    def _fetch_workbooks(self) -> SerializedAsset:
+        asset = TableauRevampAsset.WORKBOOK
+
+        workbooks = self._fetch_from_metadata_api(asset)
+        workbook_projects = self._fetch_from_tsc(asset)
+
+        return _enrich_workbooks_with_tsc(workbooks, workbook_projects)
+
+    def fetch(
+        self,
+        asset: TableauRevampAsset,
+    ) -> SerializedAsset:
+        """
+        Extract the given Tableau Asset
+        """
+        if asset == TableauRevampAsset.DATASOURCE:
+            # both APIs are required to extract datasources
+            return self._fetch_datasources()
+
+        if asset == TableauRevampAsset.WORKBOOK:
+            # both 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)
+
+        # extract most assets via Metadata API
+        return self._fetch_from_metadata_api(asset)

castor_extractor/visualization/tableau_revamp/client/credentials.py

@@ -0,0 +1,104 @@
+from enum import Enum
+from typing import Dict, Literal, Optional, overload
+
+from ....utils import from_env
+
+_AUTH_ERROR_MSG = "Need either user and password or token_name and token"
+
+# 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]: ...
+
+
+@overload
+def get_value(
+    key: CredentialsKey, kwargs: dict, optional: Literal[True]
+) -> Optional[str]: ...
+
+
+@overload
+def get_value(
+    key: CredentialsKey, kwargs: dict, optional: Literal[False]
+) -> str: ...
+
+
+def get_value(
+    key: CredentialsKey,
+    kwargs: dict,
+    optional: bool = True,
+) -> Optional[str]:
+    """
+    Returns the value of the given key:
+    - from kwargs in priority
+    - from ENV otherwise
+    Raises an error if not found (unless optional)
+    """
+
+    if key.value in kwargs:
+        return kwargs[key.value]
+
+    env_key = CREDENTIALS_ENV[key]
+    return from_env(env_key, optional)
+
+
+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
+
+    @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),
+        )

castor_extractor/visualization/tableau_revamp/client/errors.py

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

castor_extractor/visualization/tableau_revamp/client/gql_queries.py

@@ -0,0 +1,134 @@
+from typing import Dict, Tuple
+
+from ..assets import TableauRevampAsset
+
+QUERY_TEMPLATE = """
+{{
+  {resource}Connection(first: {page_size}, after: AFTER_TOKEN_SIGNAL) {{
+    nodes {{ {fields}
+    }}
+    pageInfo {{
+      hasNextPage
+      endCursor
+    }}
+    totalCount
+  }}
+}}
+"""
+
+_COLUMNS_QUERY = """
+downstreamDashboards { id }
+downstreamFields { id }
+downstreamWorkbooks { id }
+id
+name
+table { id }
+"""
+
+_DASHBOARDS_QUERY = """
+createdAt
+id
+name
+path
+tags { name }
+updatedAt
+workbook { id }
+"""
+
+_DATASOURCES_QUERY = """
+__typename
+downstreamDashboards { id }
+downstreamWorkbooks { id }
+id
+name
+... on PublishedDatasource {
+    description
+    luid
+    owner { luid }
+    site { name }
+    tags { name }
+    uri
+}
+"""
+
+_TABLES_QUERY = """
+__typename
+downstreamDashboards { id }
+downstreamDatasources { id }
+downstreamWorkbooks { id }
+id
+name
+... on DatabaseTable {
+    connectionType
+    fullName
+    schema
+}
+... on CustomSQLTable {
+    query
+}
+"""
+
+
+_WORKBOOKS_QUERY = """
+createdAt
+description
+embeddedDatasources { id }
+id
+luid
+name
+owner { luid }
+site { name }
+tags { name }
+updatedAt
+uri
+"""
+
+_FIELDS_QUERY = """
+__typename
+datasource { id }
+description
+downstreamDashboards { id }
+downstreamWorkbooks { id }
+folderName
+id
+name
+dataType
+role
+"""
+
+
+_FIELDS_QUERY_WITH_COLUMNS = f"""
+{_FIELDS_QUERY}
+columns {{
+    name
+   table {{ name }}
+}}
+"""
+
+_SHEETS_QUERY = """
+containedInDashboards { id }
+createdAt
+id
+index
+name
+updatedAt
+upstreamFields { name }
+workbook { id }
+"""
+
+
+GQL_QUERIES: Dict[TableauRevampAsset, Tuple[str, str]] = {
+    TableauRevampAsset.COLUMN: ("columns", _COLUMNS_QUERY),
+    TableauRevampAsset.DASHBOARD: ("dashboards", _DASHBOARDS_QUERY),
+    TableauRevampAsset.DATASOURCE: ("datasources", _DATASOURCES_QUERY),
+    TableauRevampAsset.SHEET: ("sheets", _SHEETS_QUERY),
+    TableauRevampAsset.TABLE: ("tables", _TABLES_QUERY),
+    TableauRevampAsset.WORKBOOK: ("workbooks", _WORKBOOKS_QUERY),
+}
+
+FIELDS_QUERIES = (
+    ("binFields", _FIELDS_QUERY),
+    ("calculatedFields", _FIELDS_QUERY),
+    ("columnFields", _FIELDS_QUERY_WITH_COLUMNS),
+    ("groupFields", _FIELDS_QUERY),
+)

castor_extractor/visualization/tableau_revamp/client/tsc_fields.py

@@ -0,0 +1,34 @@
+from typing import Dict, Set
+
+from ..assets import TableauRevampAsset
+
+# list of fields to pick in TSC response
+TSC_FIELDS: Dict[TableauRevampAsset, Set[str]] = {
+    TableauRevampAsset.DATASOURCE: {
+        "id",
+        "project_id",
+        "webpage_url",
+    },
+    TableauRevampAsset.PROJECT: {
+        "description",
+        "id",
+        "name",
+        "parent_id",
+    },
+    TableauRevampAsset.USAGE: {
+        "name",
+        "total_views",
+        "workbook_id",
+    },
+    TableauRevampAsset.USER: {
+        "email",
+        "fullname",
+        "id",
+        "name",
+        "site_role",
+    },
+    TableauRevampAsset.WORKBOOK: {
+        "id",
+        "project_id",
+    },
+}

castor_extractor/visualization/tableau_revamp/constants.py

@@ -0,0 +1,5 @@
+TABLEAU_SERVER_VERSION = "3.5"
+
+DEFAULT_PAGE_SIZE = 100
+
+DEFAULT_TIMEOUT_SECONDS = 100

castor_extractor/visualization/tableau_revamp/extract.py

@@ -0,0 +1,53 @@
+import logging
+from typing import Iterable, Tuple
+
+from ...utils import (
+    OUTPUT_DIR,
+    current_timestamp,
+    deep_serialize,
+    from_env,
+    get_output_filename,
+    write_errors_logs,
+    write_json,
+    write_summary,
+)
+from .assets import TableauRevampAsset
+from .client import TableauRevampClient
+
+logger = logging.getLogger(__name__)
+
+
+def iterate_all_data(
+    client: TableauRevampClient,
+) -> Iterable[Tuple[TableauRevampAsset, list]]:
+    """Iterate over the extracted Data from Tableau"""
+
+    logger.info("Extracting USER from Tableau API")
+    yield TableauRevampAsset.USER, deep_serialize(
+        client.fetch(TableauRevampAsset.USER)
+    )
+
+
+def extract_all(client: TableauRevampClient, **kwargs: str) -> None:
+    """
+    Extract Data from tableau
+    Store data locally in files under the output_directory
+    If errors from Tableau's API are catch store them locally in file under the output_directory
+    """
+    output_directory = kwargs.get("output_directory") or from_env(OUTPUT_DIR)
+
+    timestamp = current_timestamp()
+
+    for key, data in iterate_all_data(client):
+        filename = get_output_filename(key.value, output_directory, timestamp)
+        write_json(filename, data)
+
+    write_summary(
+        output_directory,
+        timestamp,
+        base_url=client.base_url(),
+        client_name=client.name(),
+    )
+
+    if client.errors:
+        write_errors_logs(output_directory, timestamp, client.errors)

castor_extractor/warehouse/databricks/client.py

@@ -3,7 +3,7 @@ from datetime import date
 from functools import partial
 from typing import Any, Dict, List, Optional, Set
 
-from ...utils import at_midnight, date_after
+from ...utils import at_midnight, date_after, mapping_from_rows
 from ...utils.client.api import APIClient
 from ...utils.pager import PagerOnToken
 from ..abstract.time_filter import TimeFilter
@@ -88,15 +88,22 @@ class DatabricksClient(APIClient):
         )
 
     @staticmethod
-    def _match_table_with_user(table: dict, user_id_by_email: dict) -> dict:
+    def _match_table_with_user(table: dict, user_mapping: dict) -> dict:
         table_owner_email = table.get("owner_email")
         if not table_owner_email:
             return table
-        owner_external_id = user_id_by_email.get(table_owner_email)
+        owner_external_id = user_mapping.get(table_owner_email)
         if not owner_external_id:
             return table
         return {**table, "owner_external_id": owner_external_id}
 
+    @staticmethod
+    def _get_user_mapping(users: List[dict]) -> dict:
+        return {
+            **mapping_from_rows(users, "email", "id"),
+            **mapping_from_rows(users, "user_name", "id"),
+        }
+
     def tables_and_columns(
         self, schemas: List[dict], users: List[dict]
     ) -> TablesColumns:
@@ -105,11 +112,11 @@ class DatabricksClient(APIClient):
         """
         tables: List[dict] = []
         columns: List[dict] = []
-        user_id_by_email = {user.get("email"): user.get("id") for user in users}
+        user_mapping = self._get_user_mapping(users)
         for schema in schemas:
             t_to_add, c_to_add = self._tables_columns_of_schema(schema)
             t_with_owner = [
-                self._match_table_with_user(table, user_id_by_email)
+                self._match_table_with_user(table, user_mapping)
                 for table in t_to_add
             ]
             tables.extend(t_with_owner)