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

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/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

castor_extractor/warehouse/databricks/credentials.py

@@ -1,24 +1,22 @@
 from dataclasses import field
-from typing import Optional
 
-from pydantic.dataclasses import dataclass
-from pydantic_settings import SettingsConfigDict
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
 DATABRICKS_ENV_PREFIX = "CASTOR_DATABRICKS_"
 
 
-@dataclass
-class DatabricksCredentials:
+class DatabricksCredentials(BaseSettings):
     """
     Credentials needed by Databricks client
     Requires:
     - host
+    - http_path
     - token
     """
 
     host: str
+    http_path: str
     token: str = field(metadata={"sensitive": True})
-    http_path: Optional[str] = field(default=None)
 
     model_config = SettingsConfigDict(
         env_prefix=DATABRICKS_ENV_PREFIX,

castor_extractor/warehouse/databricks/enums.py

@@ -0,0 +1,15 @@
+from enum import Enum
+
+
+class LineageEntity(Enum):
+    """Entities that can be linked in Databricks lineage"""
+
+    COLUMN = "COLUMN"
+    TABLE = "TABLE"
+
+
+class TagEntity(Enum):
+    """Entities that can be tagged in Databricks"""
+
+    COLUMN = "COLUMN"
+    TABLE = "TABLE"

castor_extractor/warehouse/databricks/extract.py

@@ -1,4 +1,5 @@
 import logging
+from datetime import date
 from typing import Optional
 
 from ...utils import AbstractStorage, LocalStorage, write_summary
@@ -16,6 +17,7 @@ from ..abstract import (
 )
 from .client import DatabricksClient
 from .credentials import DatabricksCredentials
+from .enums import LineageEntity
 
 DATABRICKS_ASSETS: SupportedAssets = {
     WarehouseAssetGroup.ADDITIONAL_LINEAGE: ADDITIONAL_LINEAGE_ASSETS,
@@ -32,6 +34,12 @@ OTimeFilter = Optional[TimeFilter]
 Paths = dict[str, str]
 
 
+def _day(time_filter: OTimeFilter) -> date:
+    if not time_filter:
+        return TimeFilter.default().day
+    return time_filter.day
+
+
 class DatabricksExtractionProcessor:
     """Databricks' API-based extraction management"""
 
@@ -96,22 +104,18 @@ class DatabricksExtractionProcessor:
         logger.info(f"Extracted {len(columns)} columns to {location}")
         return catalog_locations
 
-    def extract_lineage(self) -> Paths:
+    def extract_lineage(self, time_filter: OTimeFilter = None) -> Paths:
         if self._should_not_reextract(WarehouseAssetGroup.ADDITIONAL_LINEAGE):
             return self._existing_group_paths(
                 WarehouseAssetGroup.ADDITIONAL_LINEAGE
             )
         lineage_locations: dict[str, str] = dict()
 
-        # extract catalog
-        databases = self._client.databases()
-        schemas = self._client.schemas(databases)
-        users = self._client.users()
-        tables, columns = self._client.tables_and_columns(schemas, users)
-        logger.info("Extracted pre-requisite catalog. Next comes lineage")
+        day = _day(time_filter)
+        client = self._client.sql_client
 
         # extract table lineage
-        table_lineage = self._client.table_lineage(tables)
+        table_lineage = client.get_lineage(LineageEntity.TABLE, day)
         table_lineage_key = WarehouseAsset.ADDITIONAL_TABLE_LINEAGE.value
         location = self._storage.put(table_lineage_key, table_lineage)
         lineage_locations[table_lineage_key] = location
@@ -119,9 +123,7 @@ class DatabricksExtractionProcessor:
         logger.info(msg)
 
         # extract column lineage
-        column_lineage = self._client.column_lineage(
-            tables, columns, table_lineage
-        )
+        column_lineage = client.get_lineage(LineageEntity.COLUMN, day)
         column_lineage_key = WarehouseAsset.ADDITIONAL_COLUMN_LINEAGE.value
         location = self._storage.put(column_lineage_key, column_lineage)
         lineage_locations[column_lineage_key] = location

castor_extractor/warehouse/databricks/lineage.py

@@ -1,141 +1,69 @@
-from typing import cast
+from typing import Iterable, Optional
 
-from .types import Link, Ostr, OTimestampedLink, TimestampedLink
+from .enums import LineageEntity
 
 
-class LineageLinks:
+class LineageProcessor:
     """
     helper class that handles lineage deduplication and filtering
     """
 
-    def __init__(self):
-        self.lineage: dict[Link, Ostr] = dict()
+    def __init__(self, lineage_entity: LineageEntity):
+        self.lineage_entity = lineage_entity
 
-    def add(self, timestamped_link: TimestampedLink) -> None:
-        """
-        keep the most recent lineage link, adding to `self.lineage`
-        """
-        parent, child, timestamp = timestamped_link
-        link = (parent, child)
-        if not self.lineage.get(link):
-            self.lineage[link] = timestamp
-            return
-
-        if not timestamp:
-            return
-        # keep most recent link; cast for mypy
-        recent = max(cast(str, self.lineage[link]), cast(str, timestamp))
-        self.lineage[link] = recent
+        self.lineage: dict[tuple[str, str], dict] = dict()
 
+    def _parent_path(self, link) -> Optional[str]:
+        if self.lineage_entity == LineageEntity.TABLE:
+            return link["source_table_full_name"]
 
-def _to_table_path(table: dict) -> Ostr:
-    if table.get("name"):
-        return f"{table['catalog_name']}.{table['schema_name']}.{table['name']}"
-    return None
+        source_table = link["source_table_full_name"]
+        source_column = link["source_column_name"]
+        if not (source_table and source_column):
+            return None
 
+        return f"{source_table}.{source_column}"
 
-def _to_column_path(column: dict) -> Ostr:
-    if column.get("name"):
-        return f"{column['catalog_name']}.{column['schema_name']}.{column['table_name']}.{column['name']}"
-    return None
+    def _child_path(self, link) -> Optional[str]:
+        if self.lineage_entity == LineageEntity.TABLE:
+            return link["target_table_full_name"]
 
+        target_table = link["target_table_full_name"]
+        target_column = link["target_column_name"]
+        if not (target_table and target_column):
+            return None
 
-def _link(path_from: Ostr, path_to: Ostr, timestamp: Ostr) -> OTimestampedLink:
-    """exclude missing path and self-lineage"""
-    if (not path_from) or (not path_to):
-        return None
-    is_self_lineage = path_from.lower() == path_to.lower()
-    if is_self_lineage:
-        return None
-    return path_from, path_to, timestamp
+        return f"{target_table}.{target_column}"
 
+    def add(self, link: dict) -> None:
+        """
+        If the parent and child paths are valid, keeps the most recent lineage
+        link in the `self.lineage` map.
+        """
+        parent = self._parent_path(link)
+        child = self._child_path(link)
+        timestamp = link["event_time"]
 
-def single_table_lineage_links(
-    table_path: str, single_table_lineage: dict
-) -> list[TimestampedLink]:
-    """
-    process databricks lineage API response for a given table
-    returns a list of (parent, child, timestamp)
-
-    Note: in `upstreams` or `downstreams` we only care about `tableInfo`,
-    we could also have `notebookInfos` or `fileInfo`
-    """
-    links: list[OTimestampedLink] = []
-    # add parent:
-    for link in single_table_lineage.get("upstreams", []):
-        parent = link.get("tableInfo", {})
-        parent_path = _to_table_path(parent)
-        timestamp: Ostr = parent.get("lineage_timestamp")
-        links.append(_link(parent_path, table_path, timestamp))
-
-    # add children:
-    for link in single_table_lineage.get("downstreams", []):
-        child = link.get("tableInfo", {})
-        child_path = _to_table_path(child)
-        timestamp = child.get("lineage_timestamp")
-        links.append(_link(table_path, child_path, timestamp))
-
-    return list(filter(None, links))
-
-
-def single_column_lineage_links(
-    column_path: str, single_column_lineage: dict
-) -> list[TimestampedLink]:
-    """
-    process databricks lineage API response for a given table
-    returns a list of (parent, child, timestamp)
-
-    Note: in `upstreams` or `downstreams` we only care about `tableInfo`,
-    we could also have `notebookInfos` or `fileInfo`
-    """
-    links: list[OTimestampedLink] = []
-    # add parent:
-    for link in single_column_lineage.get("upstream_cols", []):
-        parent_path = _to_column_path(link)
-        timestamp: Ostr = link.get("lineage_timestamp")
-        links.append(_link(parent_path, column_path, timestamp))
+        if not (parent and child and parent != child):
+            return
 
-    # add children:
-    for link in single_column_lineage.get("downstream_cols", []):
-        child_path = _to_column_path(link)
-        timestamp = link.get("lineage_timestamp")
-        links.append(_link(column_path, child_path, timestamp))
+        key = (parent, child)
+        if key in self.lineage and self.lineage[key]["event_time"] > timestamp:
+            return
 
-    return list(filter(None, links))
+        self.lineage[key] = link
 
 
-def paths_for_column_lineage(
-    tables: list[dict], columns: list[dict], table_lineage: list[dict]
-) -> list[tuple[str, str]]:
+def valid_lineage(
+    lineage: Iterable[dict], lineage_entity: LineageEntity
+) -> list[dict]:
     """
-    helper providing a list of candidate columns to look lineage for:
-    we only look for column lineage where there is table lineage
+    Filters out self-lineage or lineage with a missing source or target path,
+    then deduplicates by picking the link with the most recent event timestmap.
     """
-    # mapping between table id and its path db.schema.table
-    # table["schema_id"] follows the pattern `db.schema`
-    mapping = {
-        table["id"]: ".".join([table["schema_id"], table["table_name"]])
-        for table in tables
-    }
-
-    tables_with_lineage: set[str] = set()
-    for t in table_lineage:
-        tables_with_lineage.add(t["parent_path"])
-        tables_with_lineage.add(t["child_path"])
-
-    paths_to_return: list[tuple[str, str]] = []
-    for column in columns:
-        table_path = mapping[column["table_id"]]
-        if table_path not in tables_with_lineage:
-            continue
-        column_ = (table_path, column["column_name"])
-        paths_to_return.append(column_)
-
-    return paths_to_return
-
-
-def deduplicate_lineage(lineages: list[TimestampedLink]) -> dict:
-    deduplicated_lineage = LineageLinks()
-    for timestamped_link in lineages:
-        deduplicated_lineage.add(timestamped_link)
-    return deduplicated_lineage.lineage
+    deduplicated_lineage = LineageProcessor(lineage_entity)
+
+    for link in lineage:
+        deduplicated_lineage.add(link)
+
+    return list(deduplicated_lineage.lineage.values())