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

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())

castor_extractor/warehouse/databricks/lineage_test.py

@@ -1,34 +1,89 @@
-from .lineage import LineageLinks
-from .test_constants import (
-    CLOSER_DATE,
-    OLDER_DATE,
-)
+from .enums import LineageEntity
+from .lineage import LineageProcessor, valid_lineage
+
+_OLDER_DATE = "2025-01-01 00:00:01.0"
+_CLOSER_DATE = "2025-01-01 02:02:02.0"
+
+_TABLE_LINEAGES = [
+    {
+        "source_table_full_name": "a.b.source",
+        "target_table_full_name": "a.b.target",
+        "event_time": _CLOSER_DATE,
+        "other": "more recent stuff",
+    },
+    {
+        "source_table_full_name": "a.b.source",
+        "target_table_full_name": "a.b.target",
+        "event_time": _OLDER_DATE,
+        "other": "stuff that's too old",
+    },
+    {
+        "source_table_full_name": "no target",
+        "target_table_full_name": None,
+        "event_time": _CLOSER_DATE,
+    },
+    {
+        "source_table_full_name": None,
+        "target_table_full_name": "no source",
+        "event_time": _CLOSER_DATE,
+    },
+]
+
+
+_COLUMN_LINEAGES = [
+    {
+        "source_table_full_name": "a.b.source",
+        "source_column_name": "src_col",
+        "target_table_full_name": "a.b.target",
+        "target_column_name": "trgt_col",
+        "event_time": _OLDER_DATE,
+        "other": "old stuff",
+    },
+    {
+        "source_table_full_name": "a.b.source",
+        "source_column_name": "src_col",
+        "target_table_full_name": "a.b.target",
+        "target_column_name": "trgt_col",
+        "event_time": _CLOSER_DATE,
+        "other": "newer stuff",
+    },
+    {
+        "source_table_full_name": "a.b.toto",
+        "source_column_name": "toto_col",
+        "target_table_full_name": "a.b.tata",
+        "target_column_name": "tata_col",
+        "event_time": _OLDER_DATE,
+    },
+    {
+        "source_table_full_name": "a.b.source",
+        "source_column_name": "a.b.source",
+        "target_table_full_name": None,
+        "target_column_name": None,
+        "event_time": _CLOSER_DATE,
+    },
+]
+
+
+def test_valid_lineage():
+    table_links = valid_lineage(_TABLE_LINEAGES, LineageEntity.TABLE)
+
+    assert len(table_links) == 1
+    assert table_links[0]["source_table_full_name"] == "a.b.source"
+    assert table_links[0]["target_table_full_name"] == "a.b.target"
+    assert table_links[0]["event_time"] == _CLOSER_DATE
+    assert table_links[0]["other"] == "more recent stuff"
 
 
 def test_LineageLinks_add():
-    links = LineageLinks()
-    timestamped_link = ("parent", "child", None)
-    expected_key = ("parent", "child")
-
-    links.add(timestamped_link)
-
-    assert expected_key in links.lineage
-    assert links.lineage[expected_key] is None
-
-    # we replace None by an actual timestamp
-    timestamped_link = ("parent", "child", OLDER_DATE)
-    links.add(timestamped_link)
-    assert expected_key in links.lineage
-    assert links.lineage[expected_key] == OLDER_DATE
-
-    # we update with the more recent timestamp
-    timestamped_link = ("parent", "child", CLOSER_DATE)
-    links.add(timestamped_link)
-    assert expected_key in links.lineage
-    assert links.lineage[expected_key] == CLOSER_DATE
-
-    # we keep the more recent timestamp
-    timestamped_link = ("parent", "child", OLDER_DATE)
-    links.add(timestamped_link)
-    assert expected_key in links.lineage
-    assert links.lineage[expected_key] == CLOSER_DATE
+    deduplicated_lineage = LineageProcessor(LineageEntity.COLUMN)
+    for link in _COLUMN_LINEAGES:
+        deduplicated_lineage.add(link)
+
+    lineage = deduplicated_lineage.lineage
+    assert len(lineage) == 2
+    assert ("a.b.source.src_col", "a.b.target.trgt_col") in lineage
+    assert ("a.b.toto.toto_col", "a.b.tata.tata_col") in lineage
+    assert (
+        lineage[("a.b.source.src_col", "a.b.target.trgt_col")]["other"]
+        == "newer stuff"
+    )

