kontra 0.5.2__py3-none-any.whl

kontra/engine/executors/postgres_sql.py

@@ -0,0 +1,162 @@
+# src/kontra/engine/executors/postgres_sql.py
+"""
+PostgreSQL SQL Executor - executes validation rules via SQL pushdown.
+
+Uses DatabaseSqlExecutor base class for shared compile/execute logic.
+
+Supports rules:
+  - not_null(column) - uses EXISTS (fast early termination)
+  - unique(column) - uses COUNT DISTINCT
+  - min_rows(threshold) - uses COUNT
+  - max_rows(threshold) - uses COUNT
+  - allowed_values(column, values) - uses SUM CASE
+  - freshness(column, max_age) - uses MAX
+  - range(column, min, max) - uses SUM CASE
+  - regex(column, pattern) - uses ~ operator
+  - compare(left, right, op) - uses SUM CASE
+  - conditional_not_null(column, when) - uses SUM CASE
+  - conditional_range(column, when, min, max) - uses SUM CASE
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Any, Dict, List, Tuple
+
+from kontra.connectors.handle import DatasetHandle
+from kontra.connectors.postgres import PostgresConnectionParams, get_connection
+from kontra.connectors.detection import parse_table_reference, get_default_schema, POSTGRESQL
+
+from .database_base import DatabaseSqlExecutor
+from .registry import register_executor
+
+
+@register_executor("postgres")
+class PostgresSqlExecutor(DatabaseSqlExecutor):
+    """
+    PostgreSQL SQL pushdown executor.
+
+    Inherits compile() and execute() from DatabaseSqlExecutor.
+    Provides PostgreSQL-specific connection and table handling.
+    """
+
+    DIALECT = "postgres"
+    SUPPORTED_RULES = {
+        "not_null", "unique", "min_rows", "max_rows",
+        "allowed_values", "disallowed_values",
+        "freshness", "range", "length",
+        "regex", "contains", "starts_with", "ends_with",
+        "compare", "conditional_not_null", "conditional_range",
+        "custom_sql_check", "custom_agg"
+    }
+
+    @property
+    def name(self) -> str:
+        return "postgres"
+
+    def _supports_scheme(self, scheme: str, handle: DatasetHandle) -> bool:
+        """Check if this executor supports the given URI scheme."""
+        # BYOC: check dialect for external connections
+        if scheme == "byoc" and handle.dialect == "postgresql":
+            return handle.external_conn is not None
+
+        # URI-based: handle postgres:// URIs
+        return scheme in {"postgres", "postgresql"}
+
+    @contextmanager
+    def _get_connection_ctx(self, handle: DatasetHandle):
+        """
+        Get a PostgreSQL connection context.
+
+        For BYOC, yields the external connection directly (not owned by us).
+        For URI-based, yields a new connection (owned by context manager).
+        """
+        if handle.scheme == "byoc" and handle.external_conn is not None:
+            yield handle.external_conn
+        elif handle.db_params:
+            with get_connection(handle.db_params) as conn:
+                yield conn
+        else:
+            raise ValueError("Handle has neither external_conn nor db_params")
+
+    def _get_table_reference(self, handle: DatasetHandle) -> str:
+        """
+        Get the fully-qualified table reference for PostgreSQL.
+
+        Returns: "schema"."table" format.
+        """
+        if handle.scheme == "byoc" and handle.table_ref:
+            _db, schema, table = parse_table_reference(handle.table_ref)
+            schema = schema or get_default_schema(POSTGRESQL)
+            return f"{self._esc(schema)}.{self._esc(table)}"
+        elif handle.db_params:
+            params: PostgresConnectionParams = handle.db_params
+            return f"{self._esc(params.schema)}.{self._esc(params.table)}"
+        else:
+            raise ValueError("Handle has neither table_ref nor db_params")
+
+    def _get_schema_and_table(self, handle: DatasetHandle) -> Tuple[str, str]:
+        """
+        Get schema and table name separately for custom SQL placeholder replacement.
+
+        Returns: Tuple of (schema, table_name)
+        """
+        if handle.scheme == "byoc" and handle.table_ref:
+            _db, schema, table = parse_table_reference(handle.table_ref)
+            schema = schema or get_default_schema(POSTGRESQL)
+            return schema, table
+        elif handle.db_params:
+            params: PostgresConnectionParams = handle.db_params
+            return params.schema, params.table
+        else:
+            raise ValueError("Handle has neither table_ref nor db_params")
+
+    def _get_cursor(self, conn):
+        """Get a cursor using context manager."""
+        # PostgreSQL cursors support context managers
+        return conn.cursor()
+
+    def _close_cursor(self, cursor):
+        """Close the cursor."""
+        cursor.close()
+
+    def introspect(self, handle: DatasetHandle, **kwargs) -> Dict[str, Any]:
+        """
+        Introspect the PostgreSQL table for metadata.
+
+        Returns:
+            {"row_count": int, "available_cols": [...], "staging": None}
+        """
+        table = self._get_table_reference(handle)
+
+        # Get schema and table name for information_schema query
+        if handle.scheme == "byoc" and handle.table_ref:
+            _db, schema, table_name = parse_table_reference(handle.table_ref)
+            schema = schema or get_default_schema(POSTGRESQL)
+        elif handle.db_params:
+            params: PostgresConnectionParams = handle.db_params
+            schema = params.schema
+            table_name = params.table
+        else:
+            raise ValueError("Handle has neither table_ref nor db_params")
+
+        with self._get_connection_ctx(handle) as conn:
+            with conn.cursor() as cur:
+                # Get row count
+                cur.execute(f"SELECT COUNT(*) FROM {table}")
+                row_count = cur.fetchone()
+                n = int(row_count[0]) if row_count else 0
+
+                # Get column names
+                cur.execute(
+                    """
+                    SELECT column_name
+                    FROM information_schema.columns
+                    WHERE table_schema = %s AND table_name = %s
+                    ORDER BY ordinal_position
+                    """,
+                    (schema, table_name),
+                )
+                cols = [row[0] for row in cur.fetchall()]
+
+        return {"row_count": n, "available_cols": cols, "staging": None}

kontra/engine/executors/registry.py

@@ -0,0 +1,69 @@
+# src/kontra/engine/executors/registry.py
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+
+from kontra.connectors.handle import DatasetHandle
+
+if TYPE_CHECKING:
+    from .base import SqlExecutor
+
+
+# Global registry: maps "executor_name" -> constructor()
+_EXECUTORS: Dict[str, Callable[[], SqlExecutor]] = {}
+
+
+def register_executor(name: str):
+    """
+    Decorator to register a SQL executor class.
+    """
+
+    def deco(cls: Callable[[], SqlExecutor]) -> Callable[[], SqlExecutor]:
+        if name in _EXECUTORS:
+            raise ValueError(f"Executor '{name}' is already registered.")
+        _EXECUTORS[name] = cls
+        return cls
+
+    return deco
+
+
+def pick_executor(
+    handle: DatasetHandle, sql_specs: List[Dict[str, Any]]
+) -> Optional[SqlExecutor]:
+    """
+    Find the first registered executor that supports the given handle and rules.
+    """
+    if not sql_specs:
+        return None  # Nothing to push down
+
+    # Iterate over registered constructors
+    for name, ctor in _EXECUTORS.items():
+        executor = ctor()  # Instantiate the executor
+        try:
+            if executor.supports(handle, sql_specs):
+                return executor
+        except Exception:
+            # Be conservative: ignore faulty executors
+            continue
+    return None
+
+
+def register_default_executors() -> None:
+    """
+    Eagerly import built-in executors so their @register_executor
+    decorators run and populate the registry.
+    """
+    # Local import triggers decorator side-effect
+    from . import duckdb_sql  # noqa: F401
+
+    # PostgreSQL executor (optional - requires psycopg)
+    try:
+        from . import postgres_sql  # noqa: F401
+    except ImportError:
+        pass  # psycopg not installed, skip postgres executor
+
+    # SQL Server executor (optional - requires pymssql)
+    try:
+        from . import sqlserver_sql  # noqa: F401
+    except ImportError:
+        pass  # pymssql not installed, skip sqlserver executor

kontra/engine/executors/sqlserver_sql.py

@@ -0,0 +1,163 @@
+# src/kontra/engine/executors/sqlserver_sql.py
+"""
+SQL Server SQL Executor - executes validation rules via SQL pushdown.
+
+Uses DatabaseSqlExecutor base class for shared compile/execute logic.
+
+Supports rules:
+  - not_null(column) - uses EXISTS (fast early termination)
+  - unique(column) - uses COUNT DISTINCT
+  - min_rows(threshold) - uses COUNT
+  - max_rows(threshold) - uses COUNT
+  - allowed_values(column, values) - uses SUM CASE
+  - freshness(column, max_age) - uses MAX
+  - range(column, min, max) - uses SUM CASE
+  - compare(left, right, op) - uses SUM CASE
+  - conditional_not_null(column, when) - uses SUM CASE
+  - conditional_range(column, when, min, max) - uses SUM CASE
+
+NOT supported (falls back to Polars):
+  - regex(column, pattern) - PATINDEX uses LIKE wildcards, not regex
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Any, Dict, List, Tuple
+
+from kontra.connectors.handle import DatasetHandle
+from kontra.connectors.sqlserver import SqlServerConnectionParams, get_connection
+from kontra.connectors.detection import parse_table_reference, get_default_schema, SQLSERVER
+
+from .database_base import DatabaseSqlExecutor
+from .registry import register_executor
+
+
+@register_executor("sqlserver")
+class SqlServerSqlExecutor(DatabaseSqlExecutor):
+    """
+    SQL Server SQL pushdown executor.
+
+    Inherits compile() and execute() from DatabaseSqlExecutor.
+    Provides SQL Server-specific connection and table handling.
+
+    Note: regex is NOT supported because PATINDEX uses LIKE-style wildcards.
+    Regex rules fall back to Polars execution.
+    """
+
+    DIALECT = "sqlserver"
+    # Note: regex is NOT supported - PATINDEX uses LIKE wildcards, not regex.
+    # But contains/starts_with/ends_with use LIKE, so they work!
+    SUPPORTED_RULES = {
+        "not_null", "unique", "min_rows", "max_rows",
+        "allowed_values", "disallowed_values",
+        "freshness", "range", "length",
+        "contains", "starts_with", "ends_with",  # LIKE-based, works on SQL Server
+        "compare", "conditional_not_null", "conditional_range",
+        "custom_sql_check", "custom_agg"
+    }
+
+    @property
+    def name(self) -> str:
+        return "sqlserver"
+
+    def _supports_scheme(self, scheme: str, handle: DatasetHandle) -> bool:
+        """Check if this executor supports the given URI scheme."""
+        # BYOC: check dialect for external connections
+        if scheme == "byoc" and handle.dialect == "sqlserver":
+            return handle.external_conn is not None
+
+        # URI-based: handle mssql:// or sqlserver:// URIs
+        return scheme in {"mssql", "sqlserver"}
+
+    @contextmanager
+    def _get_connection_ctx(self, handle: DatasetHandle):
+        """
+        Get a SQL Server connection context.
+
+        For BYOC, yields the external connection directly (not owned by us).
+        For URI-based, yields a new connection (owned by context manager).
+        """
+        if handle.scheme == "byoc" and handle.external_conn is not None:
+            yield handle.external_conn
+        elif handle.db_params:
+            with get_connection(handle.db_params) as conn:
+                yield conn
+        else:
+            raise ValueError("Handle has neither external_conn nor db_params")
+
+    def _get_table_reference(self, handle: DatasetHandle) -> str:
+        """
+        Get the fully-qualified table reference for SQL Server.
+
+        Returns: [schema].[table] format.
+        """
+        if handle.scheme == "byoc" and handle.table_ref:
+            _db, schema, table = parse_table_reference(handle.table_ref)
+            schema = schema or get_default_schema(SQLSERVER)
+            return f"{self._esc(schema)}.{self._esc(table)}"
+        elif handle.db_params:
+            params: SqlServerConnectionParams = handle.db_params
+            return f"{self._esc(params.schema)}.{self._esc(params.table)}"
+        else:
+            raise ValueError("Handle has neither table_ref nor db_params")
+
+    def _get_schema_and_table(self, handle: DatasetHandle) -> Tuple[str, str]:
+        """
+        Get schema and table name separately for custom SQL placeholder replacement.
+
+        Returns: Tuple of (schema, table_name)
+        """
+        if handle.scheme == "byoc" and handle.table_ref:
+            _db, schema, table = parse_table_reference(handle.table_ref)
+            schema = schema or get_default_schema(SQLSERVER)
+            return schema, table
+        elif handle.db_params:
+            params: SqlServerConnectionParams = handle.db_params
+            return params.schema, params.table
+        else:
+            raise ValueError("Handle has neither table_ref nor db_params")
+
+    def introspect(self, handle: DatasetHandle, **kwargs) -> Dict[str, Any]:
+        """
+        Introspect the SQL Server table for metadata.
+
+        Returns:
+            {"row_count": int, "available_cols": [...], "staging": None}
+        """
+        table = self._get_table_reference(handle)
+
+        # Get schema and table name for information_schema query
+        if handle.scheme == "byoc" and handle.table_ref:
+            _db, schema, table_name = parse_table_reference(handle.table_ref)
+            schema = schema or get_default_schema(SQLSERVER)
+        elif handle.db_params:
+            params: SqlServerConnectionParams = handle.db_params
+            schema = params.schema
+            table_name = params.table
+        else:
+            raise ValueError("Handle has neither table_ref nor db_params")
+
+        with self._get_connection_ctx(handle) as conn:
+            cursor = conn.cursor()
+            try:
+                # Get row count
+                cursor.execute(f"SELECT COUNT(*) FROM {table}")
+                row_count = cursor.fetchone()
+                n = int(row_count[0]) if row_count else 0
+
+                # Get column names (pymssql uses %s for parameters)
+                cursor.execute(
+                    """
+                    SELECT column_name
+                    FROM information_schema.columns
+                    WHERE table_schema = %s AND table_name = %s
+                    ORDER BY ordinal_position
+                    """,
+                    (schema, table_name),
+                )
+                cols = [row[0] for row in cursor.fetchall()]
+            finally:
+                cursor.close()
+
+        return {"row_count": n, "available_cols": cols, "staging": None}

