duckguard 2.0.0__py3-none-any.whl

duckguard/connectors/__init__.py

@@ -0,0 +1,58 @@
+"""Connectors for various data sources."""
+
+from duckguard.connectors.base import Connector, ConnectionConfig
+from duckguard.connectors.files import FileConnector, S3Connector, GCSConnector, AzureConnector
+from duckguard.connectors.factory import connect, register_connector
+
+# Database connectors (imported lazily to avoid import errors)
+__all__ = [
+    # Base classes
+    "Connector",
+    "ConnectionConfig",
+    # File connectors
+    "FileConnector",
+    "S3Connector",
+    "GCSConnector",
+    "AzureConnector",
+    # Factory functions
+    "connect",
+    "register_connector",
+]
+
+
+def __getattr__(name: str):
+    """Lazy import database connectors to avoid import errors."""
+    if name == "PostgresConnector":
+        from duckguard.connectors.postgres import PostgresConnector
+        return PostgresConnector
+    if name == "MySQLConnector":
+        from duckguard.connectors.mysql import MySQLConnector
+        return MySQLConnector
+    if name == "SQLiteConnector":
+        from duckguard.connectors.sqlite import SQLiteConnector
+        return SQLiteConnector
+    if name == "SnowflakeConnector":
+        from duckguard.connectors.snowflake import SnowflakeConnector
+        return SnowflakeConnector
+    if name == "BigQueryConnector":
+        from duckguard.connectors.bigquery import BigQueryConnector
+        return BigQueryConnector
+    if name == "RedshiftConnector":
+        from duckguard.connectors.redshift import RedshiftConnector
+        return RedshiftConnector
+    if name == "SQLServerConnector":
+        from duckguard.connectors.sqlserver import SQLServerConnector
+        return SQLServerConnector
+    if name == "DatabricksConnector":
+        from duckguard.connectors.databricks import DatabricksConnector
+        return DatabricksConnector
+    if name == "OracleConnector":
+        from duckguard.connectors.oracle import OracleConnector
+        return OracleConnector
+    if name == "MongoDBConnector":
+        from duckguard.connectors.mongodb import MongoDBConnector
+        return MongoDBConnector
+    if name == "KafkaConnector":
+        from duckguard.connectors.kafka import KafkaConnector
+        return KafkaConnector
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

duckguard/connectors/base.py

