duckguard 2.0.0__py3-none-any.whl

duckguard/connectors/oracle.py

@@ -0,0 +1,196 @@
+"""Oracle Database 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 OracleConnector(Connector):
+    """
+    Connector for Oracle Database.
+
+    Uses the oracledb (python-oracledb) package for connectivity.
+
+    Examples:
+        # Using connection string
+        data = connect(
+            "oracle://user:pass@host:1521/service_name",
+            table="orders"
+        )
+
+        # Using TNS alias
+        data = connect(
+            "oracle://user:pass@tns_alias",
+            table="orders"
+        )
+
+        # Using options
+        data = connect(
+            "oracle://host:1521/service_name",
+            table="orders",
+            user="myuser",
+            password="mypass",
+            schema="HR"
+        )
+    """
+
+    def __init__(self, engine: DuckGuardEngine | None = None):
+        super().__init__(engine)
+        self._connection = None
+
+    def connect(self, config: ConnectionConfig) -> Dataset:
+        """
+        Connect to Oracle and return a Dataset.
+
+        Args:
+            config: Connection configuration
+
+        Returns:
+            Dataset object
+        """
+        try:
+            import oracledb
+        except ImportError:
+            raise ImportError(
+                "Oracle support requires oracledb. "
+                "Install with: pip install duckguard[oracle]"
+            )
+
+        if not config.table:
+            raise ValueError("Table name is required for Oracle connections")
+
+        # Parse connection parameters
+        conn_params = self._parse_connection_string(config.source, config)
+
+        # Build connection
+        if conn_params.get("dsn"):
+            # Using DSN/TNS
+            self._connection = oracledb.connect(
+                user=conn_params.get("user"),
+                password=conn_params.get("password"),
+                dsn=conn_params["dsn"],
+            )
+        else:
+            # Using host/port/service
+            self._connection = oracledb.connect(
+                user=conn_params.get("user"),
+                password=conn_params.get("password"),
+                host=conn_params.get("host", "localhost"),
+                port=int(conn_params.get("port", 1521)),
+                service_name=conn_params.get("service_name"),
+            )
+
+        table = config.table
+        schema = config.schema or conn_params.get("schema", conn_params.get("user", "").upper())
+
+        # Build fully qualified table name
+        if schema:
+            fq_table = f'"{schema}"."{table.upper()}"'
+        else:
+            fq_table = f'"{table.upper()}"'
+
+        return OracleDataset(
+            source=fq_table,
+            engine=self.engine,
+            name=table,
+            connection=self._connection,
+        )
+
+    def _parse_connection_string(self, conn_string: str, config: ConnectionConfig) -> dict:
+        """Parse Oracle connection string."""
+        params: dict[str, Any] = {}
+
+        # Parse URL format: oracle://user:pass@host:port/service_name
+        if conn_string.lower().startswith("oracle://"):
+            parsed = urlparse(conn_string)
+
+            params["user"] = parsed.username or ""
+            params["password"] = parsed.password or ""
+            params["host"] = parsed.hostname or "localhost"
+            params["port"] = str(parsed.port) if parsed.port else "1521"
+
+            # Path is service name or SID
+            if parsed.path:
+                service = parsed.path.lstrip("/")
+                if service:
+                    params["service_name"] = service
+
+            # Check if it's a TNS alias (no port specified and no dots in hostname)
+            if not parsed.port and parsed.hostname and "." not in parsed.hostname:
+                params["dsn"] = parsed.hostname
+
+        # Override with config options
+        options = config.options or {}
+        for key in ["user", "password", "host", "port", "service_name", "dsn", "schema"]:
+            if key in options:
+                params[key] = options[key]
+
+        if config.database:
+            params["service_name"] = config.database
+        if config.schema:
+            params["schema"] = config.schema
+
+        return params
+
+    @classmethod
+    def can_handle(cls, source: str) -> bool:
+        """Check if this is an Oracle connection string."""
+        return source.lower().startswith("oracle://")
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """Oracle connector has high priority."""
+        return 55
+
+
+class OracleDataset(Dataset):
+    """Dataset that queries Oracle directly."""
+
+    def __init__(
+        self,
+        source: str,
+        engine: DuckGuardEngine,
+        name: str,
+        connection: Any,
+    ):
+        super().__init__(source=source, engine=engine, name=name)
+        self._ora_connection = connection
+
+    def _execute_query(self, sql: str) -> list[tuple[Any, ...]]:
+        """Execute a query on Oracle."""
+        cursor = self._ora_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 Oracle."""
+        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 Oracle."""
+        if self._columns_cache is None:
+            cursor = self._ora_connection.cursor()
+            try:
+                cursor.execute(f"SELECT * FROM {self._source} WHERE ROWNUM = 0")
+                self._columns_cache = [desc[0] for desc in cursor.description]
+            finally:
+                cursor.close()
+        return self._columns_cache

duckguard/connectors/postgres.py

@@ -0,0 +1,99 @@
+"""PostgreSQL connector."""
+
+from __future__ import annotations
+
+import re
+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 PostgresConnector(Connector):
+    """
+    Connector for PostgreSQL databases.
+
+    Uses DuckDB's postgres extension for efficient query pushdown.
+    """
+
+    def __init__(self, engine: DuckGuardEngine | None = None):
+        super().__init__(engine)
+        self._setup_extension()
+
+    def _setup_extension(self) -> None:
+        """Install and load the postgres extension."""
+        try:
+            self.engine.execute("INSTALL postgres")
+            self.engine.execute("LOAD postgres")
+        except Exception:
+            # Extension might already be loaded
+            pass
+
+    def connect(self, config: ConnectionConfig) -> Dataset:
+        """
+        Connect to PostgreSQL and return a Dataset.
+
+        Args:
+            config: Connection configuration
+
+        Returns:
+            Dataset object
+        """
+        if not config.table:
+            raise ValueError("Table name is required for PostgreSQL connections")
+
+        # Parse connection string
+        conn_info = self._parse_connection_string(config.source)
+
+        # Build the full table reference
+        schema = config.schema or conn_info.get("schema", "public")
+        table = config.table
+        full_table = f"{schema}.{table}"
+
+        # Create a unique alias for this connection
+        alias = f"pg_{table}"
+
+        # Attach the database
+        attach_sql = f"ATTACH '{config.source}' AS {alias} (TYPE postgres)"
+
+        try:
+            self.engine.execute(attach_sql)
+        except Exception as e:
+            if "already exists" not in str(e).lower():
+                raise
+
+        # The source reference for DuckDB
+        source_ref = f"{alias}.{full_table}"
+
+        # Register as a view for easier access
+        view_name = f"_duckguard_{table}"
+        try:
+            self.engine.execute(f"CREATE OR REPLACE VIEW {view_name} AS SELECT * FROM {source_ref}")
+        except Exception:
+            pass
+
+        return Dataset(source=view_name, engine=self.engine, name=table)
+
+    def _parse_connection_string(self, conn_string: str) -> dict[str, str]:
+        """Parse PostgreSQL connection string."""
+        parsed = urlparse(conn_string)
+
+        return {
+            "host": parsed.hostname or "localhost",
+            "port": str(parsed.port or 5432),
+            "database": parsed.path.lstrip("/") if parsed.path else "",
+            "user": parsed.username or "",
+            "password": parsed.password or "",
+            "schema": "public",
+        }
+
+    @classmethod
+    def can_handle(cls, source: str) -> bool:
+        """Check if this is a PostgreSQL connection string."""
+        return source.lower().startswith(("postgres://", "postgresql://"))
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """Database connectors have high priority."""
+        return 50

duckguard/connectors/redshift.py

@@ -0,0 +1,154 @@
+"""Amazon Redshift connector."""
+
+from __future__ import annotations
+
+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 RedshiftConnector(Connector):
+    """
+    Connector for Amazon Redshift.
+
+    Redshift is PostgreSQL-compatible, so we can use the PostgreSQL
+    extension in DuckDB or the redshift_connector package.
+
+    Examples:
+        # Using connection string
+        data = connect(
+            "redshift://user:pass@cluster.region.redshift.amazonaws.com:5439/database",
+            table="orders"
+        )
+
+        # Using options
+        data = connect(
+            "redshift://cluster.region.redshift.amazonaws.com:5439/database",
+            table="orders",
+            user="myuser",
+            password="mypass",
+            schema="public"
+        )
+    """
+
+    def __init__(self, engine: DuckGuardEngine | None = None):
+        super().__init__(engine)
+        self._setup_extension()
+
+    def _setup_extension(self) -> None:
+        """Install and load the postgres extension (Redshift compatible)."""
+        try:
+            self.engine.execute("INSTALL postgres")
+            self.engine.execute("LOAD postgres")
+        except Exception:
+            pass
+
+    def connect(self, config: ConnectionConfig) -> Dataset:
+        """
+        Connect to Redshift and return a Dataset.
+
+        Args:
+            config: Connection configuration
+
+        Returns:
+            Dataset object
+        """
+        if not config.table:
+            raise ValueError("Table name is required for Redshift connections")
+
+        # Parse connection string
+        conn_info = self._parse_connection_string(config.source, config)
+
+        table = config.table
+        schema = config.schema or conn_info.get("schema", "public")
+
+        # Create a unique alias
+        alias = f"redshift_{table}"
+
+        # Build PostgreSQL-compatible connection string for DuckDB
+        pg_conn = self._build_connection_string(conn_info)
+
+        # Attach using PostgreSQL extension (Redshift is PG-compatible)
+        attach_sql = f"ATTACH '{pg_conn}' AS {alias} (TYPE postgres)"
+
+        try:
+            self.engine.execute(attach_sql)
+        except Exception as e:
+            if "already exists" not in str(e).lower():
+                raise
+
+        # Build source reference
+        source_ref = f"{alias}.{schema}.{table}"
+
+        # Register as a view
+        view_name = f"_duckguard_redshift_{table}"
+        try:
+            self.engine.execute(f"CREATE OR REPLACE VIEW {view_name} AS SELECT * FROM {source_ref}")
+        except Exception:
+            pass
+
+        return Dataset(source=view_name, engine=self.engine, name=table)
+
+    def _parse_connection_string(self, conn_string: str, config: ConnectionConfig) -> dict:
+        """Parse Redshift connection string."""
+        # Handle redshift:// prefix
+        if conn_string.lower().startswith("redshift://"):
+            conn_string = "postgresql://" + conn_string[11:]
+
+        parsed = urlparse(conn_string)
+
+        params = {
+            "host": parsed.hostname or "",
+            "port": str(parsed.port or 5439),
+            "database": parsed.path.lstrip("/") if parsed.path else "",
+            "user": parsed.username or "",
+            "password": parsed.password or "",
+        }
+
+        # Override with config options
+        options = config.options or {}
+        for key in ["user", "password", "host", "port", "database", "schema", "sslmode"]:
+            if key in options:
+                params[key] = options[key]
+
+        if config.database:
+            params["database"] = config.database
+        if config.schema:
+            params["schema"] = config.schema
+
+        return params
+
+    def _build_connection_string(self, conn_info: dict) -> str:
+        """Build connection string for DuckDB PostgreSQL extension."""
+        parts = []
+
+        if conn_info.get("host"):
+            parts.append(f"host={conn_info['host']}")
+        if conn_info.get("port"):
+            parts.append(f"port={conn_info['port']}")
+        if conn_info.get("user"):
+            parts.append(f"user={conn_info['user']}")
+        if conn_info.get("password"):
+            parts.append(f"password={conn_info['password']}")
+        if conn_info.get("database"):
+            parts.append(f"dbname={conn_info['database']}")
+
+        # Redshift requires SSL
+        parts.append("sslmode=require")
+
+        return " ".join(parts)
+
+    @classmethod
+    def can_handle(cls, source: str) -> bool:
+        """Check if this is a Redshift connection string."""
+        source_lower = source.lower()
+        return source_lower.startswith("redshift://") or (
+            "redshift.amazonaws.com" in source_lower
+        )
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """Redshift connector has high priority."""
+        return 55

duckguard/connectors/snowflake.py

@@ -0,0 +1,226 @@
+"""Snowflake connector."""
+
+from __future__ import annotations
+
+import re
+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 SnowflakeConnector(Connector):
+    """
+    Connector for Snowflake data warehouse.
+
+    Uses the snowflake-connector-python package to connect and query,
+    then processes results with DuckDB for validation.
+
+    Examples:
+        # Using connection string
+        data = connect(
+            "snowflake://user:pass@account/database/schema",
+            table="orders"
+        )
+
+        # Using options
+        data = connect(
+            "snowflake://account",
+            table="orders",
+            user="myuser",
+            password="mypass",
+            database="mydb",
+            schema="public",
+            warehouse="compute_wh"
+        )
+    """
+
+    def __init__(self, engine: DuckGuardEngine | None = None):
+        super().__init__(engine)
+        self._connection = None
+
+    def connect(self, config: ConnectionConfig) -> Dataset:
+        """
+        Connect to Snowflake and return a Dataset.
+
+        Args:
+            config: Connection configuration
+
+        Returns:
+            Dataset object
+        """
+        try:
+            import snowflake.connector
+        except ImportError:
+            raise ImportError(
+                "Snowflake support requires snowflake-connector-python. "
+                "Install with: pip install duckguard[snowflake]"
+            )
+
+        if not config.table:
+            raise ValueError("Table name is required for Snowflake connections")
+
+        # Parse connection parameters
+        conn_params = self._parse_connection_string(config.source, config)
+
+        # Connect to Snowflake
+        self._connection = snowflake.connector.connect(**conn_params)
+
+        table = config.table
+        schema = config.schema or conn_params.get("schema", "PUBLIC")
+        database = config.database or conn_params.get("database", "")
+
+        # Build fully qualified table name
+        if database and schema:
+            fq_table = f"{database}.{schema}.{table}"
+        elif schema:
+            fq_table = f"{schema}.{table}"
+        else:
+            fq_table = table
+
+        # Create a wrapper dataset that uses Snowflake for queries
+        return SnowflakeDataset(
+            source=fq_table,
+            engine=self.engine,
+            name=table,
+            connection=self._connection,
+            conn_params=conn_params,
+        )
+
+    def _parse_connection_string(
+        self, conn_string: str, config: ConnectionConfig
+    ) -> dict[str, Any]:
+        """Parse Snowflake connection string and merge with config options."""
+        params: dict[str, Any] = {}
+
+        # Parse URL format: snowflake://user:pass@account/database/schema
+        if conn_string.lower().startswith("snowflake://"):
+            parsed = urlparse(conn_string)
+
+            params["account"] = parsed.hostname or ""
+            if parsed.username:
+                params["user"] = parsed.username
+            if parsed.password:
+                params["password"] = parsed.password
+
+            # Parse path for database/schema
+            path_parts = [p for p in parsed.path.split("/") if p]
+            if len(path_parts) >= 1:
+                params["database"] = 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 ["user", "password", "account", "warehouse", "role", "database", "schema"]:
+            if key in options:
+                params[key] = options[key]
+
+        if config.database:
+            params["database"] = config.database
+        if config.schema:
+            params["schema"] = config.schema
+
+        return params
+
+    @classmethod
+    def can_handle(cls, source: str) -> bool:
+        """Check if this is a Snowflake connection string."""
+        return source.lower().startswith("snowflake://")
+
+    @classmethod
+    def get_priority(cls) -> int:
+        """Snowflake connector has high priority."""
+        return 60
+
+
+class SnowflakeDataset(Dataset):
+    """
+    Dataset that queries Snowflake directly for statistics.
+
+    Uses query pushdown to compute aggregations in Snowflake,
+    minimizing data transfer.
+    """
+
+    def __init__(
+        self,
+        source: str,
+        engine: DuckGuardEngine,
+        name: str,
+        connection: Any,
+        conn_params: dict[str, Any],
+    ):
+        super().__init__(source=source, engine=engine, name=name)
+        self._sf_connection = connection
+        self._sf_params = conn_params
+
+    def _execute_sf_query(self, sql: str) -> list[tuple[Any, ...]]:
+        """Execute a query on Snowflake."""
+        cursor = self._sf_connection.cursor()
+        try:
+            cursor.execute(sql)
+            return cursor.fetchall()
+        finally:
+            cursor.close()
+
+    def _fetch_sf_value(self, sql: str) -> Any:
+        """Execute query and return single value."""
+        rows = self._execute_sf_query(sql)
+        return rows[0][0] if rows else None
+
+    @property
+    def row_count(self) -> int:
+        """Get row count from Snowflake."""
+        if self._row_count_cache is None:
+            sql = f"SELECT COUNT(*) FROM {self._source}"
+            self._row_count_cache = self._fetch_sf_value(sql) or 0
+        return self._row_count_cache
+
+    @property
+    def columns(self) -> list[str]:
+        """Get column names from Snowflake."""
+        if self._columns_cache is None:
+            sql = f"SELECT * FROM {self._source} LIMIT 0"
+            cursor = self._sf_connection.cursor()
+            try:
+                cursor.execute(sql)
+                self._columns_cache = [desc[0] for desc in cursor.description]
+            finally:
+                cursor.close()
+        return self._columns_cache
+
+
+class SnowflakeColumn:
+    """Column that queries Snowflake directly."""
+
+    def __init__(self, name: str, dataset: SnowflakeDataset):
+        self._name = name
+        self._dataset = dataset
+
+    @property
+    def null_percent(self) -> float:
+        """Get null percentage from Snowflake."""
+        sql = f"""
+        SELECT
+            ROUND(100.0 * SUM(CASE WHEN "{self._name}" IS NULL THEN 1 ELSE 0 END) / COUNT(*), 2)
+        FROM {self._dataset._source}
+        """
+        return self._dataset._fetch_sf_value(sql) or 0.0
+
+    @property
+    def unique_percent(self) -> float:
+        """Get unique percentage from Snowflake."""
+        sql = f"""
+        SELECT
+            ROUND(100.0 * COUNT(DISTINCT "{self._name}") / COUNT(*), 2)
+        FROM {self._dataset._source}
+        """
+        return self._dataset._fetch_sf_value(sql) or 0.0