castor-extractor 0.24.35__tar.gz → 0.24.38__tar.gz

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/CHANGELOG.md

@@ -1,5 +1,19 @@
 # Changelog
 
+## 0.24.38 - 2025-08-07
+
+* Uploader: Support US and EU zones
+
+## 0.24.37 - 2025-08-06
+
+* Sigma: extract data models, dataset sources and workbook sources
+
+## 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

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.24.35
+Version: 0.24.38
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -215,6 +215,20 @@ For any questions or bug report, contact us at [support@coalesce.io](mailto:supp
 
 # Changelog
 
+## 0.24.38 - 2025-08-07
+
+* Uploader: Support US and EU zones
+
+## 0.24.37 - 2025-08-06
+
+* Sigma: extract data models, dataset sources and workbook sources
+
+## 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

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/castor_extractor/commands/upload.py

@@ -3,6 +3,7 @@ from argparse import ArgumentParser
 
 from castor_extractor.uploader import (  # type: ignore
     FileType,
+    Zone,
     upload_any,
 )
 from castor_extractor.utils import parse_filled_arguments  # type: ignore
@@ -40,6 +41,15 @@ def _args() -> ArgumentParser:
         ),
         choices=supported_file_type,
     )
+    supported_zones = [zone.value for zone in Zone]
+    parser.add_argument(
+        "-z",
+        "--zone",
+        help="geographic zone to upload, currently supported are {}, defaults to EU".format(
+            supported_zones,
+        ),
+        choices=supported_zones,
+    )
     return parser
 
 

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/castor_extractor/uploader/__init__.py

@@ -1,2 +1,3 @@
 from .constant import FileType
+from .enums import Zone
 from .upload import upload, upload_any, upload_manifest

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/castor_extractor/uploader/constant.py

@@ -1,9 +1,13 @@
 from enum import Enum
 
 from ..utils import RetryStrategy
+from .enums import Zone
 
 # url of the gcs proxy
-INGEST_URL = "https://ingest.castordoc.com"
+INGEST_URLS = {
+    Zone.EU: "https://ingest.castordoc.com",
+    Zone.US: "https://ingest.us.castordoc.com",
+}
 
 RETRY_BASE_MS = 10_000
 RETRY_JITTER_MS = 1_000

castor_extractor-0.24.38/castor_extractor/uploader/enums.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class Zone(Enum):
+    """Geographic cluster location"""
+
+    EU = "EU"
+    US = "US"

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/castor_extractor/uploader/settings.py

@@ -4,6 +4,7 @@ from pydantic import UUID4, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from .constant import FileType
+from .enums import Zone
 
 UPLOADER_ENV_PREFIX = "CASTOR_UPLOADER_"
 
@@ -22,3 +23,4 @@ class UploaderSettings(BaseSettings):
     file_type: FileType
     source_id: UUID4
     token: str = Field(repr=False)
+    zone: Optional[Zone] = Zone.EU

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/castor_extractor/uploader/upload.py

@@ -10,13 +10,14 @@ import requests
 
 from ..utils.retry import retry
 from .constant import (
-    INGEST_URL,
+    INGEST_URLS,
     PATH_TEMPLATES,
     RETRY_BASE_MS,
     RETRY_JITTER_MS,
     RETRY_STRATEGY,
     FileType,
 )
+from .enums import Zone
 from .env import get_blob_env
 from .settings import UploaderSettings
 from .utils import iter_files
@@ -33,6 +34,7 @@ def _path_and_url(
     source_id: UUID,
     file_type: FileType,
     file_path: str,
+    zone: Zone,
 ) -> tuple[str, str]:
     now = datetime.utcnow()
     timestamp = int(now.timestamp())
@@ -44,7 +46,7 @@ def _path_and_url(
         filename=filename,
     )
 
-    url = f"{INGEST_URL}/{path}"
+    url = f"{INGEST_URLS[zone]}/{path}"
 
     return path, url
 
