castor-extractor 0.16.6__py3-none-any.whl → 0.16.11__py3-none-any.whl

castor_extractor/visualization/tableau_revamp/client/client.py

@@ -2,6 +2,7 @@ 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
@@ -12,7 +13,7 @@ from ..constants import (
 )
 from .credentials import TableauRevampCredentials
 from .errors import TableauApiError
-from .gql_queries import GQL_QUERIES, QUERY_TEMPLATE
+from .gql_queries import FIELDS_QUERIES, GQL_QUERIES, QUERY_TEMPLATE
 from .tsc_fields import TSC_FIELDS
 
 logger = logging.getLogger(__name__)
@@ -27,13 +28,18 @@ _TSC_ASSETS = (
     TableauRevampAsset.USAGE,
 )
 
+# increase the value when extraction is too slow
+# decrease the value when timeouts arise
 _CUSTOM_PAGE_SIZE: Dict[TableauRevampAsset, int] = {
+    # fields and columns are light but volumes are bigger
+    TableauRevampAsset.COLUMN: 200,
     TableauRevampAsset.FIELD: 1000,
+    TableauRevampAsset.TABLE: 50,
 }
 
 
 def _pick_fields(
-    data: SerializedAsset,
+    data: Pager,
     asset: TableauRevampAsset,
 ) -> SerializedAsset:
     fields = TSC_FIELDS[asset]
