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

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

@@ -1,5 +1,33 @@
 # Changelog
 
+## 0.25.6 - 2025-10-06
+
+* PowerBi: Support additional authentication methods
+
+## 0.25.5 - 2025-10-06
+
+* Snowflake: remove extraction of roles and grants
+
+## 0.25.4 - 2025-10-02
+
+* SqlServer: switch database to support clients running their instance on Azure
+
+## 0.25.3 - 2025-10-01
+
+* Sigma: sleep between sources requests to avoid rate limit errors
+
+## 0.25.2 - 2025-09-30
+
+* PowerBi: Support auth with private_key
+
+## 0.25.1 - 2025-09-29
+
+* Sigma: catch ReadTimeouts during elements extraction
+
+## 0.25.0 - 2025-09-15
+
+* Count: adding connector
+
 ## 0.24.57 - 2025-09-24
 
 * Sigma:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.24.57
+Version: 0.25.6
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -16,6 +16,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: all
 Provides-Extra: bigquery
+Provides-Extra: count
 Provides-Extra: databricks
 Provides-Extra: dbt
 Provides-Extra: looker
@@ -57,7 +58,7 @@ Requires-Dist: setuptools (>=78.1)
 Requires-Dist: snowflake-connector-python (>=3.4.0,<4.0.0) ; extra == "snowflake" or extra == "all"
 Requires-Dist: snowflake-sqlalchemy (!=1.2.5,<2.0.0) ; extra == "snowflake" or extra == "all"
 Requires-Dist: sqlalchemy (>=1.4,<1.5)
-Requires-Dist: sqlalchemy-bigquery[bqstorage] (>=1.0.0,<=2.0.0) ; extra == "bigquery" or extra == "all"
+Requires-Dist: sqlalchemy-bigquery[bqstorage] (>=1.0.0,<=2.0.0) ; extra == "bigquery" or extra == "count" or extra == "all"
 Requires-Dist: sqlalchemy-redshift (>=0.8.14,<0.9.0) ; extra == "redshift" or extra == "all"
 Requires-Dist: tableauserverclient (>=0.25.0,<0.26.0) ; extra == "tableau" or extra == "all"
 Requires-Dist: tqdm (>=4.0.0,<5.0.0)
