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

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

@@ -1,5 +1,17 @@
 # Changelog
 
+## 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.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.24.57
+Version: 0.25.2
 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,18 @@ For any questions or bug report, contact us at [support@coalesce.io](mailto:supp
 
 # Changelog
 
+## 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.2/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.2}/castor_extractor/commands/extract_powerbi.py

@@ -9,10 +9,21 @@ 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="file path to json certificate file with"
+        "keys: private_key, thumbprint, public_certificate",
+    )
     parser.add_argument(
         "-sc",
         "--scopes",

castor_extractor-0.25.2/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.2/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.2/castor_extractor/visualization/count/client/__init__.py

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

castor_extractor-0.25.2/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.2/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.2/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.2/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.2/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.2/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.2/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.2/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.2}/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.2}/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.2}/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,12 @@ CLIENT_APP_BASE = "https://login.microsoftonline.com"
 REST_API_BASE_PATH = "https://api.powerbi.com/v1.0/myorg"
 
 
+class PowerbiCertificate(BaseModel):
+    public_certificate: Optional[str] = None
+    private_key: str
+    thumbprint: str
+
+
 class PowerbiCredentials(BaseSettings):
     """Class to handle PowerBI rest API permissions"""
 
@@ -21,7 +27,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.2}/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.2}/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.2}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=61.2"]
 
 [tool.poetry]
 name = "castor-extractor"
-version = "0.24.57"
+version = "0.25.2"
 description = "Extract your metadata assets."
 authors = ["Castor <support@castordoc.com>"]
 license = "EULA"
@@ -72,6 +72,7 @@ databricks-sql-connector = { version = "3.6.0", optional = true }
 
 [tool.poetry.extras]
 bigquery = ["sqlalchemy-bigquery", "numpy"]
+count = ["sqlalchemy-bigquery"]
 databricks = ["databricks-sql-connector", "pandas", "numpy"]
 dbt = []
 looker = ["looker-sdk"]
@@ -122,6 +123,7 @@ deptry = ">=0.21"
 [tool.poetry.scripts]
 castor-extract-bigquery = "castor_extractor.commands.extract_bigquery:main"
 castor-extract-confluence = "castor_extractor.commands.extract_confluence:main"
+castor-extract-count = "castor_extractor.commands.extract_count:main"
 castor-extract-databricks = "castor_extractor.commands.extract_databricks:main"
 castor-extract-domo = "castor_extractor.commands.extract_domo:main"
 castor-extract-looker = "castor_extractor.commands.extract_looker:main"