castor-extractor 0.24.34__tar.gz → 0.24.36__tar.gz

{castor_extractor-0.24.34 → castor_extractor-0.24.36}/CHANGELOG.md

@@ -1,5 +1,15 @@
 # Changelog
 
+## 0.24.36 - 2025-08-04
+
+* Sigma:
+  * Refresh token before lineage extraction
+  * Disregard 403 errors during lineage extraction
+
+## 0.24.35 - 2025-07-29
+
+* Coalesce - Fix pagination issue
+
 ## 0.24.34 - 2025-07-02
 
 * SQLServer: multiple databases

{castor_extractor-0.24.34 → castor_extractor-0.24.36}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.24.34
+Version: 0.24.36
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -215,6 +215,16 @@ For any questions or bug report, contact us at [support@coalesce.io](mailto:supp
 
 # Changelog
 
+## 0.24.36 - 2025-08-04
+
+* Sigma:
+  * Refresh token before lineage extraction
+  * Disregard 403 errors during lineage extraction
+
+## 0.24.35 - 2025-07-29
+
+* Coalesce - Fix pagination issue
+
 ## 0.24.34 - 2025-07-02
 
 * SQLServer: multiple databases

{castor_extractor-0.24.34 → castor_extractor-0.24.36}/castor_extractor/transformation/coalesce/client/client.py

@@ -1,31 +1,47 @@
 import logging
+from functools import partial
 from http import HTTPStatus
-from typing import Iterator, Optional
+from typing import Callable, Optional
 
-from requests import ConnectionError
+from pydantic import ValidationError
 
 from ....utils import (
     APIClient,
     BearerAuth,
     RequestSafeMode,
     SerializedAsset,
+    fetch_all_pages,
 )
 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
+from .pagination import CoalescePagination
+
+logger = logging.getLogger(__name__)
+
+
+COALESCE_PAGE_SIZE = 300
+COALESCE_PAGE_SIZE_RUN_RESULTS = 1_000
+
+COALESCE_TIMEOUT_SECONDS = 90
 
-_LIMIT_MAX = 1_000
 _MAX_ERRORS = 200
 
-logger = logging.getLogger(__name__)
+COALESCE_SAFE_MODE = RequestSafeMode(
+    status_codes=(HTTPStatus.INTERNAL_SERVER_ERROR,),
+    max_errors=_MAX_ERRORS,
+)
 
 
-def _run_result_payload(result: dict, query_result: dict) -> dict:
+def _run_result_payload(
+    environment_id: str,
+    result: dict,
+    query_result: dict,
+) -> dict:
     return {
+        "environment_id": environment_id,
         "node_id": result["nodeID"],
         "node_name": result["name"],
         "test_name": query_result["name"],
@@ -37,13 +53,6 @@ def _run_result_payload(result: dict, query_result: dict) -> dict:
     }
 
 
-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"""
 
@@ -69,65 +78,74 @@ class CoalesceClient(APIClient):
             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
+    def _get_paginated(
+        self,
+        endpoint: str,
+        limit: int = COALESCE_PAGE_SIZE,
+        params: Optional[dict] = None,
+    ) -> Callable:
+        return partial(
+            self._get,
+            retry_on_timeout=False,  # explained in the docstring
+            endpoint=endpoint,
+            params={
+                "limit": limit,
+                **(params or dict()),
+            },
         )
-        return self._get(endpoint=endpoint)
+
+    def _fetch_environments(self) -> SerializedAsset:
+        endpoint = CoalesceEndpointFactory.environments()
+        request = self._get_paginated(endpoint=endpoint)
+        result = fetch_all_pages(request, CoalescePagination)
+        return list(result)
 
     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"]:
-            try:
-                details = self._node_details(environment_id, node["id"])
-                nodes.append({**node, **details})
-            except ConnectionError as e:
-                node_id = node["id"]
-                message = f"ConnectionError, environment: {environment_id}, node: {node_id}"
-                logger.warning(message)
-                raise e
-        return nodes
+        request = self._get_paginated(
+            endpoint=endpoint,
+            params={"detail": "true"},
+        )
+        result = fetch_all_pages(request, CoalescePagination)
+        return [
+            {
+                **node,
+                "environment_id": environment_id,
+            }
+            for node in result
+        ]
 
     def _fetch_all_nodes(self) -> SerializedAsset:
+        environments = self._fetch_environments()
+        total = len(environments)
         nodes: list[dict] = []
-        for environment in self._fetch_environments():
-            environment_id = environment["id"]
-            nodes.extend(self._fetch_env_nodes(environment_id))
+
+        for index, env in enumerate(environments):
+            env_id = env["id"]
+            logger.info(f"Fetching nodes for env #{env_id} - {index}/{total}")
+            try:
+                nodes.extend(self._fetch_env_nodes(env_id))
+            except ValidationError as e:
+                # 500 Server Error: Internal Server Error on Coalesce API
+                logger.warning(
+                    f"Skipping nodes for {env_id} due to the following Error: {e}"
+                )
+            logger.info(f"{len(nodes)} nodes extracted so far")
         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"]
+        request = self._get_paginated(
+            endpoint=endpoint,
+            params=params,
+            limit=COALESCE_PAGE_SIZE_RUN_RESULTS,
+        )
+        return list(fetch_all_pages(request, CoalescePagination))
 
     def _fetch_run_results(self, run_id: str) -> SerializedAsset:
         endpoint = CoalesceEndpointFactory.run_results(run_id)
@@ -136,51 +154,37 @@ class CoalesceClient(APIClient):
 
     def _run_results_by_run(
         self,
+        environment_id: str,
         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(
+                if query_result["type"] != "sqlTest":
+                    continue
+                run_result = _run_result_payload(
+                    environment_id,
+                    result,
                     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
+    def _fetch_all_run_results(
+        self,
+        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)
+        runs = self._fetch_runs(starting_from)
+        total = len(runs)
 
+        for index, run in enumerate(runs):
+            logger.info(f"Extracting run results ({index}/{total})")
+            run_id = run["id"]
+            environment_id = run["environmentID"]
+            current_results = self._run_results_by_run(environment_id, run_id)
+            run_results.extend(current_results)
         return run_results
 
     def fetch(

castor_extractor-0.24.36/castor_extractor/transformation/coalesce/client/pagination.py

@@ -0,0 +1,26 @@
+from typing import Optional, Union
+
+from ....utils import PaginationModel
+
+
+class CoalescePagination(PaginationModel):
+    """
+    Class to handle paginated results for Coalesce
+    See their documentation here
+    https://docs.coalesce.io/docs/api
+    """
+
+    data: list
+    next: Union[Optional[str], Optional[int]] = None
+
+    def is_last(self) -> bool:
+        """Stopping condition for the pagination"""
+        return self.next is None
+
+    def next_page_payload(self):
+        """Payload enabling to generate the request for the next page"""
+        return {"startingFrom": self.next}
+
+    def page_results(self) -> list:
+        """List of results of the current page"""
+        return self.data

{castor_extractor-0.24.34 → castor_extractor-0.24.36}/castor_extractor/utils/client/api/pagination.py

@@ -66,16 +66,19 @@ def fetch_all_pages(
     """
     page_number = 1
     response_payload = request()
+
     paginated_response = pagination_model(**response_payload)
+
     while not paginated_response.is_last():
         logger.debug(f"Fetching page number {page_number}")
         yield from paginated_response.page_results()
         next_page_parameters = paginated_response.next_page_parameters()
-        new_request = partial(request, **next_page_parameters)
+        request_with_pagination = partial(request, **next_page_parameters)
         if rate_limit:
             sleep(rate_limit)
         paginated_response = pagination_model(
-            current_page_payload=next_page_parameters, **new_request()
+            current_page_payload=next_page_parameters,
+            **request_with_pagination(),
         )
         page_number += 1
 

{castor_extractor-0.24.34 → castor_extractor-0.24.36}/castor_extractor/visualization/sigma/client/client.py

@@ -54,6 +54,13 @@ SIGMA_SAFE_MODE = RequestSafeMode(
     max_errors=_VOLUME_IGNORED,
     status_codes=_IGNORED_ERROR_CODES,
 )
+SIGMA_SAFE_MODE_LINEAGE = RequestSafeMode(
+    max_errors=_VOLUME_IGNORED,
+    status_codes=(
+        *_IGNORED_ERROR_CODES,
+        HTTPStatus.FORBIDDEN,
+    ),
+)
 _THREADS_LINEAGE = 10  # empirically found; hit the rate limit with 20 workers
 _RETRY_NUMBER = 1
 _RETRY_BASE_MS = 60_000
@@ -210,18 +217,35 @@ class SigmaClient(APIClient):
         return contexts
 
     def _get_all_lineages(self, elements: list[dict]) -> Iterator[dict]:
+        """
+        The safe mode is temporarily modified to include 403 errors.
+
+        Due to concurrency issues, we force a refresh of the token in hopes that
+        the lineage extraction takes less than the token expiration time of
+        1 hour.
+        """
+        safe_mode = self._safe_mode
+        self._safe_mode = SIGMA_SAFE_MODE_LINEAGE
+
         lineage_context = self._lineage_context(elements)
 
+        self._auth.refresh_token()
+
         with ThreadPoolExecutor(max_workers=_THREADS_LINEAGE) as executor:
             results = executor.map(self._get_lineage, lineage_context)
 
         for lineage in results:
+            if not lineage.lineage:
+                continue
+
             yield {
                 **lineage.lineage,
                 "workbook_id": lineage.context.workbook_id,
                 "element_id": lineage.context.element_id,
             }
 
+        self._safe_mode = safe_mode
+
     @staticmethod
     def _yield_deduplicated_queries(
         queries: Iterable[dict], workbook_id: str

{castor_extractor-0.24.34 → castor_extractor-0.24.36}/pyproject.toml

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

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

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

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

@@ -1,52 +0,0 @@
-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.34/castor_extractor/transformation/coalesce/client/utils_test.py

@@ -1,54 +0,0 @@
-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