castor-extractor 0.24.54__tar.gz → 0.24.57__tar.gz

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/CHANGELOG.md

@@ -1,5 +1,20 @@
 # Changelog
 
+## 0.24.57 - 2025-09-24
+
+* Sigma:
+  * fix pagination
+  * remove redundant element lineages endpoint
+  * extract data model sources
+
+## 0.24.56 - 2025-09-24
+
+* bump dependencies
+
+## 0.24.55 - 2025-09-19
+
+* Fix encoding in LocalStorage - force to utf-8
+
 ## 0.24.54 - 2025-09-18
 
 * SqlServer: fix typo in the extraction query of schemas

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.24.54
+Version: 0.24.57
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -36,7 +36,7 @@ Requires-Dist: google-api-core (>=2.1.1,<3.0.0)
 Requires-Dist: google-api-python-client (>=2.121.0,<3.0.0) ; extra == "lookerstudio" or extra == "all"
 Requires-Dist: google-auth (>=2,<3)
 Requires-Dist: google-cloud-core (>=2.1.0,<3.0.0)
-Requires-Dist: google-cloud-storage (>=2,<3)
+Requires-Dist: google-cloud-storage (>=3.1.0,<4.0.0)
 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 (>=25.0.0,<26.0.0) ; extra == "looker" or extra == "all"