@@ -61,13 +63,16 @@ def _upload(
     source_id: UUID,
     file_path: str,
     file_type: FileType,
+    zone: Optional[Zone] = Zone.EU,
 ) -> None:
     """
     Upload the given file to Google Cloud Storage (GCS)
     - Don't call GCS API directly
     - Call the ingestion proxy which handles authorisation and uploading
     """
-    path, url = _path_and_url(source_id, file_type, file_path)
+    if not zone:
+        zone = Zone.EU
+    path, url = _path_and_url(source_id, file_type, file_path, zone)
     headers = _headers(token)
     timeout, max_retries = get_blob_env()
 
@@ -97,6 +102,7 @@ def _upload(
 def upload_manifest(
     token: str,
     source_id: UUID,
+    zone: Optional[Zone],
     file_path: Optional[str] = None,
 ) -> None:
     """
@@ -106,13 +112,20 @@ def upload_manifest(
     """
     if not file_path:
         raise ValueError("file path is needed to upload a manifest")
-    _upload(token, source_id, file_path, FileType.DBT)
+    _upload(
+        token=token,
+        source_id=source_id,
+        file_path=file_path,
+        file_type=FileType.DBT,
+        zone=zone,
+    )
 
 
 def upload(
     token: str,
     source_id: UUID,
     file_type: FileType,
+    zone: Optional[Zone],
     file_path: Optional[str] = None,
     directory_path: Optional[str] = None,
 ) -> None:
@@ -133,7 +146,13 @@ def upload(
         raise ValueError(message)
 
     for file_ in files:
-        _upload(token, source_id, file_, file_type)
+        _upload(
+            token=token,
+            source_id=source_id,
+            file_path=file_,
+            file_type=file_type,
+            zone=zone,
+        )
 
 
 def upload_any(**kwargs) -> None:
@@ -156,6 +175,7 @@ def upload_any(**kwargs) -> None:
             token=settings.token,
             source_id=settings.source_id,
             file_path=settings.file_path,
+            zone=settings.zone,
         )
         return None
 
@@ -165,4 +185,5 @@ def upload_any(**kwargs) -> None:
         file_type=file_type,
         file_path=settings.file_path,
         directory_path=settings.directory_path,
+        zone=settings.zone,
     )

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/castor_extractor/uploader/upload_test.py

@@ -1,6 +1,7 @@
 from uuid import UUID
 
-from .constant import INGEST_URL, FileType
+from .constant import INGEST_URLS, FileType
+from .enums import Zone
 from .upload import _path_and_url
 
 
@@ -8,7 +9,8 @@ def test__path():
     source_id = UUID("399a8b22-3187-11ec-8d3d-0242ac130003")
     file_type = FileType.VIZ
     file_path = "filename"
+    zone = Zone.EU
 
-    path, url = _path_and_url(source_id, file_type, file_path)
+    path, url = _path_and_url(source_id, file_type, file_path, zone)
     assert path == f"visualization-{source_id}/{file_path}"
-    assert url == f"{INGEST_URL}/{path}"
+    assert url == f"{INGEST_URLS[Zone.EU]}/{path}"

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/castor_extractor/visualization/sigma/assets.py

@@ -4,10 +4,13 @@ from ...types import ExternalAsset
 class SigmaAsset(ExternalAsset):
     """Sigma assets"""
 
+    DATAMODELS = "datamodels"
     DATASETS = "datasets"
+    DATASET_SOURCES = "dataset_sources"
     ELEMENTS = "elements"
     FILES = "files"
     LINEAGES = "lineages"
     MEMBERS = "members"
     QUERIES = "queries"
     WORKBOOKS = "workbooks"
+    WORKBOOK_SOURCES = "workbook_sources"

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/castor_extractor/visualization/sigma/client/client.py

@@ -1,3 +1,4 @@
+import logging
 from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
@@ -24,6 +25,9 @@ from .pagination import (
     SIGMA_QUERIES_PAGINATION_LIMIT,
     SigmaPagination,
 )
+from .sources_transformer import SigmaSourcesTransformer
+
+logger = logging.getLogger(__name__)
 
 _CONTENT_TYPE = "application/x-www-form-urlencoded"
 
