castor-extractor 0.22.1__py3-none-any.whl → 0.22.6__py3-none-any.whl

CHANGELOG.md

@@ -1,6 +1,27 @@
 
 # Changelog
 
+## 0.22.6 - 2025-01-21
+
+* bump dependencies: looker, databricks, deptry, ...
+
+
+## 0.22.5 - 2025-01-09
+
+* Databricks: validate and deduplicate lineage links
+
+## 0.22.4 - 2025-01-08
+
+* ThoughtSpot: extract answers
+
+## 0.22.3 - 2024-12-10
+
+* Databricks: extract lineage from system tables
+
+## 0.22.2 - 2024-12-06
+
+* Sigma: multithreading to retrieve lineage
+
 ## 0.22.1 - 2024-12-05
 
 * Salesforce: deduplicate tables

castor_extractor/file_checker/file.py

@@ -123,7 +123,7 @@ class FileCheckerRun:
         """
         if not self.verbose:
             return
-        header = f"Issues detected on Row #{index+1}\n"
+        header = f"Issues detected on Row #{index + 1}\n"
         for k, v in row.items():
             header += f"{str(k):<20} {str(v):<100}\n"
         self.logger.info(header + _SEPARATOR + issue_log + _SEPARATOR)

castor_extractor/visualization/powerbi/assets.py

@@ -8,6 +8,7 @@ class PowerBiAsset(ExternalAsset):
     DASHBOARDS = "dashboards"
     DATASETS = "datasets"
     DATASET_FIELDS = "dataset_fields"
+    DATASET_RELATIONSHIPS = "dataset_relationships"
     METADATA = "metadata"
     PAGES = "pages"
     REPORTS = "reports"
@@ -19,20 +20,9 @@ class PowerBiAsset(ExternalAsset):
     def optional(cls) -> set["PowerBiAsset"]:
         return {
             PowerBiAsset.DATASET_FIELDS,
+            PowerBiAsset.DATASET_RELATIONSHIPS,
             PowerBiAsset.PAGES,
             PowerBiAsset.TABLES,
             PowerBiAsset.TILES,
             PowerBiAsset.USERS,
         }
-
-
-# Assets extracted from the Metadata file
-# They are not directly fetched from the PowerBi api.
-METADATA_ASSETS = (
-    PowerBiAsset.DATASET_FIELDS,
-    PowerBiAsset.TABLES,
-    PowerBiAsset.TILES,
-    PowerBiAsset.USERS,
-)
-
-REPORTS_ASSETS = (PowerBiAsset.PAGES,)

castor_extractor/visualization/powerbi/extract.py

@@ -11,7 +11,7 @@ from ...utils import (
     write_json,
     write_summary,
 )
-from .assets import METADATA_ASSETS, REPORTS_ASSETS, PowerBiAsset
+from .assets import PowerBiAsset
 from .client import PowerbiClient, PowerbiCredentials
 
 logger = logging.getLogger(__name__)
@@ -21,7 +21,7 @@ def iterate_all_data(
     client: PowerbiClient,
 ) -> Iterable[tuple[PowerBiAsset, Union[list, dict]]]:
     for asset in PowerBiAsset:
-        if asset in METADATA_ASSETS + REPORTS_ASSETS:
+        if asset in PowerBiAsset.optional:
             continue
 
         logger.info(f"Extracting {asset.name} from API")

castor_extractor/visualization/sigma/client/client.py

@@ -1,9 +1,11 @@
 from collections.abc import Iterator
+from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from http import HTTPStatus
 from typing import Callable, Optional
 
 import requests
+from pydantic import BaseModel
 
 from ....utils import (
     APIClient,
@@ -12,6 +14,7 @@ from ....utils import (
     build_url,
     fetch_all_pages,
     handle_response,
+    retry,
 )
 from ..assets import SigmaAsset
 from .credentials import SigmaCredentials
@@ -29,7 +32,7 @@ _DATA_ELEMENTS: tuple[str, ...] = (
 )
 
 _AUTH_TIMEOUT_S = 60
-_SIGMA_TIMEOUT = 120
+_SIGMA_TIMEOUT_S = 300
 
 _SIGMA_HEADERS = {
     "Content-Type": _CONTENT_TYPE,
@@ -47,6 +50,23 @@ SIGMA_SAFE_MODE = RequestSafeMode(
     max_errors=_VOLUME_IGNORED,
     status_codes=_IGNORED_ERROR_CODES,
 )
+_THREADS_LINEAGE = 10  # empirically found; hit the rate limit with 20 workers
+_RETRY_NUMBER = 1
+_RETRY_BASE_MS = 60_000
+
+
+class LineageContext(BaseModel):
+    """all info needed to build the endpoint for lineage retrieval"""
+
+    workbook_id: str
+    element_id: str
+
+
+class Lineage(BaseModel):
+    """holds response from lineage API and context used to retrieve it"""
+
+    lineage: dict
+    context: LineageContext
 
 
 class SigmaBearerAuth(BearerAuth):
@@ -77,7 +97,7 @@ class SigmaClient(APIClient):
             host=credentials.host,
             auth=auth,
             headers=_SIGMA_HEADERS,
-            timeout=_SIGMA_TIMEOUT,
+            timeout=_SIGMA_TIMEOUT_S,
             safe_mode=safe_mode or SIGMA_SAFE_MODE,
         )
 
@@ -133,17 +153,51 @@ class SigmaClient(APIClient):
                     page=page, workbook_id=workbook_id
                 )
 
-    def _get_all_lineages(self, elements: list[dict]) -> Iterator[dict]:
+    @retry(
+        (ConnectionError,),
+        max_retries=_RETRY_NUMBER,
+        base_ms=_RETRY_BASE_MS,
+        log_exc_info=True,
+    )
+    def _get_lineage(self, lineage_context: LineageContext) -> Lineage:
+        """
+        return the lineage from API and other ids needed to characterize
+        lineage in castor
+        """
+        workbook_id = lineage_context.workbook_id
+        element_id = lineage_context.element_id
+        endpoint = SigmaEndpointFactory.lineage(workbook_id, element_id)
+        return Lineage(lineage=self._get(endpoint), context=lineage_context)
+
+    @staticmethod
+    def _lineage_context(elements: list[dict]) -> list[LineageContext]:
+        """
+        Helper function to prepare context for lineage retrieval.
+        Elements without associated columns are skipped.
+        """
+        contexts: list[LineageContext] = []
         for element in elements:
-            workbook_id = element["workbook_id"]
-            element_id = element["elementId"]
-            lineage = self._get(
-                endpoint=SigmaEndpointFactory.lineage(workbook_id, element_id)
+            if element.get("columns") is None:
+                continue
+
+            context = LineageContext(
+                workbook_id=element["workbook_id"],
+                element_id=element["elementId"],
             )
+            contexts.append(context)
+        return contexts
+
+    def _get_all_lineages(self, elements: list[dict]) -> Iterator[dict]:
+        lineage_context = self._lineage_context(elements)
+
+        with ThreadPoolExecutor(max_workers=_THREADS_LINEAGE) as executor:
+            results = executor.map(self._get_lineage, lineage_context)
+
+        for lineage in results:
             yield {
-                **lineage,
-                "workbook_id": workbook_id,
-                "element_id": element_id,
+                **lineage.lineage,
+                "workbook_id": lineage.context.workbook_id,
+                "element_id": lineage.context.element_id,
             }
 
     def _get_all_queries(self, workbooks: list[dict]) -> Iterator[dict]:

castor_extractor/visualization/thoughtspot/assets.py

@@ -4,6 +4,8 @@ from ...types import ExternalAsset
 class ThoughtspotAsset(ExternalAsset):
     """Thoughtspot assets"""
 
+    ANSWERS = "answers"
+    ANSWER_USAGES = "answer_usages"
     LIVEBOARDS = "liveboards"
+    LIVEBOARD_USAGES = "liveboard_usages"
     LOGICAL_TABLES = "logical_tables"
-    USAGES = "usages"

castor_extractor/visualization/thoughtspot/client/client.py

@@ -30,7 +30,12 @@ _THOUGHTSPOT_HEADERS = {
     "Content-Type": "application/json",
 }
 _METADATA_BATCH_SIZE = 100
-_USAGE_LIVEBOARD_ID = "bea79810-145f-4ad0-a02c-4177a6e7d861"
+# https://docs.thoughtspot.com/cloud/latest/object-usage-liveboard
+_OBJECT_USAGE_LIVEBOARD = "Object Usage"
+_ANSWER_USAGE_VIZ = "Answer Usage, by User"
+# https://docs.thoughtspot.com/cloud/latest/user-adoption
+_USER_ADOPTION_LIVEBOARD = "User Adoption"
+_LIVEBOARD_USAGE_VIZ = "Popular Liveboards Last 30 Days"
 # By default, no errors are ignored for the moment
 THOUGHTSPOT_SAFE_MODE = RequestSafeMode()
 
@@ -69,23 +74,39 @@ class ThoughtspotClient(APIClient):
     def _metadata_search(
         self,
         metadata_type: str,
+        identifier: Optional[str] = None,
     ) -> Iterator[dict]:
+        """
+        Yields assets of the given asset type, and optionally filters on a
+        specific identifier.
+        """
         offset = 0
+
         while True:
+            search_filters = {
+                "metadata": [{"type": metadata_type}],
+                "include_details": True,
+                "record_size": _METADATA_BATCH_SIZE,
+                "record_offset": offset,
+            }
+            if identifier:
+                search_filters["metadata"] = {
+                    "identifier": identifier,
+                    "type": metadata_type,
+                }
+
             metadata = self._post(
                 ThoughtspotEndpointFactory.metadata_search(),
-                data={
-                    "metadata": [{"type": metadata_type}],
-                    "include_details": True,
-                    "record_size": _METADATA_BATCH_SIZE,
-                    "record_offset": offset,
-                },
+                data=search_filters,
             )
             yield from metadata
             if len(metadata) < _METADATA_BATCH_SIZE:
                 break
             offset = offset + _METADATA_BATCH_SIZE
 
+    def _get_all_answers(self) -> Iterator[dict]:
+        yield from self._metadata_search(metadata_type="ANSWER")
+
     def _get_all_liveboards(self) -> Iterator[dict]:
         yield from self._metadata_search(metadata_type="LIVEBOARD")
 
@@ -95,26 +116,58 @@ class ThoughtspotClient(APIClient):
     def _get_all_tables(self) -> Iterator[dict]:
         yield from self._metadata_search(metadata_type="LOGICAL_TABLE")
 
-    def _get_liveboards_usages(self) -> Iterator[dict]:
+    def _get_usages(
+        self,
+        liveboard_name: str,
+        visualization_name: str,
+    ) -> Iterator[dict]:
+        """
+        Yields the data of a given visualization in the given liveboard.
+        ThoughtSpot maintains two system liveboards with stats about data usage,
+        which are useful to compute view counts and popularity.
+        """
+        usage_liveboard = next(
+            self._metadata_search(
+                metadata_type="LIVEBOARD", identifier=liveboard_name
+            )
+        )
+        liveboard_id = usage_liveboard["metadata_id"]
+
         data = self._post(
             endpoint=ThoughtspotEndpointFactory.liveboard(),
             headers={"Accept": "application/octet-stream"},
             data={
-                "metadata_identifier": _USAGE_LIVEBOARD_ID,
+                "metadata_identifier": liveboard_id,
                 "file_format": "CSV",
-                "visualization_identifiers": [
-                    "Popular Liveboards Last 30 Days"
-                ],
+                "visualization_identifiers": [visualization_name],
             },
             handler=lambda x: x.text,
         )
         yield from usage_liveboard_reader(data)
 
-    def fetch(self, asset: ThoughtspotAsset):
+    def _get_answer_usages(self) -> Iterator[dict]:
+        return self._get_usages(
+            liveboard_name=_OBJECT_USAGE_LIVEBOARD,
+            visualization_name=_ANSWER_USAGE_VIZ,
+        )
+
+    def _get_liveboards_usages(self) -> Iterator[dict]:
+        return self._get_usages(
+            liveboard_name=_USER_ADOPTION_LIVEBOARD,
+            visualization_name=_LIVEBOARD_USAGE_VIZ,
+        )
+
+    def fetch(self, asset: ThoughtspotAsset) -> Iterator[dict]:
+        if asset == ThoughtspotAsset.ANSWERS:
+            yield from self._get_all_answers()
+
+        if asset == ThoughtspotAsset.ANSWER_USAGES:
+            yield from self._get_answer_usages()
+
         if asset == ThoughtspotAsset.LIVEBOARDS:
             yield from self._get_all_liveboards()
 
-        if asset == ThoughtspotAsset.USAGES:
+        if asset == ThoughtspotAsset.LIVEBOARD_USAGES:
             yield from self._get_liveboards_usages()
 
         if asset == ThoughtspotAsset.LOGICAL_TABLES:

castor_extractor/visualization/thoughtspot/client/utils.py

@@ -1,13 +1,17 @@
 import csv
+import re
 from collections.abc import Iterator
 from io import StringIO
 
+_END_OF_GENERATED_TEXT = r'^""$'
+
 
 def usage_liveboard_reader(usage_liveboard_csv: str) -> Iterator[dict]:
     """
     Converts a CSV string into an iterator of dictionaries after