@@ -215,6 +215,21 @@ For any questions or bug report, contact us at [support@coalesce.io](mailto:supp
 
 # Changelog
 
+## 0.24.57 - 2025-09-24
+
+* Sigma:
+  * fix pagination
+  * remove redundant element lineages endpoint
+  * extract data model sources
+
+## 0.24.56 - 2025-09-24
+
+* bump dependencies
+
+## 0.24.55 - 2025-09-19
+
+* Fix encoding in LocalStorage - force to utf-8
+
 ## 0.24.54 - 2025-09-18
 
 * SqlServer: fix typo in the extraction query of schemas

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/castor_extractor/utils/__init__.py

@@ -26,7 +26,7 @@ from .collection import (
     group_by,
     mapping_from_rows,
 )
-from .constants import OUTPUT_DIR
+from .constants import ENCODING_UTF8, OUTPUT_DIR
 from .deprecate import deprecate_python
 from .env import from_env
 from .files import explode, search_files

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/castor_extractor/utils/constants.py

@@ -1 +1,2 @@
 OUTPUT_DIR = "CASTOR_OUTPUT_DIRECTORY"
+ENCODING_UTF8 = "utf-8"

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/castor_extractor/utils/store.py

@@ -4,6 +4,7 @@ from collections.abc import Iterable, Iterator
 from io import StringIO
 from typing import Optional
 
+from .constants import ENCODING_UTF8
 from .formatter import CsvFormatter, Formatter
 from .time import current_timestamp
 from .write import timestamped_filename
@@ -58,11 +59,11 @@ class LocalStorage(AbstractStorage):
 
     def put(self, name: str, data: Iterable[dict]) -> str:
         path = self.path(name)
-        with open(path, "w") as file:
+        with open(path, "w", encoding=ENCODING_UTF8) as file:
             self._formatter.serialize(file, data)
         return path
 
     def get(self, name: str) -> Iterator[dict]:
         path = self.path(name)
-        with open(path, "r") as file:
+        with open(path, "r", encoding=ENCODING_UTF8) as file:
             return self._formatter.deserialize(StringIO(file.read()))

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/castor_extractor/utils/write.py

@@ -3,11 +3,11 @@ import logging
 import os
 import sys
 from datetime import datetime
+from importlib.metadata import version
 from typing import Any
 
-import pkg_resources
+from ..utils import ENCODING_UTF8
 
-ENCODING = "utf8"
 SUMMARY_FILENAME = "summary.json"
 
 logger = logging.getLogger(__name__)
@@ -33,15 +33,14 @@ def write_json(filename: str, data: Any):
     """
     write the data to a json file at path filename
     """
-    with open(filename, "w", encoding=ENCODING) as f:
+    with open(filename, "w", encoding=ENCODING_UTF8) as f:
         json.dump(data, f)
         logger.info(f"Wrote output file: {filename} ({f.tell()} bytes)")
 
 
 def _current_version() -> str:
     """fetch the current version of castor extractor running"""
-    packages = pkg_resources.require("castor-extractor")
-    return str(packages[0].version)
+    return version("castor-extractor")
 
 
 def write_summary(output_directory: str, ts: int, **kwargs):

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/castor_extractor/visualization/sigma/assets.py

@@ -5,11 +5,11 @@ class SigmaAsset(ExternalAsset):
     """Sigma assets"""
 
     DATAMODELS = "datamodels"
+    DATAMODEL_SOURCES = "datamodel_sources"
     DATASETS = "datasets"
     DATASET_SOURCES = "dataset_sources"
     ELEMENTS = "elements"
     FILES = "files"
-    LINEAGES = "lineages"
     MEMBERS = "members"
     QUERIES = "queries"
     WORKBOOKS = "workbooks"

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/castor_extractor/visualization/sigma/client/client.py

@@ -1,17 +1,13 @@
 import logging
 from collections.abc import Iterator
-from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from http import HTTPStatus
 from typing import Callable, Iterable, Optional
 
-from pydantic import BaseModel
-
 from ....utils import (
     APIClient,
     RequestSafeMode,
     fetch_all_pages,
-    retry,
 )
 from ..assets import SigmaAsset
 from .authentication import SigmaBearerAuth
@@ -55,38 +51,12 @@ 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
 
 
-class LineageContext(BaseModel):
-    """all info needed to build the endpoint for lineage retrieval"""
-
-    workbook_id: str
-    element_id: str
-
-
-class Lineage(BaseModel):
-    """holds response from lineage API and context used to retrieve it"""
-
-    lineage: dict
-    context: LineageContext
-
-
 class SigmaClient(APIClient):
-    def __init__(
-        self,
-        credentials: SigmaCredentials,
-        safe_mode: Optional[RequestSafeMode] = None,
-    ):
+    def __init__(self, credentials: SigmaCredentials):
         auth = SigmaBearerAuth(
             host=credentials.host,
             token_payload=credentials.token_payload,
@@ -96,7 +66,7 @@ class SigmaClient(APIClient):
             auth=auth,
             headers=_SIGMA_HEADERS,
             timeout=_SIGMA_TIMEOUT_S,
-            safe_mode=safe_mode or SIGMA_SAFE_MODE,
+            safe_mode=SIGMA_SAFE_MODE,
         )
 
     def _get_paginated(
@@ -175,68 +145,6 @@ class SigmaClient(APIClient):
                     page=page, workbook_id=workbook_id
                 )
 
-    @retry(
-        (ConnectionError,),
-        max_retries=_RETRY_NUMBER,
-        base_ms=_RETRY_BASE_MS,
-        log_exc_info=True,
-    )
-    def _get_lineage(self, lineage_context: LineageContext) -> Lineage:
-        """
-        return the lineage from API and other ids needed to characterize
-        lineage in castor
-        """
-        workbook_id = lineage_context.workbook_id
-        element_id = lineage_context.element_id
-        endpoint = SigmaEndpointFactory.lineage(workbook_id, element_id)
-        return Lineage(lineage=self._get(endpoint), context=lineage_context)
-
-    @staticmethod
-    def _lineage_context(elements: list[dict]) -> list[LineageContext]:
-        """
-        Helper function to prepare context for lineage retrieval.
-        Elements without associated columns are skipped.
-        """
-        contexts: list[LineageContext] = []
-        for element in elements:
-            if element.get("columns") is None:
-                continue
-
-            context = LineageContext(
-                workbook_id=element["workbook_id"],
-                element_id=element["elementId"],
-            )
-            contexts.append(context)
-        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)
-
-        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
@@ -266,6 +174,13 @@ class SigmaClient(APIClient):
 
             yield from self._yield_deduplicated_queries(queries, workbook_id)
 
+    def _get_all_datamodel_sources(
+        self, datamodels: list[dict]
+    ) -> Iterator[dict]:
+        yield from SigmaSourcesTransformer(
+            self, table_id_key="tableId"
+        ).get_datamodel_sources(datamodels)
+
     def _get_all_dataset_sources(self, datasets: list[dict]) -> Iterator[dict]:
         yield from SigmaSourcesTransformer(self).get_dataset_sources(datasets)
 
@@ -277,14 +192,22 @@ class SigmaClient(APIClient):
     def fetch(
         self,
         asset: SigmaAsset,
+        datamodels: Optional[list[dict]] = None,
         datasets: Optional[list[dict]] = None,
-        elements: Optional[list[dict]] = None,
         workbooks: Optional[list[dict]] = None,
     ) -> Iterator[dict]:
         """Returns the needed metadata for the queried asset"""
         if asset == SigmaAsset.DATAMODELS:
             yield from self._get_all_datamodels()
 