@@ -54,6 +58,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
@@ -128,6 +139,12 @@ class SigmaClient(APIClient):
             params={"limit": limit},
         )
 
+    def _get_all_datamodels(self) -> Iterator[dict]:
+        request = self._get_paginated(
+            endpoint=SigmaEndpointFactory.datamodels()
+        )
+        yield from fetch_all_pages(request, SigmaPagination)
+
     def _get_all_datasets(self) -> Iterator[dict]:
         request = self._get_paginated(endpoint=SigmaEndpointFactory.datasets())
         yield from fetch_all_pages(request, SigmaPagination)
@@ -210,18 +227,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
@@ -251,18 +285,36 @@ class SigmaClient(APIClient):
 
             yield from self._yield_deduplicated_queries(queries, workbook_id)
 
+    def _get_all_dataset_sources(self, datasets: list[dict]) -> Iterator[dict]:
+        yield from SigmaSourcesTransformer(self).get_dataset_sources(datasets)
+
+    def _get_all_workbook_sources(
+        self, workbooks: list[dict]
+    ) -> Iterator[dict]:
+        yield from SigmaSourcesTransformer(self).get_workbook_sources(workbooks)
+
     def fetch(
         self,
         asset: SigmaAsset,
-        workbooks: 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.DATASETS:
+        if asset == SigmaAsset.DATAMODELS:
+            yield from self._get_all_datamodels()
+
+        elif asset == SigmaAsset.DATASETS:
             yield from self._get_all_datasets()
 
+        elif asset == SigmaAsset.DATASET_SOURCES:
+            if datasets is None:
+                raise ValueError("Missing datasets to extract dataset sources")
+
+            yield from self._get_all_dataset_sources(datasets)
+
         elif asset == SigmaAsset.ELEMENTS:
-            if not workbooks:
+            if workbooks is None:
                 raise ValueError("Missing workbooks to extract elements")
 
             yield from self._get_all_elements(workbooks)
@@ -271,15 +323,16 @@ class SigmaClient(APIClient):
             yield from self._get_all_files()
 
         elif asset == SigmaAsset.LINEAGES:
-            if not elements:
+            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()
 
         elif asset == SigmaAsset.QUERIES:
-            if not workbooks:
+            if workbooks is None:
                 raise ValueError("Missing workbooks to extract queries")
 
             yield from self._get_all_queries(workbooks)
@@ -287,5 +340,13 @@ class SigmaClient(APIClient):
         elif asset == SigmaAsset.WORKBOOKS:
             yield from self._get_all_workbooks()
 
+        elif asset == SigmaAsset.WORKBOOK_SOURCES:
+            if workbooks is None:
+                raise ValueError(
+                    "Missing workbooks to extract workbook sources"
+                )
+
+            yield from self._get_all_workbook_sources(workbooks)
+
         else:
             raise ValueError(f"This asset {asset} is unknown")

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/castor_extractor/visualization/sigma/client/endpoints.py

@@ -1,6 +1,7 @@
 class SigmaEndpointFactory:
     """Wrapper class around all endpoints we're using"""
 
+    DATAMODELS = "dataModels"
     DATASETS = "datasets"
     FILES = "files"
     MEMBERS = "members"
@@ -10,10 +11,22 @@ class SigmaEndpointFactory:
     def authentication(cls) -> str:
         return "v2/auth/token"
 
+    @classmethod
+    def connection_path(cls, inode_id: str) -> str:
+        return f"v2/connections/paths/{inode_id}"
+
+    @classmethod
+    def datamodels(cls) -> str:
+        return f"v2/{cls.DATAMODELS}"
+
     @classmethod
     def datasets(cls) -> str:
         return f"v2/{cls.DATASETS}"
 
+    @classmethod
+    def dataset_sources(cls, dataset_id: str) -> str:
+        return f"v2/{cls.DATASETS}/{dataset_id}/sources"
+
     @classmethod
     def elements(cls, workbook_id: str, page_id: str) -> str:
         return f"v2/{cls.WORKBOOKS}/{workbook_id}/pages/{page_id}/elements"
@@ -41,3 +54,7 @@ class SigmaEndpointFactory:
     @classmethod
     def workbooks(cls) -> str:
         return f"v2/{cls.WORKBOOKS}"
+
+    @classmethod
+    def workbook_sources(cls, workbook_id: str) -> str:
+        return f"v2/{cls.WORKBOOKS}/{workbook_id}/sources"

castor_extractor-0.24.38/castor_extractor/visualization/sigma/client/sources_transformer.py

@@ -0,0 +1,94 @@
+import logging
+from typing import TYPE_CHECKING, Callable, Iterator
+
+from .endpoints import SigmaEndpointFactory
+
+if TYPE_CHECKING:
+    from .client import SigmaClient
+
+logger = logging.getLogger(__name__)
+
+
+class SigmaSourcesTransformer:
+    """Retrieves asset sources and enhances them with additional information."""
+
+    def __init__(self, api_client: "SigmaClient"):
+        self.api_client = api_client
+
+    def _map_table_id_to_connection_path(
+        self, all_sources: list
+    ) -> dict[str, dict]:
+        """Maps a table id to its connection and path information."""
+        logger.info("Mapping table ids to connection and path information")
+
+        unique_table_ids = {
+            source["inodeId"]
+            for asset_sources in all_sources
+            for source in asset_sources["sources"]
+            if source["type"] == "table"
+        }
+
+        return {
+            table_id: self.api_client._get(
+                endpoint=SigmaEndpointFactory.connection_path(table_id)
+            )
+            for table_id in unique_table_ids
+        }
+
+    @staticmethod
+    def _enhance_table_source(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"], {})
+        source["connectionId"] = path_info.get("connectionId")
+        source["path"] = path_info.get("path")
+        return source
+
+    def _transform_sources(
+        self, all_sources: list, table_to_path: dict
+    ) -> Iterator[dict]:
+        """
+        Yields all sources, with table sources being enhanced with additional information.
+        """
+        logger.info("Merging sources with table information")
+
+        for asset_sources in all_sources:
+            enhanced_sources = [
+                self._enhance_table_source(source, table_to_path)
+                for source in asset_sources["sources"]
+            ]
+
+            yield {
+                "asset_id": asset_sources["asset_id"],
+                "sources": enhanced_sources,
+            }
+
+    def _get_all_sources(
+        self, endpoint: Callable[[str], str], asset_ids: set[str]
+    ) -> 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))
+            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_dataset_sources(self, datasets: list[dict]) -> Iterator[dict]:
+        asset_ids = {dataset["datasetId"] for dataset in datasets}
+        yield from self._get_all_sources(
+            endpoint=SigmaEndpointFactory.dataset_sources, asset_ids=asset_ids
+        )
+
+    def get_workbook_sources(self, workbooks: list[dict]) -> Iterator[dict]:
+        asset_ids = {workbook["workbookId"] for workbook in workbooks}
+        yield from self._get_all_sources(
+            endpoint=SigmaEndpointFactory.workbook_sources, asset_ids=asset_ids
+        )