@@ -44,7 +50,7 @@ def _pick_fields(
     return [_pick(row) for row in data]
 
 
-def _enrich_with_tsc(
+def _enrich_datasources_with_tsc(
     datasources: SerializedAsset,
     tsc_datasources: SerializedAsset,
 ) -> SerializedAsset:
@@ -69,6 +75,32 @@ def _enrich_with_tsc(
     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,
@@ -176,29 +208,32 @@ class TableauRevampClient:
         asset: TableauRevampAsset,
     ) -> SerializedAsset:
 
-        if asset == TableauRevampAsset.USER:
-            data = TSC.Pager(self._server.users)
+        if asset == TableauRevampAsset.DATASOURCE:
+            data = TSC.Pager(self._server.datasources)
 
         elif asset == TableauRevampAsset.PROJECT:
             data = TSC.Pager(self._server.projects)
 
-        elif asset == TableauRevampAsset.DATASOURCE:
-            data = TSC.Pager(self._server.datasources)
-
         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 _fetch_from_metadata_api(
+    def _run_graphql_query(
         self,
-        asset: TableauRevampAsset,
+        resource: str,
+        fields: str,
+        page_size: int = DEFAULT_PAGE_SIZE,
     ) -> SerializedAsset:
-        resource, fields = GQL_QUERIES[asset]
-        page_size = _CUSTOM_PAGE_SIZE.get(asset) or DEFAULT_PAGE_SIZE
         query = QUERY_TEMPLATE.format(
             resource=resource,
             fields=fields,
@@ -207,13 +242,40 @@ class TableauRevampClient:
         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_with_tsc(datasources, datasource_projects)
+        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,
@@ -226,6 +288,10 @@ class TableauRevampClient:
             # 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)

castor_extractor/visualization/tableau_revamp/client/gql_queries.py

@@ -18,7 +18,11 @@ QUERY_TEMPLATE = """
 
 _COLUMNS_QUERY = """
 downstreamDashboards { id }
-downstreamFields { id }
+downstreamFields {
+    id
+    __typename
+    datasource { id }
+}
 downstreamWorkbooks { id }
 id
 name
@@ -37,12 +41,10 @@ workbook { id }
 
 _DATASOURCES_QUERY = """
 __typename
-createdAt
 downstreamDashboards { id }
 downstreamWorkbooks { id }
 id
 name
-updatedAt
 ... on PublishedDatasource {
     description
     luid
@@ -64,7 +66,6 @@ name
     connectionType
     fullName
     schema
-    tableType
 }
 ... on CustomSQLTable {
     query
@@ -80,7 +81,6 @@ id
 luid
 name
 owner { luid }
-projectLuid
 site { name }
 tags { name }
 updatedAt
@@ -96,16 +96,17 @@ downstreamWorkbooks { id }
 folderName
 id
 name
-... on DataField {
-    dataType
-    role
-}
-... on ColumnField {
-    columns {
-        name
-        table { name }
-    }
-}
+dataType
+role
+"""
+
+
+_FIELDS_QUERY_WITH_COLUMNS = f"""
+{_FIELDS_QUERY}
+columns {{
+    name
+   table {{ name }}
+}}
 """
 
 _SHEETS_QUERY = """
@@ -124,8 +125,14 @@ GQL_QUERIES: Dict[TableauRevampAsset, Tuple[str, str]] = {
     TableauRevampAsset.COLUMN: ("columns", _COLUMNS_QUERY),
     TableauRevampAsset.DASHBOARD: ("dashboards", _DASHBOARDS_QUERY),
     TableauRevampAsset.DATASOURCE: ("datasources", _DATASOURCES_QUERY),
-    TableauRevampAsset.FIELD: ("fields", _FIELDS_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

@@ -27,4 +27,8 @@ TSC_FIELDS: Dict[TableauRevampAsset, Set[str]] = {
         "name",
         "site_role",
     },
+    TableauRevampAsset.WORKBOOK: {
+        "id",
+        "project_id",
+    },
 }

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)

castor_extractor/warehouse/databricks/client_test.py

@@ -66,15 +66,34 @@ def test_DatabricksClient__keep_catalog():
     assert not client._keep_catalog("something_unknown")
 
 
+def test_DatabricksClient__get_user_mapping():
+    client = MockDatabricksClient()
+    users = [
+        {"id": "both", "email": "hello@world.com", "user_name": "hello world"},
+        {"id": "no_email", "email": "", "user_name": "no email"},
+        {"id": "no_name", "email": "no@name.fr", "user_name": ""},
+        {"id": "no_both", "email": "", "user_name": ""},
+        {"id": "", "email": "no@id.com", "user_name": "no id"},
+    ]
+    expected = {
+        "hello@world.com": "both",
+        "hello world": "both",
+        "no@name.fr": "no_name",
+        "no email": "no_email",
+    }
+    mapping = client._get_user_mapping(users)
+    assert mapping == expected
+
+
 def test_DatabricksClient__match_table_with_user():
     client = MockDatabricksClient()
-    users_by_email = {"bob@castordoc.com": 3}
+    user_mapping = {"bob@castordoc.com": 3}
 
     table = {"id": 1, "owner_email": "bob@castordoc.com"}
-    table_with_owner = client._match_table_with_user(table, users_by_email)
+    table_with_owner = client._match_table_with_user(table, user_mapping)
 
     assert table_with_owner == {**table, "owner_external_id": 3}
 
     table_without_owner = {"id": 1, "owner_email": None}
-    actual = client._match_table_with_user(table_without_owner, users_by_email)
+    actual = client._match_table_with_user(table_without_owner, user_mapping)
     assert actual == table_without_owner

castor_extractor/warehouse/databricks/format.py

@@ -127,13 +127,17 @@ class DatabricksFormatter:
                 return email["value"]
         return emails[0]["value"]
 
+    def _email(self, user: dict) -> Optional[str]:
+        emails = user.get("emails")
+        return self._primary(emails) if emails else None
+
     def format_user(self, raw_users: List[dict]) -> List[dict]:
         users = []
         for user in raw_users:
             users.append(
                 {
                     "id": user["id"],
-                    "email": self._primary(user["emails"]),
+                    "email": self._email(user),
                     "first_name": None,
                     "last_name": user.get("displayName") or user["userName"],
                     "user_name": user["userName"],

castor_extractor/warehouse/salesforce/client.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict, Iterator, List
+from typing import Dict, Iterator, List, Tuple
 
 from tqdm import tqdm  # type: ignore
 
@@ -96,17 +96,19 @@ class SalesforceClient(SalesforceBaseClient):
         """
         sobjects = self.fetch_sobjects()
         logger.info(f"Extracted {len(sobjects)} sobjects")
-        return self.formatter.tables(sobjects)
+        return list(self.formatter.tables(sobjects))
 
     def columns(
-        self, sobject_names: List[str], show_progress: bool = True
+        self, sobject_names: List[Tuple[str, str]], show_progress: bool = True
     ) -> List[dict]:
         """
         Get salesforce sobject fields as columns
         show_progress: optionally deactivate the tqdm progress bar
         """
         sobject_fields: Dict[str, List[dict]] = dict()
-        for sobject_name in tqdm(sobject_names, disable=not show_progress):
-            fields = self.fetch_fields(sobject_name)
-            sobject_fields[sobject_name] = fields
+        for api_name, table_name in tqdm(
+            sobject_names, disable=not show_progress
+        ):
+            fields = self.fetch_fields(api_name)
+            sobject_fields[table_name] = fields
         return self.formatter.columns(sobject_fields)

castor_extractor/warehouse/salesforce/extract.py

@@ -72,8 +72,8 @@ class SalesforceExtractionProcessor:
         catalog_locations[WarehouseAsset.TABLE.value] = location
         logger.info(f"Extracted {len(tables)} tables to {location}")
 
-        table_names = [t["table_name"] for t in tables]
-        columns = self._client.columns(table_names, show_progress)
+        sobject_names = [(t["api_name"], t["table_name"]) for t in tables]
+        columns = self._client.columns(sobject_names, show_progress)
         location = self._storage.put(WarehouseAsset.COLUMN.value, columns)
         catalog_locations[WarehouseAsset.COLUMN.value] = location
         logger.info(f"Extracted {len(columns)} columns to {location}")

castor_extractor/warehouse/salesforce/format.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, Iterator, List
 
 from .constants import SCHEMA_NAME
 
@@ -35,17 +35,35 @@ def _to_column_payload(field: dict, position: int, table_name: str) -> dict:
     }
 
 
-def _to_table_payload(table: dict) -> dict:
+def _to_table_payload(sobject: dict, table_name: str) -> dict:
     return {
-        "id": table["QualifiedApiName"],
+        "id": table_name,
+        "api_name": sobject["QualifiedApiName"],
+        "label": sobject["Label"],
         "schema_id": SCHEMA_NAME,
-        "table_name": table["QualifiedApiName"],
+        "table_name": table_name,
         "description": "",
         "tags": [],
         "type": "TABLE",
     }
 
 
+def _merge_label_and_api_name(sobject: dict) -> dict:
+    label = sobject["Label"]
+    api_name = sobject["QualifiedApiName"]
+    table_name = f"{label} ({api_name})"
+    return _to_table_payload(sobject, table_name)
+
+
+def _by_label(sobjects: List[dict]) -> Dict[str, List[dict]]:
+    by_label: Dict[str, List[dict]] = dict()
+    for sobject in sobjects:
+        label = sobject["Label"]
+        similar_sobjects = by_label.setdefault(label, [])
+        similar_sobjects.append(sobject)
+    return by_label
+
+
 class SalesforceFormatter:
     """
     Helper functions that format the response in the format to be exported as
@@ -53,9 +71,18 @@ class SalesforceFormatter:
     """
 
     @staticmethod
-    def tables(sobjects: List[dict]) -> List[dict]:
-        """formats the raw list of sobjects to tables"""
-        return [_to_table_payload(s) for s in sobjects]
+    def tables(sobjects: List[dict]) -> Iterator[dict]:
+        """
+        formats the raw list of sobjects to tables
+        if two tables share the same label, then we add the api name as well
+        """
+        by_label = _by_label(sobjects)
+        for label, similars in by_label.items():
+            if len(similars) > 1:
+                yield from [_merge_label_and_api_name(s) for s in similars]
+            else:
+                sobject = similars[0]  # unique sobject on label
+                yield _to_table_payload(sobject, label)
 
     @staticmethod
     def columns(sobject_fields: Dict[str, List[dict]]) -> List[dict]:

castor_extractor/warehouse/salesforce/format_test.py

@@ -1,4 +1,21 @@
-from .format import _field_description
+from typing import Dict, Tuple
+
+from .format import (
+    SCHEMA_NAME,
+    SalesforceFormatter,
+    _by_label,
+    _field_description,
+    _merge_label_and_api_name,
+)
+
+
+def _example_sobjects() -> Tuple[Dict[str, str], ...]:
+    """Returns 4 sobjects with 2 sharing the same label"""
+    a = {"Label": "a", "QualifiedApiName": "a_one"}
+    b = {"Label": "b", "QualifiedApiName": "b"}
+    c = {"Label": "c", "QualifiedApiName": "c"}
+    a_prime = {"Label": "a", "QualifiedApiName": "a_two"}
+    return a, b, c, a_prime
 
 
 def test__field_description():
@@ -30,3 +47,34 @@ def test__field_description():
         "- Data Sensitivity Level: bam"
     )
     assert description == expected
+
+
+def test__merge_label_and_api_name():
+    sobject = {"Label": "foo", "QualifiedApiName": "bar"}
+    payload = _merge_label_and_api_name(sobject)
+    expected_name = "foo (bar)"
+    assert payload == {
+        "id": expected_name,
+        "api_name": "bar",
+        "label": "foo",
+        "schema_id": SCHEMA_NAME,
+        "table_name": expected_name,
+        "description": "",
+        "tags": [],
+        "type": "TABLE",
+    }
+
+
+def test__by_label():
+    a, b, c, a_prime = _example_sobjects()
+    sobjects = [a, b, c, a_prime]
+    by_label = _by_label(sobjects)
+    assert by_label == {"a": [a, a_prime], "b": [b], "c": [c]}
+
+
+def test_salesforce_formatter_tables():
+    sobjects = [*_example_sobjects()]
+    tables = SalesforceFormatter.tables(sobjects)
+    expected_names = {"a (a_one)", "a (a_two)", "b", "c"}
+    payload_names = {t["table_name"] for t in tables}
+    assert payload_names == expected_names

{castor_extractor-0.16.6.dist-info → castor_extractor-0.16.11.dist-info}/METADATA

@@ -1,12 +1,12 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.16.6
+Version: 0.16.11
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
 Author: Castor
 Author-email: support@castordoc.com
-Requires-Python: >=3.8,<3.12
+Requires-Python: >=3.8,<3.13
 Classifier: License :: Other/Proprietary License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
@@ -14,6 +14,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: all
 Provides-Extra: bigquery
 Provides-Extra: dbt
@@ -34,8 +35,10 @@ Requires-Dist: google-cloud-core (>=2.1.0,<3.0.0)
 Requires-Dist: google-cloud-storage (>=2,<3)
 Requires-Dist: google-resumable-media (>=2.0.3,<3.0.0)
 Requires-Dist: googleapis-common-protos (>=1.53.0,<2.0.0)
-Requires-Dist: looker-sdk (>=22.4.0,<=23.0.0) ; extra == "looker" or extra == "all"
+Requires-Dist: looker-sdk (>=23.0.0) ; extra == "looker" or extra == "all"
 Requires-Dist: msal (>=1.20.0,<2.0.0) ; extra == "powerbi" or extra == "all"
+Requires-Dist: numpy (<1.25) ; python_version >= "3.8" and python_version < "3.9"
+Requires-Dist: numpy (>=1.26,<2) ; python_version >= "3.12" and python_version < "3.13"
 Requires-Dist: psycopg2-binary (>=2.0.0,<3.0.0) ; extra == "metabase" or extra == "postgres" or extra == "redshift" or extra == "all"
 Requires-Dist: pycryptodome (>=3.0.0,<4.0.0) ; extra == "metabase" or extra == "all"
 Requires-Dist: pydantic (>=2.6,<3.0)