kontra/engine/materializers/__init__.py

@@ -0,0 +1,14 @@
+# src/kontra/engine/materializers/__init__.py
+from .base import BaseMaterializer
+from .registry import (
+    pick_materializer,
+    register_default_materializers,
+    register_materializer,
+)
+
+__all__ = [
+    "BaseMaterializer",
+    "pick_materializer",
+    "register_materializer",
+    "register_default_materializers",
+]

kontra/engine/materializers/base.py

@@ -0,0 +1,42 @@
+# src/kontra/engine/materializers/base.py
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+import polars as pl
+
+from kontra.connectors.handle import DatasetHandle
+
+
+class BaseMaterializer:
+    """
+    Minimal base class for materializers.
+
+    Defines the interface for:
+      - Loading a source into a Polars DataFrame (with projection)
+      - Peeking the schema
+      - Reporting I/O diagnostics
+    """
+
+    materializer_name: str = "unknown"
+
+    def __init__(self, handle: DatasetHandle):
+        """
+        Initialize the materializer with a data source handle.
+
+        Args:
+            handle: The DatasetHandle containing the URI and fs_opts.
+        """
+        self.handle = handle
+
+    def schema(self) -> List[str]:
+        """Return column names without materializing data (best effort)."""
+        raise NotImplementedError
+
+    def to_polars(self, columns: Optional[List[str]]) -> pl.DataFrame:
+        """Materialize directly as a Polars DataFrame."""
+        raise NotImplementedError
+
+    def io_debug(self) -> Optional[Dict[str, Any]]:
+        """Return last I/O diagnostics for observability (or None)."""
+        return None  # Default implementation

