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

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

@@ -1,5 +1,28 @@
 # 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:
+  * 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

{castor_extractor-0.24.55 → castor_extractor-0.25.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.24.55
+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,29 @@ 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:
+  * 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

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.55 → 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.55 → 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.55 → 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.55 → 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.55 → 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.55 → castor_extractor-0.25.2}/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"