castor_extractor-0.24.38/castor_extractor/visualization/sigma/client/sources_transformer_test.py

@@ -0,0 +1,101 @@
+from unittest.mock import Mock
+
+from .sources_transformer import SigmaSourcesTransformer
+
+_ALL_SOURCES = [
+    {
+        "asset_id": "asset1",
+        "sources": [
+            {"type": "dataset", "inodeId": "1234"},  # non-table source
+            {"type": "table", "inodeId": "table1"},
+            {"type": "table", "inodeId": "table2"},
+        ],
+    },
+    {
+        "asset_id": "asset2",
+        "sources": [
+            {"type": "table", "inodeId": "table1"},  # repeated source
+        ],
+    },
+]
+
+
+_TABLE_TO_PATH = {
+    "table1": {
+        "connectionId": "conn1",
+        "path": ["db", "schema", "table1"],
+    },
+    "table2": {
+        "connectionId": "conn2",
+        "path": ["db", "schema", "table2"],
+    },
+}
+
+
+def test__map_table_id_to_connection_path():
+    transformer = SigmaSourcesTransformer(api_client=Mock())
+
+    def mock_get(endpoint):
+        if "table1" in endpoint:
+            return _TABLE_TO_PATH["table1"]
+        elif "table2" in endpoint:
+            return _TABLE_TO_PATH["table2"]
+        else:
+            raise ValueError(f"Unexpected endpoint: {endpoint}")
+
+    transformer.api_client._get.side_effect = mock_get
+
+    result = transformer._map_table_id_to_connection_path(_ALL_SOURCES)
+
+    assert len(result) == 2
+    assert result["table1"] == {
+        "connectionId": "conn1",
+        "path": ["db", "schema", "table1"],
+    }
+    assert result["table2"] == {
+        "connectionId": "conn2",
+        "path": ["db", "schema", "table2"],
+    }
+    assert transformer.api_client._get.call_count == 2
+
+
+def test__transform_sources():
+    transformer = SigmaSourcesTransformer(api_client=Mock())
+
+    result = list(transformer._transform_sources(_ALL_SOURCES, _TABLE_TO_PATH))
+
+    assert len(result) == 2
+
+    asset_1_results = result[0]
+    assert len(asset_1_results["sources"]) == 3
+    actual_sources = sorted(
+        asset_1_results["sources"], key=lambda x: x["inodeId"]
+    )
+    expected_sources = [
+        {"type": "dataset", "inodeId": "1234"},
+        {
+            "type": "table",
+            "inodeId": "table1",
+            "connectionId": "conn1",
+            "path": ["db", "schema", "table1"],
+        },
+        {
+            "type": "table",
+            "inodeId": "table2",
+            "connectionId": "conn2",
+            "path": ["db", "schema", "table2"],
+        },
+    ]
+    expected_sources = sorted(expected_sources, key=lambda x: x["inodeId"])
+    assert actual_sources == expected_sources
+
+    asset_2_results = result[1]
+    assert asset_2_results["asset_id"] == "asset2"
+    assert asset_2_results["sources"] == [
+        {
+            "type": "table",
+            "inodeId": "table1",
+            "connectionId": "conn1",
+            "path": ["db", "schema", "table1"],
+        }
+    ]

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/castor_extractor/visualization/sigma/extract.py

