castor-extractor 0.19.0__py3-none-any.whl → 0.19.6__py3-none-any.whl

castor_extractor/warehouse/databricks/client_test.py

@@ -1,14 +1,7 @@
-from datetime import date
 from unittest.mock import Mock, patch
 
-from freezegun import freeze_time
-
-from ..abstract.time_filter import TimeFilter
 from .client import (
     DatabricksClient,
-    DatabricksCredentials,
-    LineageLinks,
-    _day_hour_to_epoch_ms,
 )
 from .test_constants import (
     CLOSER_DATE,
@@ -18,74 +11,12 @@ from .test_constants import (
 )
 
 
-def test__day_hour_to_epoch_ms():
-    _day_hour_to_epoch_ms(date(2023, 2, 14), 14) == 1644847200000
-
-
-@freeze_time("2023-7-4")
-def test_DatabricksClient__hourly_time_filters():
-    credentials = DatabricksCredentials(
-        host="carthago",
-        token="delenda",
-        http_host="est",
-    )
-    client = DatabricksClient(credentials)
-
-    # default is yesterday
-    default_filters = [f for f in client._hourly_time_filters(None)]
-
-    assert len(default_filters) == 24  # number of hours in a day
-
-    first = default_filters[0]
-    start = first["filter_by"]["query_start_time_range"]["start_time_ms"]
-    last = default_filters[-1]
-    end = last["filter_by"]["query_start_time_range"]["end_time_ms"]
-    assert start == 1688342400000  # July 3, 2023 12:00:00 AM GMT
-    assert end == 1688428800000  # July 4, 2023 12:00:00 AM GMT
-
-    # custom time (from execution_date in DAG for example)
-    time_filter = TimeFilter(day=date(2020, 10, 15))
-    custom_filters = [f for f in client._hourly_time_filters(time_filter)]
-
-    assert len(custom_filters) == 24
-
-    first = custom_filters[0]
-    start = first["filter_by"]["query_start_time_range"]["start_time_ms"]
-    last = custom_filters[-1]
-    end = last["filter_by"]["query_start_time_range"]["end_time_ms"]
-    assert start == 1602720000000  # Oct 15, 2020 12:00:00 AM
-    assert end == 1602806400000  # Oct 16, 2020 12:00:00 AM
-
-    # hourly extraction: note that hour_min == hour_max
-    hourly = TimeFilter(day=date(2023, 4, 14), hour_min=4, hour_max=4)
-    hourly_filters = [f for f in client._hourly_time_filters(hourly)]
-    expected_hourly = [
-        {
-            "filter_by": {
-                "query_start_time_range": {
-                    "end_time_ms": 1681448400000,  # April 14, 2023 5:00:00 AM
-                    "start_time_ms": 1681444800000,  # April 14, 2023 4:00:00 AM
-                }
-            }
-        }
-    ]
-    assert hourly_filters == expected_hourly
-
-
 class MockDatabricksClient(DatabricksClient):
     def __init__(self):
         self._db_allowed = ["prd", "staging"]
         self._db_blocked = ["dev"]
 
 
-def test_DatabricksClient__keep_catalog():
-    client = MockDatabricksClient()
-    assert client._keep_catalog("prd")
-    assert client._keep_catalog("staging")
-    assert not client._keep_catalog("dev")
-    assert not client._keep_catalog("something_unknown")
-
-
 def test_DatabricksClient__get_user_mapping():
     client = MockDatabricksClient()
     users = [
@@ -120,7 +51,7 @@ def test_DatabricksClient__match_table_with_user():
 
 
 @patch(
-    "source.packages.extractor.castor_extractor.warehouse.databricks.client.DatabricksClient.get",
+    "source.packages.extractor.castor_extractor.warehouse.databricks.client.DatabricksAPIClient._get",
     side_effect=TABLE_LINEAGE_SIDE_EFFECT,
 )
 def test_DatabricksClient_table_lineage(mock_get):
@@ -141,32 +72,3 @@ def test_DatabricksClient_table_lineage(mock_get):
     }
     assert expected_link_1 in lineage
     assert expected_link_2 in lineage
-
-
-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

castor_extractor/warehouse/databricks/endpoints.py

@@ -0,0 +1,28 @@
+class DatabricksEndpointFactory:
+    @classmethod
+    def tables(cls):
+        return "api/2.1/unity-catalog/tables"
+
+    @classmethod
+    def schemas(cls):
+        return "api/2.1/unity-catalog/schemas"
+
+    @classmethod
+    def databases(cls):
+        return "api/2.1/unity-catalog/catalogs"
+
+    @classmethod
+    def table_lineage(cls):
+        return "api/2.0/lineage-tracking/table-lineage"
+
+    @classmethod
+    def column_lineage(cls):
+        return "api/2.0/lineage-tracking/column-lineage"
+
+    @classmethod
+    def queries(cls):
+        return "api/2.0/sql/history/queries"
+
+    @classmethod
+    def users(cls):
+        return "api/2.0/preview/scim/v2/Users"

castor_extractor/warehouse/databricks/lineage.py

@@ -0,0 +1,141 @@
+from typing import Dict, List, Set, Tuple, cast
+
+from .types import Link, Ostr, OTimestampedLink, TimestampedLink
+
+
+class LineageLinks:
+    """
+    helper class that handles lineage deduplication and filtering
+    """
+
+    def __init__(self):
+        self.lineage: Dict[Link, Ostr] = dict()
+
+    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
+
+
+def _to_table_path(table: dict) -> Ostr:
+    if table.get("name"):
+        return f"{table['catalog_name']}.{table['schema_name']}.{table['name']}"
+    return None
+
+
+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 _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
+
+
+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))
+
+    # 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))
+
+    return list(filter(None, links))
+
+
+def paths_for_column_lineage(
+    tables: List[dict], columns: List[dict], table_lineage: List[dict]
+) -> List[Tuple[str, str]]:
+    """
+    helper providing a list of candidate columns to look lineage for:
+    we only look for column lineage where there is table lineage
+    """
+    # 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

