castor-extractor 0.24.7__tar.gz → 0.24.9__tar.gz

{castor_extractor-0.24.7 → castor_extractor-0.24.9}/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## 0.24.9 - 2025-04-16
+
+* Introduce API client for **Coalesce**
+
+## 0.24.8 - 2025-04-16
+
+* Tableau - remove duplicates introduced by `offset` pagination
+
 ## 0.24.7 - 2025-04-07
 
 * Tableau - switch from `cursor` to `offset` pagination to mitigate timeout issues

{castor_extractor-0.24.7 → castor_extractor-0.24.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.24.7
+Version: 0.24.9
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -210,6 +210,14 @@ For any questions or bug report, contact us at [support@castordoc.com](mailto:su
 
 # Changelog
 
+## 0.24.9 - 2025-04-16
+
+* Introduce API client for **Coalesce**
+
+## 0.24.8 - 2025-04-16
+
+* Tableau - remove duplicates introduced by `offset` pagination
+
 ## 0.24.7 - 2025-04-07
 
 * Tableau - switch from `cursor` to `offset` pagination to mitigate timeout issues

castor_extractor-0.24.9/castor_extractor/transformation/coalesce/__init__.py

@@ -0,0 +1,2 @@
+from .assets import CoalesceAsset, CoalesceQualityAsset
+from .client import CoalesceClient, CoalesceCredentials

castor_extractor-0.24.9/castor_extractor/transformation/coalesce/assets.py

@@ -0,0 +1,18 @@
+from ...types import ExternalAsset
+
+
+class CoalesceAsset(ExternalAsset):
+    """Coalesce assets"""
+
+    NODES = "nodes"
+
+
+class CoalesceQualityAsset(ExternalAsset):
+    """
+    Coalesce Quality Assets
+    Remark: having a dedicated Enum for Quality simplifies the process of
+    searching pushed files
+    """
+
+    NODES = "nodes"
+    RUN_RESULTS = "run_results"

castor_extractor-0.24.9/castor_extractor/transformation/coalesce/client/__init__.py

@@ -0,0 +1,2 @@
+from .client import CoalesceClient
+from .credentials import CoalesceCredentials

castor_extractor-0.24.9/castor_extractor/transformation/coalesce/client/client.py

@@ -0,0 +1,180 @@
+from http import HTTPStatus
+from typing import Iterator, Optional
+
+from ....utils import APIClient, BearerAuth, RequestSafeMode, SerializedAsset
+from ..assets import CoalesceAsset, CoalesceQualityAsset
+from .credentials import CoalesceCredentials
+from .endpoint import (
+    CoalesceEndpointFactory,
+)
+from .type import NodeIDToNamesMapping
+from .utils import column_names_per_node, is_test, test_names_per_node
+
+_LIMIT_MAX = 1_000
+_MAX_ERRORS = 50
+
+
+def _run_result_payload(result: dict, query_result: dict) -> dict:
+    return {
+        "node_id": result["nodeID"],
+        "node_name": result["name"],
+        "test_name": query_result["name"],
+        "start_time": query_result["startTime"],
+        "end_time": query_result["endTime"],
+        "status": query_result["status"],
+        "success": query_result["success"],
+        "isRunning": query_result["isRunning"],
+    }
+
+
+COALESCE_SAFE_MODE = RequestSafeMode(
+    status_codes=(HTTPStatus.INTERNAL_SERVER_ERROR,),
+    max_errors=_MAX_ERRORS,
+)
+COALESCE_TIMEOUT_SECONDS = 90
+
+
+class CoalesceBearerAuth(BearerAuth):
+    """Bearer Authentication for Coalesce"""
+
+    def fetch_token(self) -> Optional[str]:
+        pass
+
+    def __init__(self, token: str):
+        self._token = token
+
+
+class CoalesceClient(APIClient):
+    """REST API client to extract data from Coalesce"""
+
+    def __init__(
+        self,
+        credentials: CoalesceCredentials,
+    ):
+        auth = CoalesceBearerAuth(token=credentials.token)
+        super().__init__(
+            host=credentials.host,
+            auth=auth,
+            safe_mode=COALESCE_SAFE_MODE,
+            timeout=COALESCE_TIMEOUT_SECONDS,
+        )
+
+    def _fetch_environments(self) -> Iterator[dict]:
+        endpoint = CoalesceEndpointFactory.environments()
+        result = self._get(endpoint=endpoint)
+        return result["data"]
+
+    def _node_details(self, environment_id: int, node_id: str) -> dict:
+        endpoint = CoalesceEndpointFactory.nodes(
+            environment_id=environment_id, node_id=node_id
+        )
+        return self._get(endpoint=endpoint)
+
+    def _fetch_env_nodes(self, environment_id: int) -> SerializedAsset:
+        endpoint = CoalesceEndpointFactory.nodes(environment_id=environment_id)
+        result = self._get(endpoint=endpoint)
+        nodes: list[dict] = []
+        for node in result["data"]:
+            details = self._node_details(environment_id, node["id"])
+            nodes.append({**node, **details})
+        return nodes
+
+    def _fetch_all_nodes(self) -> SerializedAsset:
+        nodes: list[dict] = []
+        for environment in self._fetch_environments():
+            environment_id = environment["id"]
+            nodes.extend(self._fetch_env_nodes(environment_id))
+        return nodes
+
+    def _fetch_runs(self, starting_from: str) -> SerializedAsset:
+        """
+        fetch runs, per environment;
+        we break per environment to lower the chance of exceeding the 1k limit
+        """
+        runs: list[dict] = []
+        for environment in self._fetch_environments():
+            environment_id = environment["id"]
+            runs.extend(
+                self._fetch_recent_runs_per_env(environment_id, starting_from)
+            )
+        return runs
+
+    def _fetch_recent_runs_per_env(
+        self, environment_id: int, starting_from: str
+    ) -> SerializedAsset:
+        endpoint = CoalesceEndpointFactory.runs()
+        params = {
+            "environmentID": environment_id,
+            "limit": _LIMIT_MAX,
+            "orderBy": "runEndTime",
+            "orderByDirection": "asc",
+            "startingFrom": starting_from,
+        }
+        result = self._get(endpoint=endpoint, params=params)
+        return result["data"]
+
+    def _fetch_run_results(self, run_id: str) -> SerializedAsset:
+        endpoint = CoalesceEndpointFactory.run_results(run_id)
+        result = self._get(endpoint=endpoint)
+        return result["data"]
+
+    def _run_results_by_run(
+        self,
+        run_id: str,
+        test_names: NodeIDToNamesMapping,
+        column_names: NodeIDToNamesMapping,
+    ) -> SerializedAsset:
+        run_results: list[dict] = []
+        for result in self._fetch_run_results(run_id):
+            node_id = result["nodeID"]
+            for query_result in result["queryResults"]:
+                _is_test = is_test(
+                    query_result,
+                    node_id,
+                    test_names,
+                    column_names,
+                )
+                if not _is_test:
+                    continue
+                run_result = _run_result_payload(result, query_result)
+                run_results.append(run_result)
+        return run_results
+
+    def _run_results_by_env(
+        self, environment_id: int, starting_from: str
+    ) -> SerializedAsset:
+        run_results: list[dict] = []
+        nodes = self._fetch_env_nodes(environment_id)
+        test_names = test_names_per_node(nodes)
+        column_names = column_names_per_node(nodes)
+        runs = self._fetch_recent_runs_per_env(environment_id, starting_from)
+
+        for run in runs:
+            run_id = run["id"]
+            _results = self._run_results_by_run(
+                run_id, test_names, column_names
+            )
+            run_results.extend(_results)
+        return run_results
+
+    def _fetch_all_run_results(self, starting_from: str) -> SerializedAsset:
+        run_results: list[dict] = []
+
+        for environment in self._fetch_environments():
+            environment_id = environment["id"]
+            _results = self._run_results_by_env(environment_id, starting_from)
+            run_results.extend(_results)
+
+        return run_results
+
+    def fetch(
+        self, asset: CoalesceAsset, starting_from=None
+    ) -> SerializedAsset:
+        """Extract the given Coalesce Asset"""
+        if asset in (CoalesceAsset.NODES, CoalesceQualityAsset.NODES):
+            return self._fetch_all_nodes()
+        elif asset == CoalesceQualityAsset.RUN_RESULTS:
+            return self._fetch_all_run_results(starting_from=starting_from)
+        raise AssertionError(
+            f"Asset {asset} is not supported by CoalesceClient"
+        )

castor_extractor-0.24.9/castor_extractor/transformation/coalesce/client/credentials.py

@@ -0,0 +1,23 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+CASTOR_ENV_PREFIX = "CASTOR_COALESCE_"
+
+
+class CoalesceCredentials(BaseSettings):
+    """Class to handle Coalesce rest API permissions"""
+
+    model_config = SettingsConfigDict(
+        env_prefix=CASTOR_ENV_PREFIX,
+        extra="ignore",
+        populate_by_name=True,
+    )
+
+    host: str
+    token: str = Field(repr=False)
+
+    @property
+    def token_payload(self) -> dict[str, str]:
+        return {
+            "client_secret": self.token,
+        }

castor_extractor-0.24.9/castor_extractor/transformation/coalesce/client/endpoint.py

@@ -0,0 +1,42 @@
+from typing import Optional
+
+
+class CoalesceEndpointFactory:
+    """Provide endpoints to hit Coalesce API"""
+
+    @classmethod
+    def environments(cls, environment_id: Optional[int] = None) -> str:
+        """
+        When specified, concatenate environment_id at the end to fetch details.
+        Otherwise, list existing environments.
+        """
+        base = "api/v1/environments"
+        if environment_id:
+            return base + f"/{environment_id}"
+        return base
+
+    @classmethod
+    def nodes(cls, environment_id: int, node_id: Optional[str] = None) -> str:
+        """
+        When specified, concatenate node_id at the end to fetch details.
+        Otherwise, list existing nodes in the given environment.
+        """
+        base = f"api/v1/environments/{environment_id}/nodes"
+        if node_id:
+            return base + f"/{node_id}"
+        return base
+
+    @classmethod
+    def runs(cls) -> str:
+        """
+        Get runs (additional filtering can be done in the body)
+        """
+        base = "api/v1/runs"
+        return base
+
+    @classmethod
+    def run_results(cls, run_id: str) -> str:
+        """
+        get run results (including success/fail for tests), given a run id
+        """
+        return f"api/v1/runs/{run_id}/results"

castor_extractor-0.24.9/castor_extractor/transformation/coalesce/client/type.py

@@ -0,0 +1 @@
+NodeIDToNamesMapping = dict[str, set[str]]

castor_extractor-0.24.9/castor_extractor/transformation/coalesce/client/utils.py

@@ -0,0 +1,52 @@
+from ....utils import SerializedAsset
+from .type import NodeIDToNamesMapping
+
+_NULL_SUFFIX = ": Null"
+_UNIQUE_SUFFIX = ": Unique"
+
+
+def is_test(
+    query_result: dict,
+    node_id: str,
+    test_names: NodeIDToNamesMapping,
+    column_names: NodeIDToNamesMapping,
+) -> bool:
+    """
+    checks whether a query result is a test result or not.
+
+    all this implementation can soon be replaced by checking whether
+    query_result['type'] == 'sqlTest', which should be GA Apr 28th 2025
+    """
+    # test scoped on the node (table)
+    result_name = query_result["name"]
+    if result_name in test_names.get(node_id, {}):
+        return True
+
+    # test scoped on the column
+    if result_name.endswith(_NULL_SUFFIX) or result_name.endswith(
+        _UNIQUE_SUFFIX
+    ):
+        column_name = result_name.split(":")[0]
+        if column_name in column_names.get(node_id, {}):
+            return True
+    return False
+
+
+def test_names_per_node(nodes: SerializedAsset) -> NodeIDToNamesMapping:
+    """mapping nodeID: set(testName)"""
+    mapping: dict[str, set[str]] = {}
+    for node in nodes:
+        node_id = node["id"]
+        tests = node.get("metadata", {}).get("appliedNodeTests", [])
+        mapping[node_id] = {test["name"] for test in tests}
+    return mapping
+
+
+def column_names_per_node(nodes: SerializedAsset) -> NodeIDToNamesMapping:
+    """mapping nodeID: set(columnNames)"""
+    mapping: dict[str, set[str]] = {}
+    for node in nodes:
+        node_id = node["id"]
+        columns = node.get("metadata", {}).get("columns", [])
+        mapping[node_id] = {column["name"] for column in columns}
+    return mapping

castor_extractor-0.24.9/castor_extractor/transformation/coalesce/client/utils_test.py

@@ -0,0 +1,54 @@
+from .utils import is_test
+
+
+def test_is_test():
+    test_names = {"some-uuid": {"check-mirrors", "check-seatbelt"}}
+    column_names = {"some-uuid": {"carthago", "delenda", "est"}}
+
+    happy_node_test = is_test(
+        query_result={"name": "check-mirrors"},
+        node_id="some-uuid",
+        test_names=test_names,
+        column_names=column_names,
+    )
+    assert happy_node_test is True
+
+    unknown_node_test = is_test(
+        query_result={"name": "check-engine"},
+        node_id="some-uuid",
+        test_names=test_names,
+        column_names=column_names,
+    )
+    assert unknown_node_test is False
+
+    happy_column_test_unique = is_test(
+        query_result={"name": "carthago: Unique"},
+        node_id="some-uuid",
+        test_names=test_names,
+        column_names=column_names,
+    )
+    assert happy_column_test_unique is True
+
+    happy_column_test_null = is_test(
+        query_result={"name": "carthago: Null"},
+        node_id="some-uuid",
+        test_names=test_names,
+        column_names=column_names,
+    )
+    assert happy_column_test_null is True
+
+    unknown_column_test = is_test(
+        query_result={"name": "rome: Unique"},
+        node_id="some-uuid",
+        test_names=test_names,
+        column_names=column_names,
+    )
+    assert unknown_column_test is False
+
+    unknown_node_id_test = is_test(
+        query_result={"name": "whatever: Unique"},
+        node_id="unknown-uuid",
+        test_names=test_names,
+        column_names=column_names,
+    )
+    assert unknown_node_id_test is False

{castor_extractor-0.24.7 → castor_extractor-0.24.9}/castor_extractor/visualization/tableau/client/client_metadata_api.py

@@ -92,6 +92,34 @@ def gql_query_scroll(
             break
 
 
+def _deduplicate(result_pages: Iterator[SerializedAsset]) -> SerializedAsset:
+    """
+    Sometimes assets are duplicated, which triggers UniqueViolation errors
+    during store_all down the line.
+
+    We suspect the offset pagination to be the root cause, because we had no
+    problem until recently, when we switched from cursor pagination to offset
+    pagination (for performance reasons)
+    https://help.tableau.com/current/api/metadata_api/en-us/docs/meta_api_examples.html#pagination
+
+    This is a straightforward solution to remove these duplicates directly at
+    extraction.
+    We don't show warnings because duplicates are expected, and we keep only
+    the first occurrence since those duplicates are probably identical.
+    """
+    deduplicated: SerializedAsset = []
+    seen_ids: set[str] = set()
+    for page in result_pages:
+        for asset in page:
+            asset_id = asset["id"]
+            if asset_id in seen_ids:
+                # skip duplicate
+                continue
+            deduplicated.append(asset)
+            seen_ids.add(asset_id)
+    return deduplicated
+
+
 class TableauClientMetadataApi:
     """
     Calls the MetadataAPI, using graphQL
@@ -118,7 +146,7 @@ class TableauClientMetadataApi:
             fields=fields,
             page_size=page_size,
         )
-        return [asset for page in result_pages for asset in page]
+        return _deduplicate(result_pages)
 
     def _page_size(self, asset: TableauAsset) -> int:
         return (

castor_extractor-0.24.9/castor_extractor/visualization/tableau/client/client_metadata_api_test.py

@@ -0,0 +1,31 @@
+from .client_metadata_api import _deduplicate
+
+
+def test__deduplicate():
+    result_pages = iter(
+        [
+            [
+                {"id": 1, "name": "workbook_1"},
+                {"id": 2, "name": "workbook_2"},
+            ],
+            [
+                {"id": 1, "name": "workbook_1"},
+                {"id": 3, "name": "workbook_3"},
+                {"id": 4, "name": "workbook_4"},
+            ],
+            [
+                {"id": 4, "name": "workbook_4"},
+                {"id": 5, "name": "workbook_5"},
+                {"id": 5, "name": "workbook_5"},
+                {"id": 5, "name": "workbook_5"},
+            ],
+            [
+                {"id": 1, "name": "workbook_1"},
+                {"id": 3, "name": "workbook_3"},
+            ],
+        ]
+    )
+    deduplicated = _deduplicate(result_pages)
+    assert len(deduplicated) == 5
+    deduplicated_keys = {item["id"] for item in deduplicated}
+    assert deduplicated_keys == {1, 2, 3, 4, 5}

{castor_extractor-0.24.7 → castor_extractor-0.24.9}/castor_extractor/warehouse/databricks/format.py

@@ -168,7 +168,7 @@ class DatabricksFormatter:
                 "schema_name": None,
                 "query_text": q["query_text"],
                 "user_id": q["user_id"],
-                "user_name": q["user_name"],
+                "user_name": q.get("user_name"),
                 "start_time": start_time,
                 "end_time": end_time,
             }

{castor_extractor-0.24.7 → castor_extractor-0.24.9}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=61.2"]
 
 [tool.poetry]
 name = "castor-extractor"
-version = "0.24.7"
+version = "0.24.9"
 description = "Extract your metadata assets."
 authors = ["Castor <support@castordoc.com>"]
 license = "EULA"