castor-extractor 0.16.11__tar.gz → 0.16.15__tar.gz

{castor_extractor-0.16.11 → castor_extractor-0.16.15}/CHANGELOG.md

@@ -1,5 +1,21 @@
 # Changelog
 
+## 0.16.15 - 2024-06-07
+
+* Tableau: extract database_name for CustomSQLTables
+
+## 0.16.14 - 2024-06-06
+
+* Snowflake: Extract SQL user defined function
+
+## 0.16.13 - 2024-06-05
+
+* Tableau: extract database_name for tables
+
+## 0.16.12 - 2024-06-04
+
+* Databricks: Extract lineage
+
 ## 0.16.11 - 2024-06-03
 
 * Tableau: add extra fields to optimise storage

{castor_extractor-0.16.11 → castor_extractor-0.16.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.16.11
+Version: 0.16.15
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA

{castor_extractor-0.16.11 → castor_extractor-0.16.15}/castor_extractor/utils/client/api.py

@@ -5,7 +5,7 @@ import requests
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_TIMEOUT_MS = 30_000
+DEFAULT_TIMEOUT_S = 30
 
 # https://requests.readthedocs.io/en/latest/api/#requests.request
 HttpMethod = Literal["GET", "OPTIONS", "HEAD", "POST", "PUT", "PATCH", "DELETE"]
@@ -20,7 +20,7 @@ class APIClient:
     def __init__(self, host: str, token: Optional[str] = None):
         self._host = host
         self._token = token or ""
-        self._timeout = DEFAULT_TIMEOUT_MS
+        self._timeout = DEFAULT_TIMEOUT_S
 
     @staticmethod
     def build_url(host: str, path: str):
@@ -44,7 +44,12 @@ class APIClient:
     ) -> Any:
         logger.debug(f"Calling {method} on {url}")
         result = requests.request(
-            method, url, headers=self._headers(), params=params, json=data
+            method,
+            url,
+            headers=self._headers(),
+            params=params,
+            json=data,
+            timeout=self._timeout,
         )
         result.raise_for_status()
 

{castor_extractor-0.16.11 → castor_extractor-0.16.15}/castor_extractor/utils/retry.py

@@ -68,7 +68,8 @@ class Retry(BaseModel):
         self._retry_attempts += 1
         wait_ms = self.base() + self.jitter()
         wait_s = float(wait_ms) / MS_IN_SEC
-        logger.warning(f"Attempting a new call in {wait_s} seconds")
+        msg = f"Attempting a new call in {wait_s} seconds, {self._retry_attempts} attempt(s) / {self.max_retries} max retries"
+        logger.warning(msg)
         time.sleep(wait_s)
         return True
 