castor_extractor/warehouse/databricks/lineage_test.py

@@ -0,0 +1,34 @@
+from .lineage import LineageLinks
+from .test_constants import (
+    CLOSER_DATE,
+    OLDER_DATE,
+)
+
+
+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

castor_extractor/warehouse/databricks/pagination.py

@@ -0,0 +1,22 @@
+from typing import List, Optional
+
+from pydantic import Field
+
+from ...utils import PaginationModel
+
+DATABRICKS_PAGE_SIZE = 100
+
+
+class DatabricksPagination(PaginationModel):
+    next_page_token: Optional[str] = None
+    has_next_page: bool = False
+    res: List[dict] = Field(default_factory=list)
+
+    def is_last(self) -> bool:
+        return not (self.has_next_page and self.next_page_token)
+
+    def next_page_payload(self) -> dict:
+        return {"page_token": self.next_page_token}
+
+    def page_results(self) -> list:
+        return self.res

castor_extractor/warehouse/databricks/sql_client.py

@@ -0,0 +1,90 @@
+import logging
+from collections import defaultdict
+from enum import Enum
+from typing import Optional
+
+from databricks import sql  # type: ignore
+
+from .credentials import DatabricksCredentials
+from .format import TagMapping
+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"
+
+
+class DatabricksSQLClient:
+    def __init__(
+        self,
+        credentials: DatabricksCredentials,
+        has_table_tags: bool = False,
+        has_column_tags: bool = False,
+    ):
+        self._http_path = credentials.http_path
+        self._has_table_tags = has_table_tags
+        self._has_column_tags = has_column_tags
+        self._host = credentials.host
+        self._token = credentials.token
+
+    def execute_sql(
+        self,
+        query: str,
+        params: Optional[dict] = None,
+    ):
+        """
+        Execute a SQL query on Databricks system tables and return the results.
+        https://docs.databricks.com/en/dev-tools/python-sql-connector.html
+
+        //!\\ credentials.http_path is required in order to run SQL queries
+        """
+        assert self._http_path, "HTTP_PATH is required to run SQL queries"
+        with sql.connect(
+            server_hostname=self._host,
+            http_path=self._http_path,
+            access_token=self._token,
+        ) as connection:
+            with connection.cursor() as cursor:
+                cursor.execute(query, params)
+                return cursor.fetchall()
+
+    def _needs_extraction(self, entity: TagEntity) -> bool:
+        if entity == TagEntity.TABLE:
+            return self._has_table_tags
+        if entity == TagEntity.COLUMN:
+            return self._has_column_tags
+        raise AssertionError(f"Entity not supported: {entity}")
+
+    def get_tags_mapping(self, entity: TagEntity) -> TagMapping:
+        """
+        Fetch tags of the given entity and build a mapping:
+        { path: list[tags] }
+
+        https://docs.databricks.com/en/sql/language-manual/information-schema/table_tags.html
+        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"
+        query = f"{_INFORMATION_SCHEMA_SQL}.{table}"
+        result = self.execute_sql(query)
+        mapping = defaultdict(list)
+        for row in result:
+            dict_row = row.asDict()
+            keys = ["catalog_name", "schema_name", "table_name"]
+            if entity == TagEntity.COLUMN:
+                keys.append("column_name")
+            path = build_path(dict_row, keys)
+            label = tag_label(dict_row)
+            mapping[path].append(label)
+
+        return mapping

castor_extractor/warehouse/databricks/utils.py

@@ -1,4 +1,16 @@
-from typing import Dict, List
+from datetime import date
+from typing import Dict, Iterable, List, Optional
+
+from ...utils import at_midnight
+from ..abstract import TimeFilter
+
+_DEFAULT_HOUR_MIN = 0
+_DEFAULT_HOUR_MAX = 23
+_NUM_HOURS_IN_A_DAY = 24
+
+
+def _day_hour_to_epoch_ms(day: date, hour: int) -> int:
+    return int(at_midnight(day).timestamp() * 1000) + (hour * 3600 * 1000)
 
 
 def build_path(
@@ -25,3 +37,34 @@ def tag_label(row: Dict) -> str:
     if not tag_value:
         return tag_name
     return f"{tag_name}:{tag_value}"
+
+
+def _time_filter_payload(start_time_ms: int, end_time_ms: int) -> dict:
+    return {
+        "filter_by": {
+            "query_start_time_range": {
+                "end_time_ms": end_time_ms,
+                "start_time_ms": start_time_ms,
+            }
+        }
+    }
+
+
+def hourly_time_filters(time_filter: Optional[TimeFilter]) -> Iterable[dict]:
+    """time filters to retrieve Databricks' queries: 1h duration each"""
+    # define an explicit time window
+    if not time_filter:
+        time_filter = TimeFilter.default()
+
+    assert time_filter  # for mypy
+
+    hour_min = time_filter.hour_min
+    hour_max = time_filter.hour_max
+    day = time_filter.day
+    if hour_min is None or hour_max is None:  # fallback to an entire day
+        hour_min, hour_max = _DEFAULT_HOUR_MIN, _DEFAULT_HOUR_MAX
+
+    for index in range(hour_min, min(hour_max + 1, _NUM_HOURS_IN_A_DAY)):
+        start_time_ms = _day_hour_to_epoch_ms(day, index)
+        end_time_ms = _day_hour_to_epoch_ms(day, index + 1)
+        yield _time_filter_payload(start_time_ms, end_time_ms)

