duckguard 2.0.0__py3-none-any.whl

duckguard/connectors/factory.py

@@ -0,0 +1,292 @@
+"""Factory function for creating connections."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from duckguard.connectors.base import Connector, ConnectionConfig
+from duckguard.connectors.files import FileConnector, S3Connector, GCSConnector, AzureConnector
+from duckguard.core.dataset import Dataset
+from duckguard.core.engine import DuckGuardEngine
+
+
+# Registry of available connectors
+_CONNECTORS: list[type[Connector]] = [
+    S3Connector,
+    GCSConnector,
+    AzureConnector,
+    FileConnector,
+]
+
+
+def register_connector(connector_class: type[Connector]) -> None:
+    """
+    Register a custom connector.
+
+    Args:
+        connector_class: Connector class to register
+    """
+    _CONNECTORS.append(connector_class)
+    # Sort by priority (highest first)
+    _CONNECTORS.sort(key=lambda c: c.get_priority(), reverse=True)
+
+
+def connect(
+    source: str,
+    *,
+    table: str | None = None,
+    schema: str | None = None,
+    database: str | None = None,
+    engine: DuckGuardEngine | None = None,
+    **options: Any,
+) -> Dataset:
+    """
+    Connect to a data source and return a Dataset.
+
+    This is the main entry point for connecting to data sources.
+    It automatically detects the source type and uses the appropriate connector.
+
+    Args:
+        source: Path to file, connection string, or URL
+        table: Table name (for database connections)
+        schema: Schema name (for database connections)
+        database: Database name (for database connections)
+        engine: Optional DuckGuardEngine instance
+        **options: Additional options passed to the connector
+
+    Returns:
+        Dataset object ready for validation
+
+    Examples:
+        # Connect to a CSV file
+        orders = connect("data/orders.csv")
+
+        # Connect to a Parquet file on S3
+        orders = connect("s3://bucket/orders.parquet")
+
+        # Connect to PostgreSQL
+        orders = connect("postgres://localhost/mydb", table="orders")
+
+        # Connect to Snowflake
+        orders = connect("snowflake://account/db", table="orders", schema="public")
+
+    Raises:
+        ValueError: If no connector can handle the source
+    """
+    config = ConnectionConfig(
+        source=source,
+        table=table,
+        schema=schema,
+        database=database,
+        options=options,
+    )
+
+    # Find a connector that can handle this source
+    for connector_class in _CONNECTORS:
+        if connector_class.can_handle(source):
+            connector = connector_class(engine=engine)
+            return connector.connect(config)
+
+    # Check for database connection strings
+    if _is_database_connection(source):
+        return _handle_database_connection(source, config, engine)
+
+    raise ValueError(
+        f"No connector found for source: {source}\n"
+        f"Supported formats: CSV, Parquet, JSON, Excel\n"
+        f"Supported protocols: s3://, gs://, az://, postgres://, mysql://"
+    )
+
+
+def _is_database_connection(source: str) -> bool:
+    """Check if source is a database connection string."""
+    db_prefixes = (
+        "postgres://",
+        "postgresql://",
+        "mysql://",
+        "mysql+pymysql://",
+        "sqlite://",
+        "snowflake://",
+        "bigquery://",
+        "redshift://",
+        "mssql://",
+        "sqlserver://",
+        "databricks://",
+        "oracle://",
+        "mongodb://",
+        "mongodb+srv://",
+        "kafka://",
+    )
+    source_lower = source.lower()
+
+    # Check prefixes
+    if source_lower.startswith(db_prefixes):
+        return True
+
+    # Check for SQLite file extensions
+    if source_lower.endswith((".db", ".sqlite", ".sqlite3")):
+        return True
+
+    # Check for Redshift hostname
+    if "redshift.amazonaws.com" in source_lower:
+        return True
+
+    # Check for Databricks hostname
+    if ".databricks.com" in source_lower:
+        return True
+
+    return False
+
+
+def _handle_database_connection(
+    source: str,
+    config: ConnectionConfig,
+    engine: DuckGuardEngine | None,
+) -> Dataset:
+    """Handle database connection strings."""
+    source_lower = source.lower()
+
+    # PostgreSQL
+    if source_lower.startswith(("postgres://", "postgresql://")):
+        try:
+            from duckguard.connectors.postgres import PostgresConnector
+
+            connector = PostgresConnector(engine=engine)
+            return connector.connect(config)
+        except ImportError:
+            raise ImportError(
+                "PostgreSQL support requires psycopg2. "
+                "Install with: pip install duckguard[postgres]"
+            )
+
+    # MySQL
+    if source_lower.startswith(("mysql://", "mysql+pymysql://")):
+        try:
+            from duckguard.connectors.mysql import MySQLConnector
+
+            connector = MySQLConnector(engine=engine)
+            return connector.connect(config)
+        except ImportError:
+            raise ImportError(
+                "MySQL support requires pymysql. "
+                "Install with: pip install duckguard[mysql]"
+            )
+
+    # SQLite
+    if source_lower.startswith("sqlite://") or source_lower.endswith(
+        (".db", ".sqlite", ".sqlite3")
+    ):
+        from duckguard.connectors.sqlite import SQLiteConnector
+
+        connector = SQLiteConnector(engine=engine)
+        return connector.connect(config)
+
+    # Snowflake
+    if source_lower.startswith("snowflake://"):
+        try:
+            from duckguard.connectors.snowflake import SnowflakeConnector
+
+            connector = SnowflakeConnector(engine=engine)
+            return connector.connect(config)
+        except ImportError:
+            raise ImportError(
+                "Snowflake support requires snowflake-connector-python. "
+                "Install with: pip install duckguard[snowflake]"
+            )
+
+    # BigQuery
+    if source_lower.startswith("bigquery://"):
+        try:
+            from duckguard.connectors.bigquery import BigQueryConnector
+
+            connector = BigQueryConnector(engine=engine)
+            return connector.connect(config)
+        except ImportError:
+            raise ImportError(
+                "BigQuery support requires google-cloud-bigquery. "
+                "Install with: pip install duckguard[bigquery]"
+            )
+
+    # Redshift
+    if source_lower.startswith("redshift://") or "redshift.amazonaws.com" in source_lower:
+        from duckguard.connectors.redshift import RedshiftConnector
+
+        connector = RedshiftConnector(engine=engine)
+        return connector.connect(config)
+
+    # SQL Server
+    if source_lower.startswith(("mssql://", "sqlserver://", "mssql+pyodbc://")):
+        try:
+            from duckguard.connectors.sqlserver import SQLServerConnector
+
+            connector = SQLServerConnector(engine=engine)
+            return connector.connect(config)
+        except ImportError:
+            raise ImportError(
+                "SQL Server support requires pyodbc or pymssql. "
+                "Install with: pip install duckguard[sqlserver]"
+            )
+
+    # Databricks
+    if source_lower.startswith("databricks://") or ".databricks.com" in source_lower:
+        try:
+            from duckguard.connectors.databricks import DatabricksConnector
+
+            connector = DatabricksConnector(engine=engine)
+            return connector.connect(config)
+        except ImportError:
+            raise ImportError(
+                "Databricks support requires databricks-sql-connector. "
+                "Install with: pip install duckguard[databricks]"
+            )
+
+    # Oracle
+    if source_lower.startswith("oracle://"):
+        try:
+            from duckguard.connectors.oracle import OracleConnector
+
+            connector = OracleConnector(engine=engine)
+            return connector.connect(config)
+        except ImportError:
+            raise ImportError(
+                "Oracle support requires oracledb. "
+                "Install with: pip install duckguard[oracle]"
+            )
+
+    # MongoDB
+    if source_lower.startswith(("mongodb://", "mongodb+srv://")):
+        try:
+            from duckguard.connectors.mongodb import MongoDBConnector
+
+            connector = MongoDBConnector(engine=engine)
+            return connector.connect(config)
+        except ImportError:
+            raise ImportError(
+                "MongoDB support requires pymongo. "
+                "Install with: pip install duckguard[mongodb]"
+            )
+
+    # Kafka
+    if source_lower.startswith("kafka://"):
+        try:
+            from duckguard.connectors.kafka import KafkaConnector
+
+            connector = KafkaConnector(engine=engine)
+            return connector.connect(config)
+        except ImportError:
+            raise ImportError(
+                "Kafka support requires kafka-python. "
+                "Install with: pip install duckguard[kafka]"
+            )
+
+    # For other databases, raise helpful error
+    raise ValueError(
+        f"Database connector not yet implemented for: {source}\n"
+        f"Currently supported: postgres://, mysql://, sqlite://, snowflake://, "
+        f"bigquery://, redshift://, mssql://, databricks://, oracle://, "
+        f"mongodb://, kafka://"
+    )
+
+
+# Alias for backwards compatibility
+load = connect

duckguard/connectors/files.py

@@ -0,0 +1,135 @@
+"""File-based connectors (CSV, Parquet, JSON, Excel)."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from duckguard.connectors.base import Connector, ConnectionConfig
+from duckguard.core.dataset import Dataset
+from duckguard.core.engine import DuckGuardEngine
+
+
+class FileConnector(Connector):
+    """
+    Connector for file-based data sources.
+
+    Supports:
+    - CSV files (.csv)
+    - Parquet files (.parquet, .pq)
+    - JSON files (.json, .jsonl, .ndjson)
+    - Excel files (.xlsx, .xls) - requires additional setup
+    """
+
+    SUPPORTED_EXTENSIONS = {
+        ".csv": "csv",
+        ".parquet": "parquet",
+        ".pq": "parquet",
+        ".json": "json",
+        ".jsonl": "json",
+        ".ndjson": "json",
+        ".xlsx": "excel",
+        ".xls": "excel",
+    }
+
+    def __init__(self, engine: DuckGuardEngine | None = None):
+        super().__init__(engine)
+
+    def connect(self, config: ConnectionConfig) -> Dataset:
+        """
+        Connect to a file and return a Dataset.
+
+        Args:
+            config: Connection configuration with file path
+
+        Returns:
+            Dataset object
+        """
+        path = config.source
+        ext = Path(path).suffix.lower()
+
+        # Determine file type
+        file_type = self.SUPPORTED_EXTENSIONS.get(ext)
+        if not file_type:
+            raise ValueError(f"Unsupported file type: {ext}")
+
+        # Validate file exists (for local files)
+        if not self._is_remote_path(path) and not os.path.exists(path):
+            raise FileNotFoundError(f"File not found: {path}")
+
+        # Create dataset name from filename
+        name = Path(path).stem
+
+        return Dataset(source=path, engine=self.engine, name=name)
+
+    @classmethod
+    def can_handle(cls, source: str) -> bool:
+        """Check if this connector can handle the source."""
+        # Check for file extensions
+        path = Path(source)
+        ext = path.suffix.lower()
+
+        if ext in cls.SUPPORTED_EXTENSIONS:
+            return True
+
+        # Check for S3/GCS/Azure paths with supported extensions
+        if cls._is_remote_path(source):
+            # Extract extension from remote path
+            for supported_ext in cls.SUPPORTED_EXTENSIONS:
+                if source.lower().endswith(supported_ext):
+                    return True
+
+        return False
+
+    @staticmethod
+    def _is_remote_path(path: str) -> bool:
+        """Check if path is a remote storage path."""
+        remote_prefixes = ("s3://", "gs://", "gcs://", "az://", "abfs://", "http://", "https://")
+        return path.lower().startswith(remote_prefixes)
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """File connector has default priority."""
+        return 10
+
+
+class S3Connector(FileConnector):
+    """Connector for S3 paths."""
+
+    @classmethod
+    def can_handle(cls, source: str) -> bool:
+        """Check if this is an S3 path."""
+        return source.lower().startswith("s3://")
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """S3 connector has higher priority."""
+        return 20
+
+
+class GCSConnector(FileConnector):
+    """Connector for Google Cloud Storage paths."""
+
+    @classmethod
+    def can_handle(cls, source: str) -> bool:
+        """Check if this is a GCS path."""
+        return source.lower().startswith(("gs://", "gcs://"))
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """GCS connector has higher priority."""
+        return 20
+
+
+class AzureConnector(FileConnector):
+    """Connector for Azure Blob Storage paths."""
+
+    @classmethod
+    def can_handle(cls, source: str) -> bool:
+        """Check if this is an Azure path."""
+        return source.lower().startswith(("az://", "abfs://"))
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """Azure connector has higher priority."""
+        return 20