+        elif asset == SigmaAsset.DATAMODEL_SOURCES:
+            if datamodels is None:
+                raise ValueError(
+                    "Missing data models to extract data model sources"
+                )
+
+            yield from self._get_all_datamodel_sources(datamodels)
+
         elif asset == SigmaAsset.DATASETS:
             yield from self._get_all_datasets()
 
@@ -303,12 +226,6 @@ class SigmaClient(APIClient):
         elif asset == SigmaAsset.FILES:
             yield from self._get_all_files()
 
-        elif asset == SigmaAsset.LINEAGES:
-            if elements is None:
-                raise ValueError("Missing elements to extract lineage")
-
-            yield from self._get_all_lineages(elements)
-
         elif asset == SigmaAsset.MEMBERS:
             yield from self._get_all_members()
 

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/castor_extractor/visualization/sigma/client/endpoints.py

@@ -19,6 +19,10 @@ class SigmaEndpointFactory:
     def datamodels(cls) -> str:
         return f"v2/{cls.DATAMODELS}"
 
+    @classmethod
+    def datamodel_sources(cls, datamodel_id: str) -> str:
+        return f"v2/{cls.DATAMODELS}/{datamodel_id}/sources"
+
     @classmethod
     def datasets(cls) -> str:
         return f"v2/{cls.DATASETS}"

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/castor_extractor/visualization/sigma/client/pagination.py

@@ -10,7 +10,7 @@ SIGMA_QUERIES_PAGINATION_LIMIT = 50
 
 
 class SigmaPagination(PaginationModel):
-    next_page: Optional[str] = "0"
+    next_page: Optional[str] = None
     entries: list = Field(default_factory=list)
 
     model_config = ConfigDict(
@@ -27,3 +27,23 @@ class SigmaPagination(PaginationModel):
 
     def page_results(self) -> list:
         return self.entries
+
+
+class SigmaTokenPagination(PaginationModel):
+    next_page_token: Optional[str] = ""  # noqa: S105
+    entries: list = Field(default_factory=list)
+
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )
+
+    def is_last(self) -> bool:
+        return not self.next_page_token
+
+    def next_page_payload(self) -> dict:
+        return {"pageToken": self.next_page_token}
+
+    def page_results(self) -> list:
+        return self.entries

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/castor_extractor/visualization/sigma/client/sources_transformer.py

@@ -2,8 +2,9 @@ import logging
 from http import HTTPStatus
 from typing import TYPE_CHECKING, Callable, Iterator
 
-from ....utils import retry_request
+from ....utils import fetch_all_pages, retry_request
 from .endpoints import SigmaEndpointFactory
+from .pagination import SigmaTokenPagination
 
 if TYPE_CHECKING:
     from .client import SigmaClient
@@ -17,8 +18,11 @@ SIGMA_CONNECTION_PATH_SLEEP_MS = 30_000  # 30 seconds
 class SigmaSourcesTransformer:
     """Retrieves asset sources and enhances them with additional information."""
 
-    def __init__(self, api_client: "SigmaClient"):
+    def __init__(
+        self, api_client: "SigmaClient", table_id_key: str = "inodeId"
+    ):
         self.api_client = api_client
+        self.table_id_key = table_id_key
 
     @retry_request(
         status_codes=(HTTPStatus.TOO_MANY_REQUESTS,),
@@ -38,9 +42,9 @@ class SigmaSourcesTransformer:
         logger.info("Mapping table ids to connection and path information")
 
         unique_table_ids = {
-            source["inodeId"]
+            source[self.table_id_key]
             for asset_sources in all_sources
-            for source in asset_sources["sources"]
+            for source in asset_sources.get("sources", [])
             if source["type"] == "table"
         }
 
@@ -49,15 +53,14 @@ class SigmaSourcesTransformer:
             for table_id in unique_table_ids
         }
 
-    @staticmethod
-    def _enhance_table_source(source: dict, table_to_path: dict) -> dict:
+    def _enhance_table_source(self, source: dict, table_to_path: dict) -> dict:
         """
         Combines a single table source with its connection and path information.
         """
         if source["type"] != "table":
             return source
 
-        path_info = table_to_path.get(source["inodeId"], {})
+        path_info = table_to_path.get(source[self.table_id_key], {})
         source["connectionId"] = path_info.get("connectionId")
         source["path"] = path_info.get("path")
         return source