castor_extractor/warehouse/databricks/utils_test.py

@@ -1,4 +1,14 @@
-from .utils import build_path, tag_label
+from datetime import date
+
+from freezegun import freeze_time
+
+from ..abstract import TimeFilter
+from .utils import (
+    _day_hour_to_epoch_ms,
+    build_path,
+    hourly_time_filters,
+    tag_label,
+)
 
 
 def test_build_path():
@@ -23,3 +33,50 @@ def test_tag_label():
         "tag_value": "fou",
     }
     assert tag_label(row) == "fi:fou"
+
+
+def test__day_hour_to_epoch_ms():
+    assert _day_hour_to_epoch_ms(date(2023, 2, 14), 14) == 1676383200000
+
+
+@freeze_time("2023-7-4")
+def test_hourly_time_filters():
+    # default is yesterday
+    default_filters = [f for f in hourly_time_filters(None)]
+
+    assert len(default_filters) == 24  # number of hours in a day
+
+    first = default_filters[0]
+    start = first["filter_by"]["query_start_time_range"]["start_time_ms"]
+    last = default_filters[-1]
+    end = last["filter_by"]["query_start_time_range"]["end_time_ms"]
+    assert start == 1688342400000  # July 3, 2023 12:00:00 AM GMT
+    assert end == 1688428800000  # July 4, 2023 12:00:00 AM GMT
+
+    # custom time (from execution_date in DAG for example)
+    time_filter = TimeFilter(day=date(2020, 10, 15))
+    custom_filters = [f for f in hourly_time_filters(time_filter)]
+
+    assert len(custom_filters) == 24
+
+    first = custom_filters[0]
+    start = first["filter_by"]["query_start_time_range"]["start_time_ms"]
+    last = custom_filters[-1]
+    end = last["filter_by"]["query_start_time_range"]["end_time_ms"]
+    assert start == 1602720000000  # Oct 15, 2020 12:00:00 AM
+    assert end == 1602806400000  # Oct 16, 2020 12:00:00 AM
+
+    # hourly extraction: note that hour_min == hour_max
+    hourly = TimeFilter(day=date(2023, 4, 14), hour_min=4, hour_max=4)
+    hourly_filters = [f for f in hourly_time_filters(hourly)]
+    expected_hourly = [
+        {
+            "filter_by": {
+                "query_start_time_range": {
+                    "end_time_ms": 1681448400000,  # April 14, 2023 5:00:00 AM
+                    "start_time_ms": 1681444800000,  # April 14, 2023 4:00:00 AM
+                }
+            }
+        }
+    ]
+    assert hourly_filters == expected_hourly