@@ -215,6 +216,34 @@ For any questions or bug report, contact us at [support@coalesce.io](mailto:supp
 
 # Changelog
 
+## 0.25.6 - 2025-10-06
+
+* PowerBi: Support additional authentication methods
+
+## 0.25.5 - 2025-10-06
+
+* Snowflake: remove extraction of roles and grants
+
+## 0.25.4 - 2025-10-02
+
+* SqlServer: switch database to support clients running their instance on Azure
+
+## 0.25.3 - 2025-10-01
+
+* Sigma: sleep between sources requests to avoid rate limit errors
+
+## 0.25.2 - 2025-09-30
+
+* PowerBi: Support auth with private_key
+
+## 0.25.1 - 2025-09-29
+
+* Sigma: catch ReadTimeouts during elements extraction
+
+## 0.25.0 - 2025-09-15
+
+* Count: adding connector
+
 ## 0.24.57 - 2025-09-24
 
 * Sigma:

castor_extractor-0.25.6/castor_extractor/commands/extract_count.py

@@ -0,0 +1,22 @@
+from argparse import ArgumentParser
+
+from castor_extractor.utils import parse_filled_arguments  # type: ignore
+from castor_extractor.visualization import count  # type: ignore
+
+
+def main():
+    parser = ArgumentParser()
+
+    parser.add_argument(
+        "-c",
+        "--credentials",
+        help="GCP credentials as string",
+    )
+    parser.add_argument("-o", "--output", help="Directory to write to")
+    parser.add_argument(
+        "-d",
+        "--dataset_id",
+        help="dataset id, where count info is stored for the current customer",
+    )
+
+    count.extract_all(**parse_filled_arguments(parser))

{castor_extractor-0.24.57 → castor_extractor-0.25.6}/castor_extractor/commands/extract_powerbi.py

@@ -9,10 +9,25 @@ logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
 
 def main():
     parser = ArgumentParser()
+    auth_group = parser.add_mutually_exclusive_group(required=True)
 
     parser.add_argument("-t", "--tenant_id", help="PowerBi tenant ID")
     parser.add_argument("-c", "--client_id", help="PowerBi client ID")
-    parser.add_argument("-s", "--secret", help="PowerBi password")
+    auth_group.add_argument(
+        "-s",
+        "--secret",
+        help="PowerBi password as a string",
+    )
+    auth_group.add_argument(
+        "-cert",
+        "--certificate",
+        help=(
+            "Path to certificate file for authentication. "
+            "Accepts: X.509 certificates (.pem, .crt), "
+            "PKCS#12 files (.p12, .pfx), or JSON files containing "
+            "certificate data or custom authentication secrets."
+        ),
+    )
     parser.add_argument(
         "-sc",
         "--scopes",

{castor_extractor-0.24.57 → castor_extractor-0.25.6}/castor_extractor/utils/client/abstract.py

@@ -26,6 +26,17 @@ class AbstractSourceClient(ABC):
     def execute(self, query: ExtractionQuery) -> Iterator[dict]:
         pass
 
+    def set_database(self, database: Optional[str]) -> None:
+        """
+        Set the active database for this client.
+
+        Some source technologies require establishing a new connection
+        or using a different connection URI when switching databases.
+
+        Default behaviour is to do nothing.
+        """
+        pass
+
 
 class SqlalchemyClient(AbstractSourceClient, ABC):
     def __init__(self, credentials: dict):

{castor_extractor-0.24.57 → castor_extractor-0.25.6}/castor_extractor/utils/client/api/pagination.py

@@ -58,7 +58,7 @@ class PaginationModel(BaseModel):
 def fetch_all_pages(
     request: Callable,
     pagination_model: type[PaginationModel],
-    rate_limit: Optional[int] = None,
+    rate_limit: Optional[float] = None,
 ) -> Iterator:
     """
     Method to return all results of a Paginated API based on the

castor_extractor-0.25.6/castor_extractor/visualization/count/__init__.py

@@ -0,0 +1,3 @@
+from .assets import CountAsset
+from .client import CountClient, CountCredentials
+from .extract import extract_all

castor_extractor-0.25.6/castor_extractor/visualization/count/assets.py

@@ -0,0 +1,11 @@
+from ...types import ExternalAsset
+
+
+class CountAsset(ExternalAsset):
+    """Count assets"""
+
+    CANVASES = "canvases"
+    CANVAS_PERMISSIONS = "canvas_permissions"
+    CELLS = "cells"
+    PROJECTS = "projects"
+    USERS = "users"

castor_extractor-0.25.6/castor_extractor/visualization/count/client/__init__.py

@@ -0,0 +1,2 @@
+from .client import CountClient
+from .credentials import CountCredentials

castor_extractor-0.25.6/castor_extractor/visualization/count/client/client.py

@@ -0,0 +1,50 @@
+import logging
+from dataclasses import asdict
+from typing import Any, Iterator
+
+from ....utils import load_file
+from ....warehouse.bigquery import BigQueryClient
+from ..assets import (
+    CountAsset,
+)
+from .credentials import CountCredentials
+
+logger = logging.getLogger(__name__)
+
+_QUERIES_FOLDER = "queries"
+
+
+class CountClient(BigQueryClient):
+    """
+    Count.co does not currently provide an official API.
+    Instead, metadata such as dashboards, users, and queries is made available through
+    special metadata tables stored in BigQuery.
+
+    This client extends `BigQueryClient` to access and interact with those metadata tables.
+    """
+
+    def __init__(self, credentials: CountCredentials):
+        super().__init__(asdict(credentials))
+        self.project_id = credentials.project_id
+        self.dataset_id = credentials.dataset_id
+
+    def _load_query(self, asset: CountAsset) -> str:
+        query = load_file(
+            f"{_QUERIES_FOLDER}/{asset.name.lower()}.sql", __file__
+        )
+        return query.format(
+            project_id=self.project_id, dataset_id=self.dataset_id
+        )
+
+    def fetch(self, asset: CountAsset) -> Iterator[dict[str, Any]]:
+        """
+        Fetch the asset given as param, by running a BigQuery query.
+        """
+        logger.info(f"Running BigQuery query to fetch: {asset.name}")
+
+        query_str = self._load_query(asset)
+        job = self.client.query(query_str)
+        results = job.result()
+
+        for row in results:
+            yield dict(row)

castor_extractor-0.25.6/castor_extractor/visualization/count/client/credentials.py

@@ -0,0 +1,10 @@
+from pydantic.dataclasses import dataclass
+
+from ....warehouse.bigquery import BigQueryCredentials
+
+
+@dataclass
+class CountCredentials(BigQueryCredentials):
+    """Count credentials extending BigQuery credentials with additional dataset information"""
+
+    dataset_id: str

castor_extractor-0.25.6/castor_extractor/visualization/count/client/queries/canvas_permissions.sql

@@ -0,0 +1,6 @@
+SELECT
+    canvas_key,
+    type,
+    role,
+    user_key
+FROM `{project_id}.{dataset_id}.canvas_permissions`

castor_extractor-0.25.6/castor_extractor/visualization/count/client/queries/canvases.sql

@@ -0,0 +1,6 @@
+SELECT
+    key,
+    project_key,
+    title
+FROM `{project_id}.{dataset_id}.canvases`
+

castor_extractor-0.25.6/castor_extractor/visualization/count/client/queries/cells.sql

@@ -0,0 +1,8 @@
+SELECT
+    key,
+    canvas_key,
+    name,
+    type,
+    connection_key
+FROM `{project_id}.{dataset_id}.cells`
+

castor_extractor-0.25.6/castor_extractor/visualization/count/client/queries/projects.sql

@@ -0,0 +1,5 @@
+SELECT
+    key,
+    name
+FROM `{project_id}.{dataset_id}.projects`
+

castor_extractor-0.25.6/castor_extractor/visualization/count/client/queries/users.sql

@@ -0,0 +1,8 @@
+SELECT
+    key,
+    created_at,
+    name,
+    email,
+    role
+FROM `{project_id}.{dataset_id}.users`
+

castor_extractor-0.25.6/castor_extractor/visualization/count/extract.py

@@ -0,0 +1,54 @@
+import logging
+from typing import Iterable, Iterator, Union
+
+from ...utils import (
+    OUTPUT_DIR,
+    current_timestamp,
+    deep_serialize,
+    from_env,
+    get_output_filename,
+    write_json,
+    write_summary,
+)
+from .assets import (
+    CountAsset,
+)
+from .client import (
+    CountClient,
+    CountCredentials,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def iterate_all_data(
+    client: CountClient,
+) -> Iterable[tuple[CountAsset, Union[list, Iterator, dict]]]:
+    """Iterate over the extracted data from count"""
+
+    for asset in CountAsset:
+        logger.info(f"Extracting {asset.value} from API")
+        data = client.fetch(asset)
+        yield asset, deep_serialize(data)
+
+
+def extract_all(**kwargs) -> None:
+    """
+    Extract data from count BigQuery project
+    Store the output files locally under the given output_directory
+    """
+    _output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
+    dataset_id = kwargs.get("dataset_id")
+    if not dataset_id:
+        raise ValueError("dataset_id is required")
+
+    credentials = CountCredentials(**kwargs)
+    client = CountClient(credentials=credentials)
+
+    ts = current_timestamp()
+
+    for key, data in iterate_all_data(client):
+        filename = get_output_filename(key.name.lower(), _output_directory, ts)
+        write_json(filename, list(data))
+
+    write_summary(_output_directory, ts)

{castor_extractor-0.24.57 → castor_extractor-0.25.6}/castor_extractor/visualization/powerbi/client/__init__.py

@@ -3,5 +3,6 @@ from .credentials import (
     CLIENT_APP_BASE,
     DEFAULT_SCOPE,
     REST_API_BASE_PATH,
+    PowerbiCertificate,
     PowerbiCredentials,
 )

{castor_extractor-0.24.57 → castor_extractor-0.25.6}/castor_extractor/visualization/powerbi/client/authentication.py

@@ -1,11 +1,24 @@
+from typing import Optional, Union
+
 import msal  # type: ignore
 
 from ....utils import BearerAuth
 from .constants import Keys
-from .credentials import PowerbiCredentials
+from .credentials import PowerbiCertificate, PowerbiCredentials
 from .endpoints import PowerBiEndpointFactory
 
 
+def _get_client_credential(
+    secret: Optional[str], certificate: Optional[PowerbiCertificate]
+) -> Union[str, dict]:
+    if secret:
+        return secret
+    if certificate:
+        return certificate.model_dump()
+
+    raise ValueError("Either certificate or secret must be provided.")
+
+
 class PowerBiBearerAuth(BearerAuth):
     def __init__(self, credentials: PowerbiCredentials):
         self.credentials = credentials
@@ -14,10 +27,15 @@ class PowerBiBearerAuth(BearerAuth):
             api_base=self.credentials.api_base,
         )
         authority = endpoint_factory.authority(self.credentials.tenant_id)
+
+        client_credential = _get_client_credential(
+            self.credentials.secret, self.credentials.certificate
+        )
+
         self.app = msal.ConfidentialClientApplication(
             client_id=self.credentials.client_id,
             authority=authority,
-            client_credential=self.credentials.secret,
+            client_credential=client_credential,
         )
 
     def fetch_token(self):

{castor_extractor-0.24.57 → castor_extractor-0.25.6}/castor_extractor/visualization/powerbi/client/credentials.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from pydantic import Field, field_validator
+from pydantic import BaseModel, field_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 DEFAULT_SCOPE = "https://analysis.windows.net/powerbi/api/.default"
@@ -10,6 +10,20 @@ CLIENT_APP_BASE = "https://login.microsoftonline.com"
 REST_API_BASE_PATH = "https://api.powerbi.com/v1.0/myorg"
 
 
+class PowerbiCertificate(BaseModel):
+    """
+    Supports all dict credentials formats supported by PowerBI
+    https://learn.microsoft.com/en-us/python/api/msal/msal.application.confidentialclientapplication
+    """
+
+    client_assertion: Optional[str] = None
+    passphrase: Optional[str] = None
+    private_key: Optional[str] = None
+    private_key_pfx_path: Optional[str] = None
+    public_certificate: Optional[str] = None
+    thumbprint: Optional[str] = None
+
+
 class PowerbiCredentials(BaseSettings):
     """Class to handle PowerBI rest API permissions"""
 
@@ -21,7 +35,8 @@ class PowerbiCredentials(BaseSettings):
 
     client_id: str
     tenant_id: str
-    secret: str = Field(repr=False)
+    secret: Optional[str] = None
+    certificate: Optional[PowerbiCertificate] = None
     api_base: str = REST_API_BASE_PATH
     login_url: str = CLIENT_APP_BASE
     scopes: list[str] = [DEFAULT_SCOPE]

{castor_extractor-0.24.57 → castor_extractor-0.25.6}/castor_extractor/visualization/powerbi/extract.py

@@ -1,6 +1,7 @@
+import json
 import logging
 from collections.abc import Iterable
-from typing import Union
+from typing import Optional, Union
 
 from ...utils import (
     OUTPUT_DIR,
@@ -12,11 +13,22 @@ from ...utils import (
     write_summary,
 )
 from .assets import PowerBiAsset
-from .client import PowerbiClient, PowerbiCredentials
+from .client import PowerbiCertificate, PowerbiClient, PowerbiCredentials
 
 logger = logging.getLogger(__name__)
 
 
+def _load_certificate(
+    certificate: Optional[str],
+) -> Optional[PowerbiCertificate]:
+    if not certificate:
+        return None
+
+    with open(certificate) as file:
+        cert = json.load(file)
+        return PowerbiCertificate(**cert)
+
+
 def iterate_all_data(
     client: PowerbiClient,
 ) -> Iterable[tuple[PowerBiAsset, Union[list, dict]]]:
@@ -36,7 +48,15 @@ def extract_all(**kwargs) -> None:
     Store the output files locally under the given output_directory
     """
     _output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
-    creds = PowerbiCredentials(**kwargs)
+    creds = PowerbiCredentials(
+        client_id=kwargs.get("client_id"),
+        tenant_id=kwargs.get("tenant_id"),
+        secret=kwargs.get("secret"),
+        certificate=_load_certificate(kwargs.get("certificate")),
+        api_base=kwargs.get("api_base"),
+        login_url=kwargs.get("login_url"),
+        scopes=kwargs.get("scopes"),
+    )
     client = PowerbiClient(creds)
     ts = current_timestamp()
 

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

@@ -4,6 +4,8 @@ from functools import partial
 from http import HTTPStatus
 from typing import Callable, Iterable, Optional
 
+from requests import ReadTimeout
+
 from ....utils import (
     APIClient,
     RequestSafeMode,
@@ -114,6 +116,31 @@ class SigmaClient(APIClient):
         request = self._get_paginated(endpoint=SigmaEndpointFactory.workbooks())
         yield from fetch_all_pages(request, SigmaPagination)
 
+    @staticmethod
+    def _safe_fetch_elements(
+        elements: Iterator[dict],
+        workbook_id: str,
+        page_id: str,
+    ) -> Iterator[dict]:
+        """
+        Safely iterates over elements with ReadTimeout handling. In case of
+        said error, it skips the entire rest of the page.
+        """
+        try:
+            for element in elements:
+                if element.get("type") not in _DATA_ELEMENTS:
+                    continue
+                yield {
+                    **element,
+                    "workbook_id": workbook_id,
+                    "page_id": page_id,
+                }
+        except ReadTimeout:
+            logger.warning(
+                f"ReadTimeout for page {page_id} in workbook {workbook_id}"
+            )
+            return
+
     def _get_elements_per_page(
         self, page: dict, workbook_id: str
     ) -> Iterator[dict]:
@@ -122,14 +149,7 @@ class SigmaClient(APIClient):
             SigmaEndpointFactory.elements(workbook_id, page_id)
         )
         elements = fetch_all_pages(request, SigmaPagination)
-        for element in elements:
-            if element.get("type") not in _DATA_ELEMENTS:
-                continue
-            yield {
-                **element,
-                "workbook_id": workbook_id,
-                "page_id": page_id,
-            }
+        yield from self._safe_fetch_elements(elements, workbook_id, page_id)
 
     def _get_all_elements(self, workbooks: list[dict]) -> Iterator[dict]:
         for workbook in workbooks:

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

@@ -1,5 +1,6 @@
 import logging
 from http import HTTPStatus
+from time import sleep
 from typing import TYPE_CHECKING, Callable, Iterator
 
 from ....utils import fetch_all_pages, retry_request
@@ -13,6 +14,7 @@ logger = logging.getLogger(__name__)
 
 SIGMA_CONNECTION_PATH_MAX_RETRY = 1
 SIGMA_CONNECTION_PATH_SLEEP_MS = 30_000  # 30 seconds
+SIGMA_SOURCES_SLEEP_MS = 2 / 1000  # 0.2 seconds
 
 
 class SigmaSourcesTransformer:
@@ -84,6 +86,29 @@ class SigmaSourcesTransformer:
                 "sources": enhanced_sources,
             }
 
+    def _fetch_asset_sources(
+        self,
+        asset_id: str,
+        endpoint: Callable[[str], str],
+        with_pagination: bool,
+    ) -> dict:
+        """Fetches sources for a single asset."""
+        endpoint_url = endpoint(asset_id)
+
+        if with_pagination:
+            request = self.api_client._get_paginated(endpoint=endpoint_url)
+            pages_generator = fetch_all_pages(
+                request=request,
+                pagination_model=SigmaTokenPagination,
+                rate_limit=SIGMA_SOURCES_SLEEP_MS,
+            )
+            sources = list(pages_generator)
+        else:
+            sources = self.api_client._get(endpoint=endpoint_url)
+            sleep(SIGMA_SOURCES_SLEEP_MS)
+
+        return {"asset_id": asset_id, "sources": sources}
+
     def _get_all_sources(
         self,
         endpoint: Callable[[str], str],
@@ -91,16 +116,10 @@ class SigmaSourcesTransformer:
         with_pagination: bool = False,
     ) -> Iterator[dict]:
         """Returns transformed sources for the given assets"""
-        all_sources = []
-
-        for asset_id in asset_ids:
-            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})
+        all_sources = [
+            self._fetch_asset_sources(asset_id, endpoint, with_pagination)
+            for asset_id in asset_ids
+        ]
 
         table_to_path = self._map_table_id_to_connection_path(all_sources)
 