kontra/engine/materializers/duckdb.py

@@ -0,0 +1,110 @@
+# src/kontra/engine/materializers/duckdb.py
+from __future__ import annotations
+
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+import polars as pl
+import duckdb
+
+# --- Kontra Imports ---
+from kontra.engine.backends.duckdb_session import create_duckdb_connection
+from kontra.engine.backends.duckdb_utils import (
+    esc_ident,
+    lit_str,
+)
+from kontra.connectors.handle import DatasetHandle
+
+from .base import BaseMaterializer  # Import from new base file
+from .registry import register_materializer
+
+
+@register_materializer("duckdb")
+class DuckDBMaterializer(BaseMaterializer):
+    """
+    Column-pruned materialization via DuckDB httpfs → Arrow → Polars.
+
+    Guarantees:
+      - **Format aware**: Parquet via read_parquet(), CSV via read_csv_auto().
+      - **Projection**: SELECT only requested columns at source (true pruning).
+      - **Low copy**: Arrow table handoff → Polars DataFrame.
+      - **Remote support**: S3/HTTP via DuckDB httpfs (loaded in session factory).
+
+    Scope:
+      - I/O only. Does NOT execute rule SQL; the SQL executor handles pushdown.
+    """
+
+    def __init__(self, handle: DatasetHandle):
+        super().__init__(handle)
+        self.source = handle.uri
+        self._io_debug_enabled = bool(os.getenv("KONTRA_IO_DEBUG"))
+        self._last_io_debug: Optional[Dict[str, Any]] = None
+        self.con = create_duckdb_connection(self.handle)
+
+    # ---------- Materializer API ----------
+
+    def schema(self) -> List[str]:
+        """
+        Return column names without materializing data (best effort, format-aware).
+        """
+        read_fn = self._get_read_function()
+        cur = self.con.execute(
+            f"SELECT * FROM {read_fn}({lit_str(self.source)}) LIMIT 0"
+        )
+        return [d[0] for d in cur.description] if cur.description else []
+
+    def to_polars(self, columns: Optional[List[str]]) -> pl.DataFrame:
+        """
+        Materialize the requested columns as a Polars DataFrame via Arrow.
+        """
+        # Route through Arrow for consistent, low-copy materialization.
+        import pyarrow as pa  # noqa: F401
+
+        cols_sql = (
+            ", ".join(esc_ident(c) for c in (columns or [])) if columns else "*"
+        )
+        read_func = self._get_read_function()
+
+        t0 = time.perf_counter()
+        query = f"SELECT {cols_sql} FROM {read_func}({lit_str(self.source)})"
+        cur = self.con.execute(query)
+        table = cur.fetch_arrow_table()
+        t1 = time.perf_counter()
+
+        if self._io_debug_enabled:
+            self._last_io_debug = {
+                "materializer": "duckdb",
+                "mode": "duckdb_project_to_arrow",
+                "columns_requested": list(columns or []),
+                "column_count": len(columns or []),
+                "elapsed_ms": int((t1 - t0) * 1000),
+            }
+        else:
+            self._last_io_debug = None
+
+        return pl.from_arrow(table)
+
+    def io_debug(self) -> Optional[Dict[str, Any]]:
+        return self._last_io_debug
+
+    # ---------- Internals ----------
+
+    def _get_read_function(self) -> str:
+        """
+        Return the correct DuckDB read function based on file format.
+
+        Notes:
+          - For CSV, we prefer DuckDB's read_csv_auto() for performance.
+          - CSV options (delimiter, header, etc.) can be threaded from
+            connector handle in future (TODO), but auto inference is robust
+            for most standardized data lake dumps.
+        """
+        fmt = (self.handle.format or "").lower()
+        if fmt == "parquet":
+            return "read_parquet"
+        if fmt == "csv":
+            return "read_csv_auto"
+
+        # Fallback: attempt format autodetection; Parquet is most common.
+        return "read_parquet"

kontra/engine/materializers/factory.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+from typing import Optional
+from ...connectors.capabilities import ConnectorCapabilities as CC  # re-export path differs when imported from engine
+from .duckdb import DuckDBMaterializer
+from .polars_connector import PolarsConnectorMaterializer
+
+def is_s3_uri(val: str | None) -> bool:
+    return isinstance(val, str) and val.lower().startswith("s3://")
+
+class MaterializerFactory:
+    @staticmethod
+    def from_source(source: str, connector, caps: int, prefer_remote_pruning: bool):
+        """
+        Choose the best materializer for a given source and connector capabilities.
+
+        Strategy (v1):
+          - If remote S3 and we prefer pruning → DuckDBMaterializer (httpfs + Arrow)
+          - Else → PolarsConnectorMaterializer (direct connector.load)
+        """
+        if is_s3_uri(source) and prefer_remote_pruning and (caps & (CC.PUSHDOWN | CC.REMOTE_PARTIAL)):
+            return DuckDBMaterializer(source)
+        return PolarsConnectorMaterializer(source, connector)