castor_extractor/warehouse/mysql/client.py

@@ -1,5 +1,3 @@
-import pymysql  # type: ignore
-
 from ...utils import SqlalchemyClient, uri_encode
 
 SERVER_URI = "{user}:{password}@{host}:{port}"

castor_extractor/warehouse/salesforce/client.py

@@ -1,14 +1,16 @@
 import logging
-from typing import Dict, Iterator, List, Optional, Tuple
+from functools import partial
+from typing import Dict, List, Optional, Tuple
 
 from tqdm import tqdm  # type: ignore
 
+from ...utils import fetch_all_pages
 from ...utils.salesforce import SalesforceBaseClient, SalesforceCredentials
 from .format import SalesforceFormatter
+from .pagination import SalesforceSQLPagination, format_sobject_query
 from .soql import (
     DESCRIPTION_QUERY_TPL,
     SOBJECT_FIELDS_QUERY_TPL,
-    SOBJECTS_QUERY_TPL,
 )
 
 logger = logging.getLogger(__name__)
@@ -19,9 +21,6 @@ class SalesforceClient(SalesforceBaseClient):
     Salesforce API client to extract sobjects
     """
 
-    # Implicit (hard-coded in Salesforce) limitation when using SOQL of 2,000 rows
-    LIMIT_RECORDS_PER_PAGE = 2000
-
     def __init__(self, credentials: SalesforceCredentials):
         super().__init__(credentials)
         self.formatter = SalesforceFormatter()
@@ -30,74 +29,28 @@ class SalesforceClient(SalesforceBaseClient):
     def name() -> str:
         return "Salesforce"
 
-    def _format_query(self, query_template: str, start_durable_id: str) -> str:
-        return query_template.format(
-            start_durable_id=start_durable_id,
-            limit=self.LIMIT_RECORDS_PER_PAGE,
-        )
-
-    def _next_records(
-        self, url: str, query_template: str, start_durable_id: str = "0000"
-    ) -> List[dict]:
-        query = self._format_query(
-            query_template, start_durable_id=start_durable_id
-        )
-        records, _ = self._call(
-            url, params={"q": query}, processor=self._query_processor
-        )
-        return records
-
-    def _is_last_page(self, records: List[dict]) -> bool:
-        return len(records) < self.LIMIT_RECORDS_PER_PAGE
-
-    def _should_query_next_page(
-        self, records: List[dict], page_number: int
-    ) -> bool:
-        return not (
-            self._is_last_page(records)
-            or self._has_reached_pagination_limit(page_number)
-        )
-
-    def _query_all(self, query_template: str) -> Iterator[dict]:
-        """
-        Run a SOQL query over salesforce API
-
-        Note, pagination is performed via a LIMIT in the SOQL query and requires
-        that ids are sorted. The SOQL query must support `limit` and
-        `start_durable_id` as parameters.
-        """
-        url = self.query_url
-        logger.info("querying page 0")
-        records = self._next_records(url, query_template)
-        yield from records
-
-        page_count = 1
-        while self._should_query_next_page(records, page_count):
-            logger.info(f"querying page {page_count}")
-            last_durable_id = records[-1]["DurableId"]
-            records = self._next_records(
-                url, query_template, start_durable_id=last_durable_id
-            )
-            yield from records
-            page_count += 1
-
     def fetch_sobjects(self) -> List[dict]:
         """Fetch all sobjects"""
         logger.info("Extracting sobjects")
-        return list(self._query_all(SOBJECTS_QUERY_TPL))
+        query = format_sobject_query()
+        request_ = partial(
+            self._get, endpoint=self.query_endpoint, params={"q": query}
+        )
+        results = fetch_all_pages(request_, SalesforceSQLPagination)
+        return list(results)
 
     def fetch_fields(self, sobject_name: str) -> List[dict]:
         """Fetches fields of a given sobject"""
         query = SOBJECT_FIELDS_QUERY_TPL.format(
             entity_definition_id=sobject_name
         )
-        response = self._call(self.tooling_url, params={"q": query})
+        response = self._get(self.tooling_endpoint, params={"q": query})
         return response["records"]
 
     def fetch_description(self, table_name: str) -> Optional[str]:
         """Retrieve description of a table"""
         query = DESCRIPTION_QUERY_TPL.format(table_name=table_name)
-        response = self._call(self.tooling_url, params={"q": query})
+        response = self._get(self.tooling_endpoint, params={"q": query})
         if not response["records"]:
             return None
         return response["records"][0]["Description"]