@@ -93,6 +94,7 @@ def retry(
             try:
                 return None, callable(*args, **kwargs)
             except exceptions_ as err:
+                logger.warning(f"Exception within {callable.__name__}")
                 return err, None
 
         def _func(*args, **kwargs) -> Any:

{castor_extractor-0.16.11 → castor_extractor-0.16.15}/castor_extractor/visualization/tableau_revamp/client/client.py

@@ -31,8 +31,11 @@ _TSC_ASSETS = (
 # 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,
+    # for some clients, extraction of columns tend to hit the node limit
+    # https://community.tableau.com/s/question/0D54T00000YuK60SAF/metadata-query-nodelimitexceeded-error
+    # the workaround is to reduce pagination
+    TableauRevampAsset.COLUMN: 50,
+    # fields are light but volumes are bigger
     TableauRevampAsset.FIELD: 1000,
     TableauRevampAsset.TABLE: 50,
 }

{castor_extractor-0.16.11 → castor_extractor-0.16.15}/castor_extractor/visualization/tableau_revamp/client/gql_queries.py

@@ -63,12 +63,21 @@ downstreamWorkbooks { id }
 id
 name
 ... on DatabaseTable {
-    connectionType
     fullName
     schema
+    database {
+        connectionType
+        id
+        name
+    }
 }
 ... on CustomSQLTable {
     query
+    database {
+        connectionType
+        id
+        name
+    }
 }
 """
 

{castor_extractor-0.16.11 → castor_extractor-0.16.15}/castor_extractor/warehouse/abstract/__init__.py

@@ -1,6 +1,8 @@
 from .asset import (
+    ADDITIONAL_LINEAGE_ASSETS,
     CATALOG_ASSETS,
     EXTERNAL_LINEAGE_ASSETS,
+    FUNCTIONS_ASSETS,
     QUERIES_ASSETS,
     VIEWS_ASSETS,
     SupportedAssets,

{castor_extractor-0.16.11 → castor_extractor-0.16.15}/castor_extractor/warehouse/abstract/asset.py

@@ -7,6 +7,8 @@ from ...types import ExternalAsset, classproperty
 class WarehouseAsset(ExternalAsset):
     """Assets that can be extracted from warehouses"""
 
+    ADDITIONAL_COLUMN_LINEAGE = "additional_column_lineage"
+    ADDITIONAL_TABLE_LINEAGE = "additional_table_lineage"
     COLUMN = "column"
     COLUMN_LINEAGE = "column_lineage"  # specific to snowflake
     DATABASE = "database"
@@ -19,12 +21,15 @@ class WarehouseAsset(ExternalAsset):
     ROLE = "role"
     SCHEMA = "schema"
     TABLE = "table"
+    FUNCTION = "function"
     USER = "user"
     VIEW_DDL = "view_ddl"
 
     @classproperty
     def optional(cls) -> Set["WarehouseAsset"]:
         return {
+            WarehouseAsset.ADDITIONAL_COLUMN_LINEAGE,
+            WarehouseAsset.ADDITIONAL_TABLE_LINEAGE,
             WarehouseAsset.EXTERNAL_COLUMN_LINEAGE,
             WarehouseAsset.EXTERNAL_TABLE_LINEAGE,
         }
@@ -33,8 +38,10 @@ class WarehouseAsset(ExternalAsset):
 class WarehouseAssetGroup(Enum):
     """Groups of assets that can be extracted together"""
 
+    ADDITIONAL_LINEAGE = "additional_lineage"
     CATALOG = "catalog"
     EXTERNAL_LINEAGE = "external_lineage"
+    FUNCTION = "function"
     QUERY = "query"
     ROLE = "role"
     SNOWFLAKE_LINEAGE = "snowflake_lineage"
@@ -53,6 +60,7 @@ CATALOG_ASSETS = (
 )
 
 # shared by technologies supporting queries
+FUNCTIONS_ASSETS = (WarehouseAsset.FUNCTION,)
 QUERIES_ASSETS = (WarehouseAsset.QUERY,)
 VIEWS_ASSETS = (WarehouseAsset.VIEW_DDL,)
 
@@ -61,6 +69,11 @@ EXTERNAL_LINEAGE_ASSETS = (
     WarehouseAsset.EXTERNAL_TABLE_LINEAGE,
 )
 
+ADDITIONAL_LINEAGE_ASSETS = (
+    WarehouseAsset.ADDITIONAL_COLUMN_LINEAGE,
+    WarehouseAsset.ADDITIONAL_TABLE_LINEAGE,
+)
+
 NON_EXTRACTABLE_ASSETS = {WarehouseAssetGroup.EXTERNAL_LINEAGE}
 
 

castor_extractor-0.16.15/castor_extractor/warehouse/databricks/client.py

@@ -0,0 +1,465 @@
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from datetime import date
+from functools import partial
+from typing import Any, Dict, List, Optional, Set, Tuple, cast
+
+import requests
+
+from ...utils import (
+    SafeMode,
+    at_midnight,
+    date_after,
+    mapping_from_rows,
+    retry,
+    safe_mode,
+)
+from ...utils.client.api import APIClient
+from ...utils.pager import PagerOnToken
+from ..abstract.time_filter import TimeFilter
+from .credentials import DatabricksCredentials
+from .format import DatabricksFormatter
+from .types import Link, Ostr, OTimestampedLink, TablesColumns, TimestampedLink
+
+logger = logging.getLogger(__name__)
+
+_MAX_NUMBER_OF_LINEAGE_ERRORS = 1000
+_MAX_THREADS = 10
+_RETRY_ATTEMPTS = 3
+_RETRY_BASE_MS = 1000
+_RETRY_EXCEPTIONS = [
+    requests.exceptions.ConnectTimeout,
+]
+
+safe_params = SafeMode((BaseException,), _MAX_NUMBER_OF_LINEAGE_ERRORS)
+
+
+def _day_to_epoch_ms(day: date) -> int:
+    return int(at_midnight(day).timestamp() * 1000)
+
+
+def _day_hour_to_epoch_ms(day: date, hour: int) -> int:
+    return int(at_midnight(day).timestamp() * 1000) + (hour * 3600 * 1000)
+
+
+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
+        else:
+            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
+
+
+class DatabricksClient(APIClient):
+    """Databricks Client"""
+
+    def __init__(
+        self,
+        credentials: DatabricksCredentials,
+        db_allowed: Optional[Set[str]] = None,
+        db_blocked: Optional[Set[str]] = None,
+    ):
+        super().__init__(host=credentials.host, token=credentials.token)
+        self._db_allowed = db_allowed
+        self._db_blocked = db_blocked
+        self.formatter = DatabricksFormatter()
+
+    @staticmethod
+    def name() -> str:
+        return "Databricks"
+
+    def _keep_catalog(self, catalog: str) -> bool:
+        """
+        Helper function to determine if we should keep the Databricks catalog
+        which is a CastorDoc database
+        """
+        if self._db_allowed and catalog not in self._db_allowed:
+            return False
+        if self._db_blocked and catalog in self._db_blocked:
+            return False
+        return True
+
+    def databases(self) -> List[dict]:
+        path = "api/2.1/unity-catalog/catalogs"
+        content = self.get(path=path)
+        _databases = self.formatter.format_database(content.get("catalogs", []))
+        return [d for d in _databases if self._keep_catalog(d["database_name"])]
+
+    def _schemas_of_database(self, database: dict) -> List[dict]:
+        path = "api/2.1/unity-catalog/schemas"
+        payload = {"catalog_name": database["database_name"]}
+        content = self.get(path=path, payload=payload)
+        schemas = content.get("schemas", [])
+        return self.formatter.format_schema(schemas, database)
+
+    def schemas(self, databases: List[dict]) -> List[dict]:
+        """
+        Get the databricks schemas (also sometimes called databases)
+        (which correspond to the schemas in Castor)
+        leveraging the unity catalog API
+        """
+        return [
+            schema
+            for database in databases
+            for schema in self._schemas_of_database(database)
+        ]
+
+    def _tables_columns_of_schema(self, schema: dict) -> TablesColumns:
+        path = "api/2.1/unity-catalog/tables"
+        payload = {
+            "catalog_name": schema["database_id"],
+            "schema_name": schema["schema_name"],
+        }
+        content = self.get(path=path, payload=payload)
+        return self.formatter.format_table_column(
+            content.get("tables", []), schema
+        )
+
+    @staticmethod
+    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_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:
+        """
+        Get the databricks tables & columns leveraging the unity catalog API
+        """
+        tables: List[dict] = []
+        columns: List[dict] = []
+        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_mapping)
+                for table in t_to_add
+            ]
+            tables.extend(t_with_owner)
+            columns.extend(c_to_add)
+        return tables, columns
+
+    @staticmethod
+    def _to_table_path(table: dict) -> Ostr:
+        if table.get("name"):
+            return f"{table['catalog_name']}.{table['schema_name']}.{table['name']}"
+        return None
+
+    @staticmethod
+    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(
+        self, 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(
+        self, 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 = self._to_table_path(parent)
+            timestamp: Ostr = parent.get("lineage_timestamp")
+            links.append(self._link(parent_path, table_path, timestamp))
+
+        # add children:
+        for link in single_table_lineage.get("downstreams", []):
+            child = link.get("tableInfo", {})
+            child_path = self._to_table_path(child)
+            timestamp = child.get("lineage_timestamp")
+            links.append(self._link(table_path, child_path, timestamp))
+
+        return list(filter(None, links))
+
+    @safe_mode(safe_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
+        """
+        path = "api/2.0/lineage-tracking/table-lineage"
+        payload = {"table_name": table_path, "include_entity_lineage": True}
+        content = self.get(path=path, payload=payload)
+        return self._single_table_lineage_links(table_path, content)
+
+    def _deduplicate_lineage(self, lineages: List[TimestampedLink]) -> dict:
+        deduplicated_lineage = LineageLinks()
+        for timestamped_link in lineages:
+            deduplicated_lineage.add(timestamped_link)
+        return deduplicated_lineage.lineage
+
+    def table_lineage(self, tables: List[dict]) -> List[dict]:
+        """
+        Wrapper function that retrieves all table lineage
+        """
+        # retrieve table lineage
+        with ThreadPoolExecutor(max_workers=_MAX_THREADS) as executor:
+            table_paths = [
+                ".".join([table["schema_id"], table["table_name"]])
+                for table in tables
+            ]
+            results = executor.map(self.get_single_table_lineage, table_paths)
+        lineages = [link for links in results for link in links]
+        deduplicated = self._deduplicate_lineage(lineages)
+        return self.formatter.format_lineage(deduplicated)
+
+    @staticmethod
+    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 _single_column_lineage_links(
+        self, 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 = self._to_column_path(link)
+            timestamp: Ostr = link.get("lineage_timestamp")
+            links.append(self._link(parent_path, column_path, timestamp))
+
+        # add children:
+        for link in single_column_lineage.get("downstream_cols", []):
+            child_path = self._to_column_path(link)
+            timestamp = link.get("lineage_timestamp")
+            links.append(self._link(column_path, child_path, timestamp))
+
+        return list(filter(None, links))
+
+    @safe_mode(safe_params, lambda: [])
+    @retry(
+        exceptions=_RETRY_EXCEPTIONS,
+        max_retries=_RETRY_ATTEMPTS,
+        base_ms=_RETRY_BASE_MS,
+    )
+    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
+        api_path = "api/2.0/lineage-tracking/column-lineage"
+        payload = {
+            "table_name": table_path,
+            "column_name": column_name,
+            "include_entity_lineage": True,
+        }
+        content = self.get(path=api_path, payload=payload)
+        column_path = f"{table_path}.{column_name}"
+        return self._single_column_lineage_links(column_path, content)
+
+    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 = self._paths_for_column_lineage(
+            tables, columns, table_lineage
+        )
+        lineages: List[TimestampedLink] = [
+            link
+            for paths in candidate_paths
+            for link in self.get_single_column_lineage(paths)
+        ]
+        deduplicated = self._deduplicate_lineage(lineages)
+        return self.formatter.format_lineage(deduplicated)
+
+    @staticmethod
+    def _time_filter(time_filter: Optional[TimeFilter]) -> dict:
+        """time filter to retrieve Databricks' queries"""
+        # 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 not None and hour_max is not None:  # specific window
+            start_time_ms = _day_hour_to_epoch_ms(day, hour_min)
+            # note: in practice, hour_min == hour_max (hourly query ingestion)
+            end_time_ms = _day_hour_to_epoch_ms(day, hour_max + 1)
+        else:  # fallback to an extraction of the entire day
+            start_time_ms = _day_to_epoch_ms(day)
+            end_time_ms = _day_to_epoch_ms(date_after(day, 1))
+
+        return {
+            "filter_by": {
+                "query_start_time_range": {
+                    "end_time_ms": end_time_ms,
+                    "start_time_ms": start_time_ms,
+                }
+            }
+        }
+
+    def query_payload(
+        self,
+        page_token: Optional[str] = None,
+        max_results: Optional[int] = None,
+        time_range_filter: Optional[dict] = None,
+    ) -> dict:
+        """helper method to build the payload used to retrieve queries"""
+        # in payload: You can provide only one of 'page_token' or 'filter_by'
+        if page_token:
+            payload: Dict[str, Any] = {"page_token": page_token}
+        else:
+            if time_range_filter:
+                payload = {**time_range_filter}
+            else:
+                payload = self._time_filter(None)  # default to yesterday
+        if max_results:
+            payload["max_results"] = max_results
+        return payload
+
+    def _scroll_queries(
+        self,
+        page_token: Optional[str] = None,
+        max_results: Optional[int] = None,
+        time_range_filter: Optional[dict] = None,
+    ) -> dict:
+        """
+        Callback to scroll the queries api
+        https://docs.databricks.com/api/workspace/queryhistory/list
+        max_results: Limit the number of results returned in one page.
+            The default is 100. (both on our side and Databricks')
+        """
+        path = "api/2.0/sql/history/queries"
+        payload = self.query_payload(page_token, max_results, time_range_filter)
+        content = self.get(path=path, payload=payload)
+        return content if content else {}
+
+    def queries(self, time_filter: Optional[TimeFilter] = None) -> List[dict]:
+        """get all queries"""
+        # add a time filter (by default: yesterday)
+        time_range_filter = self._time_filter(time_filter)
+        _time_filtered_scroll_queries = partial(
+            self._scroll_queries,
+            time_range_filter=time_range_filter,
+        )
+
+        # retrieve all queries using pagination
+        raw_queries = PagerOnToken(_time_filtered_scroll_queries).all()
+
+        return self.formatter.format_query(raw_queries)
+
+    def users(self) -> List[dict]:
+        """
+        retrieve user from api
+        """
+        path = "api/2.0/preview/scim/v2/Users"
+        content = self.get(path=path)
+        return self.formatter.format_user(content.get("Resources", []))
+
+    def _view_ddl(self, schema: dict) -> List[dict]:
+        path = "api/2.1/unity-catalog/tables"
+        payload = {
+            "catalog_name": schema["database_id"],
+            "schema_name": schema["schema_name"],
+            "omit_columns": True,
+        }
+        content = self.get(path=path, payload=payload)
+        return self.formatter.format_view_ddl(content.get("tables", []), schema)
+
+    def view_ddl(self, schemas: List[dict]) -> List[dict]:
+        """retrieve view ddl"""
+        view_ddl: List[dict] = []
+        for schema in schemas:
+            v_to_add = self._view_ddl(schema)
+            view_ddl.extend(v_to_add)
+        return view_ddl