@@ -22,14 +22,30 @@ def iterate_all_data(
 ) -> Iterable[tuple[SigmaAsset, Union[list, Iterator, dict]]]:
     """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))
+
     logger.info("Extracting DATASETS from API")
-    datasets = client.fetch(SigmaAsset.DATASETS)
+    datasets = list(client.fetch(SigmaAsset.DATASETS))
     yield SigmaAsset.DATASETS, list(deep_serialize(datasets))
 
+    logger.info("Extracting DATASET SOURCES from API")
+    dataset_sources = client.fetch(
+        SigmaAsset.DATASET_SOURCES, datasets=datasets
+    )
+    yield SigmaAsset.DATASET_SOURCES, list(deep_serialize(dataset_sources))
+
     logger.info("Extracting WORKBOOKS from API")
     workbooks = list(client.fetch(SigmaAsset.WORKBOOKS))
     yield SigmaAsset.WORKBOOKS, list(deep_serialize(workbooks))
 
+    logger.info("Extracting WORKBOOK SOURCES from API")
+    workbook_sources = client.fetch(
+        SigmaAsset.WORKBOOK_SOURCES, workbooks=workbooks
+    )
+    yield SigmaAsset.WORKBOOKS, list(deep_serialize(workbook_sources))
+
     logger.info("Extracting FILES from API")
     files = client.fetch(SigmaAsset.FILES)
     yield SigmaAsset.FILES, list(deep_serialize(files))

{castor_extractor-0.24.35 → castor_extractor-0.24.38}/pyproject.toml

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