@@ -0,0 +1,80 @@
+"""Base connector interface."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+from duckguard.core.dataset import Dataset
+from duckguard.core.engine import DuckGuardEngine
+
+
+@dataclass
+class ConnectionConfig:
+    """Configuration for a data source connection."""
+
+    source: str
+    table: str | None = None
+    schema: str | None = None
+    database: str | None = None
+    options: dict[str, Any] | None = None
+
+    def __post_init__(self) -> None:
+        if self.options is None:
+            self.options = {}
+
+
+class Connector(ABC):
+    """
+    Base class for data source connectors.
+
+    Connectors handle the logic of connecting to different data sources
+    and creating Dataset objects.
+    """
+
+    def __init__(self, engine: DuckGuardEngine | None = None):
+        """
+        Initialize the connector.
+
+        Args:
+            engine: Optional DuckGuardEngine instance
+        """
+        self.engine = engine or DuckGuardEngine.get_instance()
+
+    @abstractmethod
+    def connect(self, config: ConnectionConfig) -> Dataset:
+        """
+        Connect to a data source and return a Dataset.
+
+        Args:
+            config: Connection configuration
+
+        Returns:
+            Dataset object
+        """
+        pass
+
+    @classmethod
+    @abstractmethod
+    def can_handle(cls, source: str) -> bool:
+        """
+        Check if this connector can handle the given source.
+
+        Args:
+            source: Source string (path, URL, connection string)
+
+        Returns:
+            True if this connector can handle the source
+        """
+        pass
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """
+        Get the priority of this connector (higher = checked first).
+
+        Returns:
+            Priority value
+        """
+        return 0

duckguard/connectors/bigquery.py

@@ -0,0 +1,171 @@
+"""BigQuery connector."""
+
+from __future__ import annotations
+
+from typing import Any
+from urllib.parse import urlparse
+
+from duckguard.connectors.base import Connector, ConnectionConfig
+from duckguard.core.dataset import Dataset
+from duckguard.core.engine import DuckGuardEngine
+
+
+class BigQueryConnector(Connector):
+    """
+    Connector for Google BigQuery.
+
+    Uses the google-cloud-bigquery package to connect and query,
+    then processes results with DuckDB for validation.
+
+    Examples:
+        # Using connection string
+        data = connect(
+            "bigquery://project-id/dataset",
+            table="orders"
+        )
+
+        # Using options with service account
+        data = connect(
+            "bigquery://project-id",
+            table="orders",
+            dataset="my_dataset",
+            credentials_path="/path/to/service-account.json"
+        )
+    """
+
+    def __init__(self, engine: DuckGuardEngine | None = None):
+        super().__init__(engine)
+        self._client = None
+
+    def connect(self, config: ConnectionConfig) -> Dataset:
+        """
+        Connect to BigQuery and return a Dataset.
+
+        Args:
+            config: Connection configuration
+
+        Returns:
+            Dataset object
+        """
+        try:
+            from google.cloud import bigquery
+        except ImportError:
+            raise ImportError(
+                "BigQuery support requires google-cloud-bigquery. "
+                "Install with: pip install duckguard[bigquery]"
+            )
+
+        if not config.table:
+            raise ValueError("Table name is required for BigQuery connections")
+
+        # Parse connection parameters
+        conn_params = self._parse_connection_string(config.source, config)
+
+        # Initialize BigQuery client
+        if conn_params.get("credentials_path"):
+            self._client = bigquery.Client.from_service_account_json(
+                conn_params["credentials_path"]
+            )
+        else:
+            self._client = bigquery.Client(project=conn_params.get("project"))
+
+        table = config.table
+        dataset = conn_params.get("dataset", "")
+        project = conn_params.get("project", self._client.project)
+
+        # Build fully qualified table name
+        if project and dataset:
+            fq_table = f"`{project}.{dataset}.{table}`"
+        elif dataset:
+            fq_table = f"`{dataset}.{table}`"
+        else:
+            fq_table = f"`{table}`"
+
+        return BigQueryDataset(
+            source=fq_table,
+            engine=self.engine,
+            name=table,
+            client=self._client,
+        )
+
+    def _parse_connection_string(
+        self, conn_string: str, config: ConnectionConfig
+    ) -> dict[str, Any]:
+        """Parse BigQuery connection string and merge with config options."""
+        params: dict[str, Any] = {}
+
+        # Parse URL format: bigquery://project-id/dataset
+        if conn_string.lower().startswith("bigquery://"):
+            parsed = urlparse(conn_string)
+
+            params["project"] = parsed.hostname or ""
+
+            # Parse path for dataset
+            path_parts = [p for p in parsed.path.split("/") if p]
+            if len(path_parts) >= 1:
+                params["dataset"] = path_parts[0]
+
+        # Override with config options
+        options = config.options or {}
+        for key in ["project", "dataset", "credentials_path", "location"]:
+            if key in options:
+                params[key] = options[key]
+
+        if config.database:
+            params["dataset"] = config.database
+        if config.schema:
+            params["dataset"] = config.schema
+
+        return params
+
+    @classmethod
+    def can_handle(cls, source: str) -> bool:
+        """Check if this is a BigQuery connection string."""
+        return source.lower().startswith("bigquery://")
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """BigQuery connector has high priority."""
+        return 60
+
+
+class BigQueryDataset(Dataset):
+    """Dataset that queries BigQuery directly."""
+
+    def __init__(
+        self,
+        source: str,
+        engine: DuckGuardEngine,
+        name: str,
+        client: Any,
+    ):
+        super().__init__(source=source, engine=engine, name=name)
+        self._bq_client = client
+
+    def _execute_bq_query(self, sql: str) -> list[Any]:
+        """Execute a query on BigQuery."""
+        query_job = self._bq_client.query(sql)
+        return list(query_job.result())
+
+    def _fetch_bq_value(self, sql: str) -> Any:
+        """Execute query and return single value."""
+        rows = self._execute_bq_query(sql)
+        return rows[0][0] if rows else None
+
+    @property
+    def row_count(self) -> int:
+        """Get row count from BigQuery."""
+        if self._row_count_cache is None:
+            sql = f"SELECT COUNT(*) FROM {self._source}"
+            self._row_count_cache = self._fetch_bq_value(sql) or 0
+        return self._row_count_cache
+
+    @property
+    def columns(self) -> list[str]:
+        """Get column names from BigQuery."""
+        if self._columns_cache is None:
+            sql = f"SELECT * FROM {self._source} LIMIT 0"
+            query_job = self._bq_client.query(sql)
+            result = query_job.result()
+            self._columns_cache = [field.name for field in result.schema]
+        return self._columns_cache

duckguard/connectors/databricks.py

@@ -0,0 +1,201 @@
+"""Databricks connector."""
+
+from __future__ import annotations
+
+from typing import Any
+from urllib.parse import parse_qs, urlparse
+
+from duckguard.connectors.base import Connector, ConnectionConfig
+from duckguard.core.dataset import Dataset
+from duckguard.core.engine import DuckGuardEngine
+
+
+class DatabricksConnector(Connector):
+    """
+    Connector for Databricks SQL Warehouse and Unity Catalog.
+
+    Uses the databricks-sql-connector package for efficient querying.
+
+    Examples:
+        # Using connection string
+        data = connect(
+            "databricks://workspace.cloud.databricks.com/catalog/schema",
+            table="orders",
+            token="dapi..."
+        )
+
+        # Using options
+        data = connect(
+            "databricks://workspace.cloud.databricks.com",
+            table="orders",
+            catalog="main",
+            schema="default",
+            http_path="/sql/1.0/warehouses/abc123",
+            token="dapi..."
+        )
+    """
+
+    def __init__(self, engine: DuckGuardEngine | None = None):
+        super().__init__(engine)
+        self._connection = None
+
+    def connect(self, config: ConnectionConfig) -> Dataset:
+        """
+        Connect to Databricks and return a Dataset.
+
+        Args:
+            config: Connection configuration
+
+        Returns:
+            Dataset object
+        """
+        try:
+            from databricks import sql as databricks_sql
+        except ImportError:
+            raise ImportError(
+                "Databricks support requires databricks-sql-connector. "
+                "Install with: pip install duckguard[databricks]"
+            )
+
+        if not config.table:
+            raise ValueError("Table name is required for Databricks connections")
+
+        # Parse connection parameters
+        conn_params = self._parse_connection_string(config.source, config)
+
+        # Validate required parameters
+        if not conn_params.get("server_hostname"):
+            raise ValueError("Databricks server hostname is required")
+        if not conn_params.get("http_path"):
+            raise ValueError("Databricks http_path is required (SQL Warehouse path)")
+        if not conn_params.get("access_token"):
+            raise ValueError("Databricks access token is required")
+
+        # Connect to Databricks
+        self._connection = databricks_sql.connect(
+            server_hostname=conn_params["server_hostname"],
+            http_path=conn_params["http_path"],
+            access_token=conn_params["access_token"],
+        )
+
+        table = config.table
+        catalog = conn_params.get("catalog", "main")
+        schema = config.schema or conn_params.get("schema", "default")
+
+        # Build fully qualified table name
+        fq_table = f"`{catalog}`.`{schema}`.`{table}`"
+
+        return DatabricksDataset(
+            source=fq_table,
+            engine=self.engine,
+            name=table,
+            connection=self._connection,
+        )
+
+    def _parse_connection_string(
+        self, conn_string: str, config: ConnectionConfig
+    ) -> dict[str, Any]:
+        """Parse Databricks connection string and merge with config options."""
+        params: dict[str, Any] = {}
+
+        # Parse URL format: databricks://workspace.cloud.databricks.com/catalog/schema
+        if conn_string.lower().startswith("databricks://"):
+            parsed = urlparse(conn_string)
+
+            params["server_hostname"] = parsed.hostname or ""
+
+            # Parse path for catalog/schema
+            path_parts = [p for p in parsed.path.split("/") if p]
+            if len(path_parts) >= 1:
+                params["catalog"] = path_parts[0]
+            if len(path_parts) >= 2:
+                params["schema"] = path_parts[1]
+
+            # Parse query parameters
+            if parsed.query:
+                query_params = parse_qs(parsed.query)
+                for key, values in query_params.items():
+                    params[key] = values[0] if len(values) == 1 else values
+
+        # Override with config options
+        options = config.options or {}
+        for key in [
+            "server_hostname",
+            "http_path",
+            "access_token",
+            "token",
+            "catalog",
+            "schema",
+        ]:
+            if key in options:
+                # Handle token alias
+                if key == "token":
+                    params["access_token"] = options[key]
+                else:
+                    params[key] = options[key]
+
+        if config.database:
+            params["catalog"] = config.database
+        if config.schema:
+            params["schema"] = config.schema
+
+        return params
+
+    @classmethod
+    def can_handle(cls, source: str) -> bool:
+        """Check if this is a Databricks connection string."""
+        source_lower = source.lower()
+        return source_lower.startswith("databricks://") or ".databricks.com" in source_lower
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """Databricks connector has high priority."""
+        return 60
+
+
+class DatabricksDataset(Dataset):
+    """Dataset that queries Databricks directly."""
+
+    def __init__(
+        self,
+        source: str,
+        engine: DuckGuardEngine,
+        name: str,
+        connection: Any,
+    ):
+        super().__init__(source=source, engine=engine, name=name)
+        self._db_connection = connection
+
+    def _execute_query(self, sql: str) -> list[tuple[Any, ...]]:
+        """Execute a query on Databricks."""
+        cursor = self._db_connection.cursor()
+        try:
+            cursor.execute(sql)
+            return cursor.fetchall()
+        finally:
+            cursor.close()
+
+    def _fetch_value(self, sql: str) -> Any:
+        """Execute query and return single value."""
+        rows = self._execute_query(sql)
+        return rows[0][0] if rows else None
+
+    @property
+    def row_count(self) -> int:
+        """Get row count from Databricks."""
+        if self._row_count_cache is None:
+            sql = f"SELECT COUNT(*) FROM {self._source}"
+            self._row_count_cache = self._fetch_value(sql) or 0
+        return self._row_count_cache
+
+    @property
+    def columns(self) -> list[str]:
+        """Get column names from Databricks."""
+        if self._columns_cache is None:
+            cursor = self._db_connection.cursor()
+            try:
+                cursor.execute(f"SELECT * FROM {self._source} LIMIT 0")
+                self._columns_cache = [desc[0] for desc in cursor.description]
+            finally:
+                cursor.close()
+        return self._columns_cache