@@ -82,19 +85,35 @@ class SigmaSourcesTransformer:
             }
 
     def _get_all_sources(
-        self, endpoint: Callable[[str], str], asset_ids: set[str]
+        self,
+        endpoint: Callable[[str], str],
+        asset_ids: set[str],
+        with_pagination: bool = False,
     ) -> Iterator[dict]:
         """Returns transformed sources for the given assets"""
         all_sources = []
 
         for asset_id in asset_ids:
-            sources = self.api_client._get(endpoint=endpoint(asset_id))
+            endpoint_url = endpoint(asset_id)
+            if with_pagination:
+                request = self.api_client._get_paginated(endpoint=endpoint_url)
+                sources = list(fetch_all_pages(request, SigmaTokenPagination))
+            else:
+                sources = self.api_client._get(endpoint=endpoint_url)
             all_sources.append({"asset_id": asset_id, "sources": sources})
 
         table_to_path = self._map_table_id_to_connection_path(all_sources)
 
         yield from self._transform_sources(all_sources, table_to_path)
 
+    def get_datamodel_sources(self, datamodels: list[dict]) -> Iterator[dict]:
+        asset_ids = {datamodel["dataModelId"] for datamodel in datamodels}
+        yield from self._get_all_sources(
+            endpoint=SigmaEndpointFactory.datamodel_sources,
+            asset_ids=asset_ids,
+            with_pagination=True,
+        )
+
     def get_dataset_sources(self, datasets: list[dict]) -> Iterator[dict]:
         asset_ids = {dataset["datasetId"] for dataset in datasets}
         yield from self._get_all_sources(

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/castor_extractor/visualization/sigma/extract.py

@@ -23,8 +23,14 @@ def iterate_all_data(
     """Iterate over the extracted data from Sigma"""
 
     logger.info("Extracting DATA MODELS from API")
-    datamodels = client.fetch(SigmaAsset.DATAMODELS)
-    yield SigmaAsset.DATASETS, list(deep_serialize(datamodels))
+    datamodels = list(client.fetch(SigmaAsset.DATAMODELS))
+    yield SigmaAsset.DATASETS, deep_serialize(datamodels)
+
+    logger.info("Extracting DATAMODEL SOURCES from API")
+    datamodel_sources = client.fetch(
+        SigmaAsset.DATAMODEL_SOURCES, datamodels=datamodels
+    )
+    yield SigmaAsset.DATAMODEL_SOURCES, list(deep_serialize(datamodel_sources))
 
     logger.info("Extracting DATASETS from API")
     datasets = list(client.fetch(SigmaAsset.DATASETS))
@@ -62,10 +68,6 @@ def iterate_all_data(
     elements = list(client.fetch(SigmaAsset.ELEMENTS, workbooks=workbooks))
     yield SigmaAsset.ELEMENTS, list(deep_serialize(elements))
 
-    logging.info("Extracting LINEAGES data from API")
-    lineages = client.fetch(SigmaAsset.LINEAGES, elements=elements)
-    yield SigmaAsset.LINEAGES, list(deep_serialize(lineages))
-
 
 def extract_all(**kwargs) -> None:
     """

{castor_extractor-0.24.54 → castor_extractor-0.24.57}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=61.2"]
 
 [tool.poetry]
 name = "castor-extractor"
-version = "0.24.54"
+version = "0.24.57"
 description = "Extract your metadata assets."
 authors = ["Castor <support@castordoc.com>"]
 license = "EULA"
@@ -32,7 +32,7 @@ google-api-core = ">=2.1.1, <3.0.0"
 google-api-python-client = {version = "^2.121.0", optional = true}
 google-auth = "^2"
 google-cloud-core = ">=2.1.0, <3.0.0"
-google-cloud-storage = "^2"
+google-cloud-storage = "^3.1.0"
 google-resumable-media = "^2.0.3"
 googleapis-common-protos = "^1.53.0"
 looker-sdk = { version = "^25.0.0", optional = true }
@@ -197,15 +197,16 @@ pytest-freezegun = "freezegun"
 DEP004 = ["pytest", "requests_mock", "freezegun"]
 # SQL Alchemy needs some dependencies to be installed
 DEP002 = [
+  "numpy",                      # See comment on numpy in dependencies
+  "pandas",                     # See comment on pandas in dependencies
   "psycopg2-binary",
+  "pymssql",
+  "pymysql",
+  "setuptools",
   "snowflake-connector-python",
   "snowflake-sqlalchemy",
   "sqlalchemy-bigquery",
   "sqlalchemy-redshift",
-  "pymysql",
-  "pymssql",
-  "numpy",                      # See comment on numpy in dependencies
-  "pandas",                     # See comment on pandas in dependencies
 ]