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

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

@@ -1,5 +1,19 @@
 # Changelog
 
+## 0.25.9 - 2025-10-09
+
+* Snowflake: raise an exception when no database is available for extraction
+* Databricks: raise an exception when no database is available for extraction
+* BigQuery: raise an exception when no project is available for extraction
+
+## 0.25.8 - 2025-10-09
+
+* Count: extracting queries and canvas_loads
+
+## 0.25.7 - 2025-10-07
+
+* SqlServer: Ensure database consistency between query and engine
+
 ## 0.25.6 - 2025-10-06
 
 * PowerBi: Support additional authentication methods

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.25.6
+Version: 0.25.9
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -216,6 +216,20 @@ For any questions or bug report, contact us at [support@coalesce.io](mailto:supp
 
 # Changelog
 
+## 0.25.9 - 2025-10-09
+
+* Snowflake: raise an exception when no database is available for extraction
+* Databricks: raise an exception when no database is available for extraction
+* BigQuery: raise an exception when no project is available for extraction
+
+## 0.25.8 - 2025-10-09
+
+* Count: extracting queries and canvas_loads
+
+## 0.25.7 - 2025-10-07
+
+* SqlServer: Ensure database consistency between query and engine
+
 ## 0.25.6 - 2025-10-06
 
 * PowerBi: Support additional authentication methods

{castor_extractor-0.25.6 → castor_extractor-0.25.9}/castor_extractor/commands/extract_sqlserver.py

@@ -11,7 +11,6 @@ def main():
 
     parser.add_argument("-H", "--host", help="MSSQL Host")
     parser.add_argument("-P", "--port", help="MSSQL Port")
-    parser.add_argument("-d", "--database", help="MSSQL Database")
     parser.add_argument("-u", "--user", help="MSSQL User")
     parser.add_argument("-p", "--password", help="MSSQL Password")
 
@@ -47,7 +46,6 @@ def main():
 
     sqlserver.extract_all(
         host=args.host,
-        database=args.database,
         port=args.port,
         user=args.user,
         password=args.password,

castor_extractor-0.25.9/castor_extractor/exceptions.py

@@ -0,0 +1,5 @@
+class NoDatabaseProvidedException(Exception):
+    def __init__(self):
+        super().__init__("""No database eligible for extraction.
+If you are using the db_allow/db_block options, please make sure to use the correct case.
+        """)

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

@@ -26,17 +26,6 @@ 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.25.9/castor_extractor/utils/client/query.py

@@ -0,0 +1,20 @@
+from typing import Optional
+
+
+class ExtractionQuery:
+    """
+    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,
+        database: Optional[str] = None,
+    ):
+        self.statement = statement
+        self.params = params
+        self.database = database

{castor_extractor-0.25.6 → castor_extractor-0.25.9}/castor_extractor/visualization/count/assets.py

@@ -5,7 +5,9 @@ class CountAsset(ExternalAsset):
     """Count assets"""
 
     CANVASES = "canvases"
+    CANVAS_LOADS = "canvas_loads"
     CANVAS_PERMISSIONS = "canvas_permissions"
     CELLS = "cells"
     PROJECTS = "projects"
+    QUERIES = "queries"
     USERS = "users"

{castor_extractor-0.25.6 → castor_extractor-0.25.9}/castor_extractor/visualization/count/client/client.py

@@ -3,6 +3,7 @@ from dataclasses import asdict
 from typing import Any, Iterator
 
 from ....utils import load_file
+from ....warehouse.abstract import TimeFilter
 from ....warehouse.bigquery import BigQueryClient
 from ..assets import (
     CountAsset,
@@ -27,18 +28,21 @@ class CountClient(BigQueryClient):
         super().__init__(asdict(credentials))
         self.project_id = credentials.project_id
         self.dataset_id = credentials.dataset_id
+        self.time_filter = TimeFilter.default()  # setting current date - 1
 
     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
+            project_id=self.project_id,
+            dataset_id=self.dataset_id,
+            extract_date=self.time_filter.day,
         )
 
     def fetch(self, asset: CountAsset) -> Iterator[dict[str, Any]]:
         """
-        Fetch the asset given as param, by running a BigQuery query.
+        Fetch the asset given as a param by running a BigQuery query.
         """
         logger.info(f"Running BigQuery query to fetch: {asset.name}")
 

castor_extractor-0.25.9/castor_extractor/visualization/count/client/queries/canvas_loads.sql

@@ -0,0 +1,14 @@
+WITH view_count AS(
+    SELECT
+        canvas_key,
+        DATE(loaded_at) AS load_day
+    FROM `{project_id}.{dataset_id}.canvas_loads`
+    WHERE TRUE
+        AND date(loaded_at) = '{extract_date}'
+)
+
+SELECT
+    canvas_key,
+    COUNT(*) AS view_count
+FROM view_count
+GROUP BY canvas_key

castor_extractor-0.25.9/castor_extractor/visualization/count/client/queries/queries.sql

@@ -0,0 +1,23 @@
+WITH ranked_queries AS (
+    SELECT
+        cell_key,
+        connection_key,
+        query,
+        started_at,
+        ROW_NUMBER() OVER (
+            PARTITION BY cell_key
+            ORDER BY started_at DESC
+        ) AS rank
+    FROM `{project_id}.{dataset_id}.queries`
+    WHERE TRUE
+        AND DATE(started_at) = '{extract_date}'
+)
+
+SELECT
+    cell_key,
+    connection_key,
+    query,
+    started_at
+FROM ranked_queries
+WHERE TRUE
+    AND rank = 1

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

@@ -268,3 +268,7 @@ class SigmaClient(APIClient):
 
         else:
             raise ValueError(f"This asset {asset} is unknown")
+
+    def test_connection(self) -> None:
+        """Use credentials & verify requesting the API doesn't raise an error"""
+        self._auth.refresh_token()

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

@@ -48,10 +48,9 @@ class SQLExtractionProcessor:
         # dict > set > dict
         return [dict(t) for t in {tuple(d.items()) for d in data}]
 
-    def _fetch(self, query: ExtractionQuery) -> Iterator[dict]:
+    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)
 
@@ -63,7 +62,7 @@ class SQLExtractionProcessor:
         for i, query in enumerate(queries):
             logger.info(f"Extracting {asset.value}: query {i + 1}/{total}")
             # concatenate results of all queries
-            data = chain(data, self._fetch(query))
+            data = chain(data, self.fetch(query))
 
         if self._query_builder.needs_deduplication(asset):
             # cast the list to iterator, but the streaming pipeline is broken in that case

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

@@ -3,6 +3,7 @@ import os
 from abc import ABC, abstractmethod
 from typing import Optional
 
+from ...utils import ExtractionQuery
 from .asset import WarehouseAsset
 from .time_filter import TimeFilter
 
@@ -20,25 +21,6 @@ class AssetNotSupportedError(NotImplementedError):
         super().__init__(msg)
 
 
-class ExtractionQuery:
-    """
-    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,
-        database: Optional[str] = None,
-    ):
-        self.statement = statement
-        self.params = params
-        self.database = database
-
-
 class AbstractQueryBuilder(ABC):
     """
     Build queries necessary to extract warehouse assets.

{castor_extractor-0.25.6 → castor_extractor-0.25.9}/castor_extractor/warehouse/bigquery/extract.py

@@ -2,6 +2,10 @@ import json
 import logging
 from typing import cast
 
+from castor_extractor.exceptions import (  # type: ignore
+    NoDatabaseProvidedException,
+)
+
 from ...logger import add_logging_file_handler
 from ...utils import LocalStorage, SafeMode, from_env, write_summary
 from ..abstract import (
@@ -62,6 +66,9 @@ def extract_all(**kwargs) -> None:
         db_blocked=kwargs.get("db_blocked"),
         dataset_blocked=kwargs.get("dataset_blocked"),
     )
+    projects = client.get_projects()
+    if not projects:
+        raise NoDatabaseProvidedException
 
     logger.info(f"Available projects: {client.get_projects()}\n")
 

{castor_extractor-0.25.6 → castor_extractor-0.25.9}/castor_extractor/warehouse/databricks/extract.py

@@ -2,6 +2,10 @@ import logging
 from datetime import date
 from typing import Optional
 
+from castor_extractor.exceptions import (  # type: ignore
+    NoDatabaseProvidedException,
+)
+
 from ...utils import AbstractStorage, LocalStorage, write_summary
 from ..abstract import (
     ADDITIONAL_LINEAGE_ASSETS,
@@ -177,6 +181,9 @@ def extract_all(**kwargs) -> None:
         db_blocked=kwargs.get("db_blocked"),
     )
 
+    if not client.databases():
+        raise NoDatabaseProvidedException
+
     storage = LocalStorage(directory=output_directory)
 
     extractor = DatabricksExtractionProcessor(

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

@@ -1,5 +1,9 @@
 import logging
 
+from castor_extractor.exceptions import (  # type: ignore
+    NoDatabaseProvidedException,
+)
+
 from ...utils import LocalStorage, from_env, write_summary
 from ..abstract import (
     CATALOG_ASSETS,
@@ -60,6 +64,14 @@ def _credentials(params: dict) -> dict:
     raise ValueError("missing password or private key")
 
 
+def _get_database_names(
+    extractor: SQLExtractionProcessor, query_builder: SnowflakeQueryBuilder
+) -> set[str]:
+    db_query = query_builder.build(WarehouseAsset.DATABASE)
+    databases = list(extractor.fetch(db_query[0]))
+    return {db["database_name"] for db in databases}
+
+
 def extract_all(**kwargs) -> None:
     """
     Extract all assets from Snowflake and store the results in CSV files
@@ -85,6 +97,12 @@ def extract_all(**kwargs) -> None:
         storage=storage,
     )
 
+    database_names = _get_database_names(extractor, query_builder)
+    if not database_names:
+        raise NoDatabaseProvidedException
+
+    logger.info(f"Available databases: {database_names}\n")
+
     for group in extractable_asset_groups(SNOWFLAKE_ASSETS):
         for asset in group:
             logger.info(f"Extracting `{asset.value.upper()}` ...")

{castor_extractor-0.25.6 → castor_extractor-0.25.9}/castor_extractor/warehouse/sqlserver/client.py

@@ -4,7 +4,8 @@ from typing import Optional
 
 from sqlalchemy import create_engine, text
 
-from ...utils import ExtractionQuery, SqlalchemyClient, uri_encode
+from ...utils import SqlalchemyClient, uri_encode
+from ..abstract import ExtractionQuery
 
 logger = logging.getLogger(__name__)
 
@@ -50,12 +51,25 @@ class MSSQLClient(SqlalchemyClient):
         )
         return uri
 
+    def _set_database(self, database: Optional[str]) -> None:
+        """
+        To support SQL Server running on Azure, we must change the engine's
+        URI depending on the case:
+        - URI/database for database-scoped queries
+        - URI otherwise
+        https://learn.microsoft.com/en-us/sql/t-sql/language-elements/use-transact-sql?view=sql-server-2017#database_name
+        https://github.com/dbeaver/dbeaver/issues/5258
+        """
+        database_uri = f"{self._uri}/{database}" if database else self._uri
+        self._engine = create_engine(database_uri, **self._options)
+
     def execute(self, query: ExtractionQuery) -> Iterator[dict]:
         """
         Re-implements the SQLAlchemyClient execute function to ensure we consume
         the cursor before calling connection.close() as it wipes out the data
         otherwise
         """
+        self._set_database(query.database)
         connection = self.connect()
         try:
             proxy = connection.execute(text(query.statement), query.params)
@@ -64,18 +78,6 @@ class MSSQLClient(SqlalchemyClient):
         finally:
             self.close()
 
-    def set_database(self, database: Optional[str]) -> None:
-        """
-        To support SQL Server running on Azure, we must change the engine's
-        URI for each database-scoped query
-        https://chatgpt.com/share/68de93c3-9550-8001-9d54-c5da86faa43c
-        """
-        if not database:
-            return
-        self.close()
-        database_uri = f"{self._uri}/{database}"
-        self._engine = create_engine(database_uri, **self._options)
-
     def get_databases(self) -> list[str]:
         result = self.execute(
             ExtractionQuery("SELECT name FROM sys.databases", {})
@@ -128,7 +130,11 @@ class MSSQLClient(SqlalchemyClient):
         FROM
             [{database}].sys.database_query_store_options
         """
-        query = ExtractionQuery(sql, {})
+        query = ExtractionQuery(
+            statement=sql,
+            params={},
+            database=database,
+        )
         # 2 = READ_WRITE, which means the Query Store is activated
         return next(self.execute(query))["desired_state"] == 2
 

{castor_extractor-0.25.6 → castor_extractor-0.25.9}/castor_extractor/warehouse/sqlserver/extract.py

@@ -32,7 +32,6 @@ MSSQL_USER = "CASTOR_MSSQL_USER"
 MSSQL_PASSWORD = "CASTOR_MSSQL_PASSWORD"  # noqa: S105
 MSSQL_HOST = "CASTOR_MSSQL_HOST"
 MSSQL_PORT = "CASTOR_MSSQL_PORT"
-MSSQL_DATABASE = "CASTOR_MSSQL_DATABASE"
 
 
 def _credentials(params: dict) -> dict:
@@ -43,7 +42,6 @@ def _credentials(params: dict) -> dict:
         "password": params.get("password") or from_env(MSSQL_PASSWORD),
         "host": params.get("host") or from_env(MSSQL_HOST),
         "port": params.get("port") or from_env(MSSQL_PORT),
-        "database": params.get("database") or from_env(MSSQL_DATABASE),
     }
 
 

{castor_extractor-0.25.6 → castor_extractor-0.25.9}/pyproject.toml

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

castor_extractor-0.25.6/castor_extractor/utils/client/query.py

@@ -1,6 +0,0 @@
-class ExtractionQuery:
-    """contains both statement and parameters for the query"""
-
-    def __init__(self, statement: str, params: dict):
-        self.statement = statement
-        self.params = params