-    ignoring the first 6 lines, using the 7th line as the header.
-    First 6 lines looks like the following:
+    ignoring the generated text that preceeds the actual CSV header row.
+    The generated block ends with a row containing only two double quotes.
+    Here is an example:
 
         "Data extract produced by Castor on 09/19/2024 06:54"
         "Filters applied on data :"
@@ -15,11 +19,13 @@ def usage_liveboard_reader(usage_liveboard_csv: str) -> Iterator[dict]:
         "Pinboard NOT IN [mlm - availability pinboard,null]"
         "Timestamp >= 20240820 00:00:00 < 20240919 00:00:00"
         "Timestamp >= 20240919 00:00:00 < 20240920 00:00:00"
+        ""
 
     """
     csv_file = StringIO(usage_liveboard_csv)
 
-    for _ in range(7):
-        next(csv_file)
+    line = next(csv_file)
+    while not re.match(_END_OF_GENERATED_TEXT, line.strip()):
+        line = next(csv_file)
 
     yield from csv.DictReader(csv_file)

castor_extractor/visualization/thoughtspot/client/utils_test.py

@@ -2,7 +2,7 @@ from .utils import (
     usage_liveboard_reader,
 )
 
-VALID_CSV = '''"Data extract produced by Castor on 09/19/2024 06:54"
+VALID_CSV_1 = '''"Data extract produced by Castor on 09/19/2024 06:54"
 "Filters applied on data :"
 "User Action IN [pinboard_embed_view,pinboard_tspublic_no_runtime_filter,pinboard_tspublic_runtime_filter,pinboard_view]"
 "Pinboard NOT IN [mlm - availability pinboard,null]"
@@ -16,6 +16,13 @@ VALID_CSV = '''"Data extract produced by Castor on 09/19/2024 06:54"
 "September test","25","2"'''
 
 
+VALID_CSV_2 = '''"Data extract produced by Castor on 01/07/2025 16:07"
+"Filters applied on data :"
+"Timestamp >= 20241208 00:00:00 < 20250107 00:00:00"
+""
+"Answer name","User name","Number of unique users","Count of object interactions"
+"toto","tata","1","666"'''
+
 # Invalid CSV input (missing data rows)
 INVALID_CSV = '''"Data extract produced by Castor on 09/19/2024 06:54"
 "Filters applied on data :"
@@ -27,7 +34,7 @@ INVALID_CSV = '''"Data extract produced by Castor on 09/19/2024 06:54"
 
 
 def test_usage_liveboard_reader():
-    expected_output = [
+    expected_output_1 = [
         {
             "Pinboard": "Market Report",
             "Pinboard Views": "559",
@@ -49,9 +56,20 @@ def test_usage_liveboard_reader():
             "Unique Number of User": "2",
         },
     ]
+    expected_output_2 = [
+        {
+            "Answer name": "toto",
+            "User name": "tata",
+            "Number of unique users": "1",
+            "Count of object interactions": "666",
+        }
+    ]
+
+    result = list(usage_liveboard_reader(VALID_CSV_1))
+    assert result == expected_output_1
 
-    result = list(usage_liveboard_reader(VALID_CSV))
-    assert result == expected_output
+    result = list(usage_liveboard_reader(VALID_CSV_2))
+    assert result == expected_output_2
 
     result = list(usage_liveboard_reader(INVALID_CSV))
     assert result == []  # Expect an empty result since there is no data

castor_extractor/warehouse/abstract/extract.py

@@ -60,7 +60,7 @@ class SQLExtractionProcessor:
         total = len(queries)
 
         for i, query in enumerate(queries):
-            logger.info(f"Extracting {asset.value}: query {i+1}/{total}")
+            logger.info(f"Extracting {asset.value}: query {i + 1}/{total}")
             # concatenate results of all queries
             data = chain(data, self._fetch(query))
 

castor_extractor/warehouse/bigquery/client.py

@@ -66,9 +66,9 @@ class BigQueryClient(SqlalchemyClient):
         return BIGQUERY_URI
 
     def _credentials(self) -> Credentials:
-        assert (
-            CREDENTIALS_INFO_KEY in self._options
-        ), "Missing BigQuery credentials in engine's options"
+        assert CREDENTIALS_INFO_KEY in self._options, (
+            "Missing BigQuery credentials in engine's options"
+        )
         credentials = self._options[CREDENTIALS_INFO_KEY]
         return Credentials.from_service_account_info(credentials)
 

castor_extractor/warehouse/databricks/api_client.py

@@ -1,8 +1,6 @@
 import logging
-from collections.abc import Iterator
 from functools import partial
-from http import HTTPStatus
-from typing import Optional
+from typing import Iterator, Optional
 
 import requests
 
@@ -14,16 +12,14 @@ from ...utils import (
     fetch_all_pages,
     handle_response,
     retry,
-    retry_request,
     safe_mode,
 )
 from ..abstract import TimeFilter
 from .credentials import DatabricksCredentials
 from .endpoints import DatabricksEndpointFactory
 from .format import DatabricksFormatter, TagMapping
-from .lineage import single_column_lineage_links, single_table_lineage_links
 from .pagination import DATABRICKS_PAGE_SIZE, DatabricksPagination
-from .types import TablesColumns, TimestampedLink
+from .types import TablesColumns
 from .utils import hourly_time_filters
 
 logger = logging.getLogger(__name__)
@@ -132,60 +128,6 @@ class DatabricksAPIClient(APIClient):
             column_tags=column_tags,
         )
 
-    @safe_mode(safe_lineage_params, lambda: [])
-    @retry(
-        exceptions=_RETRY_EXCEPTIONS,
-        max_retries=_RETRY_ATTEMPTS,
-        base_ms=_RETRY_BASE_MS,
-    )
-    @retry_request(
-        status_codes=(HTTPStatus.TOO_MANY_REQUESTS,),
-        max_retries=_RETRY_ATTEMPTS,
-    )
-    def get_single_column_lineage(
-        self,
-        names: tuple[str, str],
-    ) -> list[TimestampedLink]:
-        """
-        Helper function used in get_lineage_links.
-        Call data lineage API and return the content of the result
-
-        eg table_path: broward_prd.bronze.account_adjustments
-        FYI: Maximum rate of 10 requests per SECOND
-        """
-        table_path, column_name = names
-        payload = {
-            "table_name": table_path,
-            "column_name": column_name,
-            "include_entity_lineage": True,
-        }
-        content = self._get(
-            DatabricksEndpointFactory.column_lineage(), params=payload
-        )
-        column_path = f"{table_path}.{column_name}"
-        return single_column_lineage_links(column_path, content)
-
-    @safe_mode(safe_lineage_params, lambda: [])
-    @retry(
-        exceptions=_RETRY_EXCEPTIONS,
-        max_retries=_RETRY_ATTEMPTS,
-        base_ms=_RETRY_BASE_MS,
-    )
-    def get_single_table_lineage(
-        self, table_path: str
-    ) -> list[TimestampedLink]:
-        """
-        Helper function used in get_lineage_links.
-        Call data lineage API and return the content of the result
-        eg table_path: broward_prd.bronze.account_adjustments
-        FYI: Maximum rate of 50 requests per SECOND
-        """
-        payload = {"table_name": table_path, "include_entity_lineage": True}
-        content = self._get(
-            DatabricksEndpointFactory.table_lineage(), params=payload
-        )
-        return single_table_lineage_links(table_path, content)
-
     @safe_mode(safe_query_params, lambda: [])
     @retry(
         exceptions=_RETRY_EXCEPTIONS,

castor_extractor/warehouse/databricks/client.py

@@ -1,17 +1,14 @@
 import logging
-from concurrent.futures import ThreadPoolExecutor
 from typing import Optional
 
-from ...utils import (
-    mapping_from_rows,
-)
+from ...utils import mapping_from_rows
 from ..abstract import TimeFilter
 from .api_client import DatabricksAPIClient
 from .credentials import DatabricksCredentials
+from .enums import TagEntity
 from .format import DatabricksFormatter
-from .lineage import deduplicate_lineage, paths_for_column_lineage
-from .sql_client import DatabricksSQLClient, TagEntity
-from .types import TablesColumns, TimestampedLink
+from .sql_client import DatabricksSQLClient
+from .types import TablesColumns
 
 logger = logging.getLogger(__name__)
 
@@ -95,46 +92,6 @@ class DatabricksClient:
             columns.extend(c_to_add)
         return tables, columns
 
-    def table_lineage(self, tables: list[dict]) -> list[dict]:
-        """
-        Wrapper function that retrieves all table lineage
-        """
-        # retrieve table lineage
-        with ThreadPoolExecutor(max_workers=_THREADS_TABLE_LINEAGE) as executor:
-            table_paths = [
-                ".".join([table["schema_id"], table["table_name"]])
-                for table in tables
-            ]
-            results = executor.map(
-                self.api_client.get_single_table_lineage, table_paths
-            )
-        lineages = [link for links in results for link in links]
-        deduplicated = deduplicate_lineage(lineages)
-        return self.formatter.format_lineage(deduplicated)
-
-    def column_lineage(
-        self, tables: list[dict], columns: list[dict], table_lineage: list[dict]
-    ) -> list[dict]:
-        """
-        Wrapper function that retrieves all column lineage
-        we only try to retrieve column lineage if we found table lineage
-        """
-        candidate_paths = paths_for_column_lineage(
-            tables, columns, table_lineage
-        )
-        # retrieve column lineage
-        with ThreadPoolExecutor(
-            max_workers=_THREADS_COLUMN_LINEAGE
-        ) as executor:
-            results = executor.map(
-                self.api_client.get_single_column_lineage, candidate_paths
-            )
-        lineages: list[TimestampedLink] = [
-            link for links in results for link in links
-        ]
-        deduplicated = deduplicate_lineage(lineages)
-        return self.formatter.format_lineage(deduplicated)
-
     def queries(self, time_filter: Optional[TimeFilter] = None) -> list[dict]:
         return self.api_client.queries(time_filter)
 

castor_extractor/warehouse/databricks/client_test.py

@@ -1,14 +1,4 @@
-from unittest.mock import Mock, patch
-
-from .client import (
-    DatabricksClient,
-)
-from .test_constants import (
-    CLOSER_DATE,
-    MOCK_TABLES_FOR_TABLE_LINEAGE,
-    OLDER_DATE,
-    TABLE_LINEAGE_SIDE_EFFECT,
-)
+from .client import DatabricksClient
 
 
 class MockDatabricksClient(DatabricksClient):
@@ -48,27 +38,3 @@ def test_DatabricksClient__match_table_with_user():
     table_without_owner = {"id": 1, "owner_email": None}
     actual = client._match_table_with_user(table_without_owner, user_mapping)
     assert actual == table_without_owner
-
-
-@patch(
-    "source.packages.extractor.castor_extractor.warehouse.databricks.client.DatabricksAPIClient._get",
-    side_effect=TABLE_LINEAGE_SIDE_EFFECT,
-)
-def test_DatabricksClient_table_lineage(mock_get):
-    client = DatabricksClient(Mock())
-
-    lineage = client.table_lineage(MOCK_TABLES_FOR_TABLE_LINEAGE)
-    assert len(lineage) == 2
-
-    expected_link_1 = {
-        "parent_path": "dev.silver.pre_analytics",
-        "child_path": "dev.silver.analytics",
-        "timestamp": OLDER_DATE,
-    }
-    expected_link_2 = {
-        "parent_path": "dev.bronze.analytics",
-        "child_path": "dev.silver.analytics",
-        "timestamp": CLOSER_DATE,
-    }
-    assert expected_link_1 in lineage
-    assert expected_link_2 in lineage