{castor_extractor-0.24.57 → castor_extractor-0.25.6}/castor_extractor/warehouse/abstract/extract.py

@@ -51,6 +51,7 @@ class SQLExtractionProcessor:
     def _fetch(self, query: ExtractionQuery) -> Iterator[dict]:
         default: Callable[[], Iterator] = lambda: iter(())  # type: ignore
         decorator = safe_mode(self._safe_mode, default)
+        self._client.set_database(query.database)
         decorated_execute = decorator(self._client.execute)
         return decorated_execute(query)
 

{castor_extractor-0.24.57 → castor_extractor-0.25.6}/castor_extractor/warehouse/abstract/query.py

@@ -21,11 +21,22 @@ class AssetNotSupportedError(NotImplementedError):
 
 
 class ExtractionQuery:
-    """contains both statement and parameters for the query"""
+    """
+    Contains useful context to run the query:
+    - the sql statement itself
+    - parameters { ... }
+    - optionally, the target database (can be used to change the engine's URI)
+    """
 
-    def __init__(self, statement: str, params: dict):
+    def __init__(
+        self,
+        statement: str,
+        params: dict,
+        database: Optional[str] = None,
+    ):
         self.statement = statement
         self.params = params
+        self.database = database
 
 
 class AbstractQueryBuilder(ABC):

{castor_extractor-0.24.57 → castor_extractor-0.25.6}/castor_extractor/warehouse/snowflake/extract.py

@@ -24,12 +24,7 @@ SNOWFLAKE_ASSETS: SupportedAssets = {
     WarehouseAssetGroup.FUNCTION: FUNCTIONS_ASSETS,
     WarehouseAssetGroup.QUERY: QUERIES_ASSETS,
     WarehouseAssetGroup.VIEW_DDL: VIEWS_ASSETS,
-    WarehouseAssetGroup.ROLE: (
-        WarehouseAsset.GRANT_TO_ROLE,
-        WarehouseAsset.GRANT_TO_USER,
-        WarehouseAsset.ROLE,
-        WarehouseAsset.USER,
-    ),
+    WarehouseAssetGroup.ROLE: (WarehouseAsset.USER,),
     WarehouseAssetGroup.SNOWFLAKE_LINEAGE: (WarehouseAsset.COLUMN_LINEAGE,),
     WarehouseAssetGroup.EXTERNAL_LINEAGE: EXTERNAL_LINEAGE_ASSETS,
 }