castor_extractor/warehouse/databricks/sql_client.py

@@ -1,24 +1,24 @@
 import logging
 from collections import defaultdict
-from enum import Enum
+from datetime import date
 from typing import Optional
 
 from databricks import sql  # type: ignore
 
 from .credentials import DatabricksCredentials
+from .enums import LineageEntity, TagEntity
 from .format import TagMapping
+from .lineage import valid_lineage
 from .utils import build_path, tag_label
 
 logger = logging.getLogger(__name__)
 
 _INFORMATION_SCHEMA_SQL = "SELECT * FROM system.information_schema"
 
-
-class TagEntity(Enum):
-    """Entities that can be tagged in Databricks"""
-
-    COLUMN = "COLUMN"
-    TABLE = "TABLE"
+_LINEAGE_SQL_TPL = """
+SELECT * FROM system.access.{table_name}
+WHERE event_date = :day
+"""
 
 
 class DatabricksSQLClient:
@@ -71,7 +71,6 @@ class DatabricksSQLClient:
         https://docs.databricks.com/en/sql/language-manual/information-schema/column_tags.html
         """
         if not self._needs_extraction(entity):
-            # extracting tags require additional credentials (http_path)
             return dict()
 
         table = f"{entity.value.lower()}_tags"
@@ -88,3 +87,19 @@ class DatabricksSQLClient:
             mapping[path].append(label)
 
         return mapping
+
+    def get_lineage(
+        self, lineage_entity: LineageEntity, day: date
+    ) -> list[dict]:
+        """
+        Fetch {TABLE|COLUMN} lineage of the given day, via system tables
+        https://docs.databricks.com/en/admin/system-tables/lineage.html
+        """
+        table_name = f"{lineage_entity.value.lower()}_lineage"
+        query = _LINEAGE_SQL_TPL.format(table_name=table_name)
+        params = {"day": day}
+        result = self.execute_sql(query, params)
+        data = []
+        for row in result:
+            data.append(row.asDict())
+        return valid_lineage(data, lineage_entity)

castor_extractor/warehouse/databricks/types.py

@@ -1,8 +1 @@
-from typing import Optional
-
-Link = tuple[str, str]
 TablesColumns = tuple[list[dict], list[dict]]
-Ostr = Optional[str]
-TimestampedLink = tuple[str, str, Ostr]
-
-OTimestampedLink = Optional[TimestampedLink]

{castor_extractor-0.22.1.dist-info → castor_extractor-0.22.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.22.1
+Version: 0.22.6
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -38,7 +38,7 @@ 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 (>=24.16.0,<24.17.0) ; extra == "looker" or extra == "all"
+Requires-Dist: looker-sdk (>=25.0.0,<26.0.0) ; extra == "looker" or extra == "all"
 Requires-Dist: msal (>=1.20.0,<2.0.0) ; extra == "powerbi" or extra == "all"
 Requires-Dist: numpy (<2) ; extra == "bigquery" or extra == "databricks" or extra == "all"
 Requires-Dist: numpy (>=1.26) ; (python_version >= "3.12" and python_version < "3.13") and (extra == "bigquery" or extra == "databricks" or extra == "all")
@@ -57,7 +57,7 @@ Requires-Dist: snowflake-sqlalchemy (!=1.2.5,<2.0.0) ; extra == "snowflake" or e
 Requires-Dist: sqlalchemy (>=1.4,<1.5)
 Requires-Dist: sqlalchemy-bigquery[bqstorage] (>=1.0.0,<=2.0.0) ; extra == "bigquery" or extra == "all"
 Requires-Dist: sqlalchemy-redshift (>=0.8.14,<0.9.0) ; extra == "redshift" or extra == "all"
-Requires-Dist: tableauserverclient (==0.25.0) ; extra == "tableau" or extra == "all"
+Requires-Dist: tableauserverclient (>=0.25.0,<0.26.0) ; extra == "tableau" or extra == "all"
 Requires-Dist: tqdm (>=4.0.0,<5.0.0)
 Requires-Dist: typing-extensions (>=4,<5)
 Requires-Dist: websocket-client (>=1,<2) ; extra == "qlik" or extra == "all"
@@ -207,6 +207,27 @@ For any questions or bug report, contact us at [support@castordoc.com](mailto:su
 
 # 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