kontra 0.5.2__py3-none-any.whl

kontra/engine/materializers/polars_connector.py

@@ -0,0 +1,131 @@
+# src/kontra/engine/materializers/polars_connector.py
+from __future__ import annotations
+
+"""
+PolarsConnectorMaterializer
+
+Purpose
+-------
+Local, dependency-light materializer that produces a Polars DataFrame from
+file-based datasets (Parquet/CSV). Supports column projection. Does *not*
+require the legacy connectors package.
+
+Design
+------
+- First tries legacy `ConnectorFactory` (for back-compat if present).
+- Otherwise uses native Polars lazy scans:
+    - scan_*  → optional .select(projection) → collect()
+
+Notes
+-----
+Polars ≥ 1.34 removed `columns=` from scan_*; apply projection via `.select()`.
+"""
+
+from typing import Any, Dict, List, Optional
+
+import polars as pl
+
+from kontra.connectors.handle import DatasetHandle
+from .base import BaseMaterializer
+from .registry import register_materializer
+
+
+def _infer_format(uri: str, explicit: Optional[str]) -> str:
+    """Resolve file format from explicit handle.format or file extension."""
+    if explicit:
+        return explicit.lower()
+    low = uri.lower()
+    if low.endswith(".parquet"):
+        return "parquet"
+    if low.endswith(".csv"):
+        return "csv"
+    return ""
+
+
+@register_materializer("polars-connector")
+class PolarsConnectorMaterializer(BaseMaterializer):
+    """
+    Minimal, deterministic materializer for local files.
+
+    Responsibilities
+    ----------------
+    - Cheap schema peek (names only)
+    - DataFrame materialization with optional projection
+    - No side effects; no hidden state
+    """
+
+    name = "polars-connector"
+
+    def __init__(self, handle: DatasetHandle):
+        super().__init__(handle)
+        self._io_debug: Optional[Dict[str, Any]] = None  # retained for parity with duckdb materializer
+
+    # ------------------------------------------------------------------ #
+    # Introspection
+    # ------------------------------------------------------------------ #
+
+    def schema(self) -> List[str]:
+        """
+        Return column names using a lazy scan. Never raises — empty list on failure.
+        """
+        uri = self.handle.uri
+        fmt = _infer_format(uri, getattr(self.handle, "format", None))
+
+        try:
+            if fmt == "parquet":
+                return list(pl.scan_parquet(uri).collect_schema().names())
+            if fmt == "csv":
+                return list(pl.scan_csv(uri).collect_schema().names())
+        except Exception:
+            pass
+        return []
+
+    # ------------------------------------------------------------------ #
+    # Materialization
+    # ------------------------------------------------------------------ #
+
+    def to_polars(self, columns: Optional[List[str]]) -> pl.DataFrame:
+        """
+        Materialize dataset into a Polars DataFrame.
+
+        Strategy
+        --------
+        1) Attempt legacy connectors path (if installed) to preserve behavior.
+        2) Otherwise, native Polars scan with projection via `.select()`.
+        """
+        # --- Legacy path (optional/back-compat) --------------------------------
+        try:
+            from kontra.connectors.factory import ConnectorFactory  # type: ignore
+
+            connector = ConnectorFactory.from_source(self.handle.uri)
+            # The legacy API accepts `columns=` (best-effort).
+            return connector.load(self.handle.uri, columns=columns)
+        except (ImportError, ModuleNotFoundError):
+            # Fall back to native Polars path
+            pass
+
+        # --- Native Polars path -------------------------------------------------
+        uri = self.handle.uri
+        fmt = _infer_format(uri, getattr(self.handle, "format", None))
+
+        if fmt == "parquet":
+            lf = pl.scan_parquet(uri)
+        elif fmt == "csv":
+            # Add CSV options here if your data requires (delimiter, nulls, dtypes).
+            lf = pl.scan_csv(uri)
+        else:
+            raise IOError(f"Unsupported format for PolarsConnectorMaterializer: {uri}")
+
+        if columns:
+            lf = lf.select([pl.col(c) for c in columns])
+
+        # NOTE: streaming=True is deprecated; default engine suffices for tests and CI.
+        return lf.collect()
+
+    # ------------------------------------------------------------------ #
+    # Diagnostics
+    # ------------------------------------------------------------------ #
+
+    def io_debug(self) -> Optional[Dict[str, Any]]:
+        """Reserved hook for I/O diagnostics (none for this materializer)."""
+        return None

kontra/engine/materializers/postgres.py

@@ -0,0 +1,157 @@
+# src/kontra/engine/materializers/postgres.py
+"""
+PostgreSQL Materializer - loads PostgreSQL tables to Polars DataFrames.
+
+Uses psycopg3's efficient binary COPY protocol for streaming data.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+import polars as pl
+
+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 contextlib import contextmanager
+
+from .base import BaseMaterializer
+from .registry import register_materializer
+
+
+@contextmanager
+def _get_connection_ctx(handle: DatasetHandle):
+    """
+    Get a connection context for either BYOC or URI-based handles.
+
+    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:
+        # BYOC: yield external connection directly, don't close it
+        yield handle.external_conn
+    elif handle.db_params:
+        # URI-based: use our connection manager
+        with get_connection(handle.db_params) as conn:
+            yield conn
+    else:
+        raise ValueError("Handle has neither external_conn nor db_params")
+
+
+@register_materializer("postgres")
+class PostgresMaterializer(BaseMaterializer):
+    """
+    Materialize PostgreSQL tables as Polars DataFrames with column projection.
+
+    Features:
+      - Efficient data loading via psycopg3
+      - Column projection at source (SELECT only needed columns)
+      - Binary protocol for faster transfers (when available)
+      - BYOC (Bring Your Own Connection) support
+    """
+
+    def __init__(self, handle: DatasetHandle):
+        super().__init__(handle)
+
+        self._is_byoc = handle.scheme == "byoc" and handle.external_conn is not None
+
+        if self._is_byoc:
+            # BYOC: get table info from handle
+            if not handle.table_ref:
+                raise ValueError("BYOC handle missing table_ref")
+            _db, schema, table = parse_table_reference(handle.table_ref)
+            self._schema_name = schema or get_default_schema(POSTGRESQL)
+            self._table_name = table
+            self._qualified_table = f'{_esc_ident(self._schema_name)}.{_esc_ident(self._table_name)}'
+        elif handle.db_params:
+            # URI-based: use params
+            self.params: PostgresConnectionParams = handle.db_params
+            self._schema_name = self.params.schema
+            self._table_name = self.params.table
+            self._qualified_table = self.params.qualified_table
+        else:
+            raise ValueError("PostgreSQL handle missing db_params or external_conn")
+
+        self._io_debug_enabled = bool(os.getenv("KONTRA_IO_DEBUG"))
+        self._last_io_debug: Optional[Dict[str, Any]] = None
+
+    def schema(self) -> List[str]:
+        """Return column names without loading data."""
+        with _get_connection_ctx(self.handle) as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    SELECT column_name
+                    FROM information_schema.columns
+                    WHERE table_schema = %s AND table_name = %s
+                    ORDER BY ordinal_position
+                    """,
+                    (self._schema_name, self._table_name),
+                )
+                return [row[0] for row in cur.fetchall()]
+
+    def to_polars(self, columns: Optional[List[str]]) -> pl.DataFrame:
+        """
+        Load table data as a Polars DataFrame with optional column projection.
+
+        Supports both URI-based connections (handle.db_params) and
+        BYOC connections (handle.external_conn).
+
+        Args:
+            columns: List of columns to load. If None, loads all columns.
+
+        Returns:
+            Polars DataFrame with the requested columns.
+        """
+        t0 = time.perf_counter()
+
+        # Build column list for SELECT
+        if columns:
+            cols_sql = ", ".join(_esc_ident(c) for c in columns)
+        else:
+            cols_sql = "*"
+
+        query = f"SELECT {cols_sql} FROM {self._qualified_table}"
+
+        with _get_connection_ctx(self.handle) as conn:
+            with conn.cursor() as cur:
+                cur.execute(query)
+                # Fetch all rows - for large tables, consider chunked loading
+                rows = cur.fetchall()
+                col_names = [desc[0] for desc in cur.description] if cur.description else []
+
+        t1 = time.perf_counter()
+
+        # Convert to Polars DataFrame
+        if rows:
+            df = pl.DataFrame(rows, schema=col_names, orient="row")
+        else:
+            # Empty DataFrame with correct schema
+            df = pl.DataFrame(schema={name: pl.Utf8 for name in col_names})
+
+        if self._io_debug_enabled:
+            self._last_io_debug = {
+                "materializer": "postgres",
+                "mode": "psycopg_fetch" if not self._is_byoc else "byoc_fetch",
+                "table": self._qualified_table,
+                "columns_requested": list(columns or []),
+                "column_count": len(columns or col_names),
+                "row_count": len(rows) if rows else 0,
+                "elapsed_ms": int((t1 - t0) * 1000),
+            }
+        else:
+            self._last_io_debug = None
+
+        return df
+
+    def io_debug(self) -> Optional[Dict[str, Any]]:
+        return self._last_io_debug
+
+
+def _esc_ident(name: str) -> str:
+    """Escape a PostgreSQL identifier (column/table name)."""
+    # Double any internal quotes and wrap in quotes
+    return '"' + name.replace('"', '""') + '"'

kontra/engine/materializers/registry.py

@@ -0,0 +1,138 @@
+# src/kontra/engine/materializers/registry.py
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Dict, List
+
+from kontra.connectors.handle import DatasetHandle
+
+if TYPE_CHECKING:
+    # Import from the new base file
+    from .base import BaseMaterializer as Materializer
+    from .duckdb import DuckDBMaterializer  # noqa: F401
+    from .polars_connector import PolarsConnectorMaterializer  # noqa: F401
+    from .postgres import PostgresMaterializer  # noqa: F401
+    from .sqlserver import SqlServerMaterializer  # noqa: F401
+
+
+# Registry: materializer_name -> ctor(handle) function
+_MATS: Dict[str, Callable[[DatasetHandle], Materializer]] = {}
+# Simple order for picking when multiple can handle a handle
+_ORDER: List[str] = []
+
+
+def register_materializer(name: str):
+    """
+    Decorator to register a materializer class under a stable name.
+    The class must implement the Materializer protocol.
+    """
+
+    def deco(cls: Callable[[DatasetHandle], Materializer]) -> Callable[
+        [DatasetHandle], Materializer
+    ]:
+        if name in _MATS:
+            raise ValueError(f"Materializer '{name}' is already registered.")
+        _MATS[name] = cls
+        if name not in _ORDER:
+            _ORDER.append(name)
+        cls.materializer_name = name  # friendly label for stats.io
+        return cls
+
+    return deco
+
+
+def pick_materializer(handle: DatasetHandle) -> Materializer:
+    """
+    Choose the best materializer for the given dataset handle.
+
+    Policy (v1.4 - BYOC support):
+      - BYOC handles use the materializer matching their dialect.
+      - PostgreSQL URIs use the PostgreSQL materializer.
+      - SQL Server URIs use the SQL Server materializer.
+      - Remote files (s3, http) with known formats use DuckDB materializer.
+      - Otherwise, fall back to PolarsConnector materializer.
+
+    This logic is INDEPENDENT of the projection flag.
+    """
+    # BYOC: route based on dialect
+    if handle.scheme == "byoc":
+        if handle.dialect == "postgresql":
+            ctor = _MATS.get("postgres")
+            if ctor:
+                return ctor(handle)
+            raise RuntimeError(
+                "PostgreSQL materializer not registered. "
+                "Ensure psycopg is installed: pip install 'psycopg[binary]'"
+            )
+        elif handle.dialect == "sqlserver":
+            ctor = _MATS.get("sqlserver")
+            if ctor:
+                return ctor(handle)
+            raise RuntimeError(
+                "SQL Server materializer not registered. "
+                "Ensure pymssql is installed: pip install pymssql"
+            )
+        else:
+            raise RuntimeError(
+                f"Unsupported BYOC dialect: {handle.dialect}. "
+                "Supported: postgresql, sqlserver"
+            )
+
+    # PostgreSQL: use dedicated materializer
+    if handle.scheme in ("postgres", "postgresql"):
+        ctor = _MATS.get("postgres")
+        if ctor:
+            return ctor(handle)
+        raise RuntimeError(
+            "PostgreSQL materializer not registered. "
+            "Ensure psycopg is installed: pip install 'psycopg[binary]'"
+        )
+
+    # SQL Server: use dedicated materializer
+    if handle.scheme in ("mssql", "sqlserver"):
+        ctor = _MATS.get("sqlserver")
+        if ctor:
+            return ctor(handle)
+        raise RuntimeError(
+            "SQL Server materializer not registered. "
+            "Ensure pymssql is installed: pip install pymssql"
+        )
+
+    # Remote files with known formats: use DuckDB for efficient I/O
+    # Includes S3, HTTP(S), and Azure (ADLS Gen2, Azure Blob)
+    is_remote = handle.scheme in ("s3", "http", "https", "abfs", "abfss", "az")
+    is_known_format = handle.format in ("parquet", "csv")
+
+    if is_remote and is_known_format:
+        ctor = _MATS.get("duckdb")
+        if ctor:
+            return ctor(handle)
+
+    # Fallback for local files or unknown formats
+    ctor = _MATS.get("polars-connector")
+    if not ctor:
+        raise RuntimeError(
+            "No default materializer registered (polars-connector missing)"
+        )
+    return ctor(handle)
+
+
+def register_default_materializers() -> None:
+    """
+    Eagerly import built-in materializers so their @register_materializer
+    decorators run and populate the registry.
+    """
+    # Local imports to trigger decorator side-effects
+    from . import duckdb  # noqa: F401
+    from . import polars_connector  # noqa: F401
+
+    # PostgreSQL materializer (optional - requires psycopg)
+    try:
+        from . import postgres  # noqa: F401
+    except ImportError:
+        pass  # psycopg not installed, skip postgres materializer
+
+    # SQL Server materializer (optional - requires pymssql)
+    try:
+        from . import sqlserver  # noqa: F401
+    except ImportError:
+        pass  # pymssql not installed, skip sqlserver materializer

kontra/engine/materializers/sqlserver.py

@@ -0,0 +1,160 @@
+# src/kontra/engine/materializers/sqlserver.py
+"""
+SQL Server Materializer - loads SQL Server tables to Polars DataFrames.
+
+Uses pymssql for database connectivity.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+import polars as pl
+
+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 contextlib import contextmanager
+
+from .base import BaseMaterializer
+from .registry import register_materializer
+
+
+@contextmanager
+def _get_connection_ctx(handle: DatasetHandle):
+    """
+    Get a connection context for either BYOC or URI-based handles.
+
+    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:
+        # BYOC: yield external connection directly, don't close it
+        yield handle.external_conn
+    elif handle.db_params:
+        # URI-based: use our connection manager
+        with get_connection(handle.db_params) as conn:
+            yield conn
+    else:
+        raise ValueError("Handle has neither external_conn nor db_params")
+
+
+@register_materializer("sqlserver")
+class SqlServerMaterializer(BaseMaterializer):
+    """
+    Materialize SQL Server tables as Polars DataFrames with column projection.
+
+    Features:
+      - Efficient data loading via pymssql
+      - Column projection at source (SELECT only needed columns)
+      - BYOC (Bring Your Own Connection) support
+    """
+
+    materializer_name = "sqlserver"
+
+    def __init__(self, handle: DatasetHandle):
+        super().__init__(handle)
+
+        self._is_byoc = handle.scheme == "byoc" and handle.external_conn is not None
+
+        if self._is_byoc:
+            # BYOC: get table info from handle
+            if not handle.table_ref:
+                raise ValueError("BYOC handle missing table_ref")
+            _db, schema, table = parse_table_reference(handle.table_ref)
+            self._schema_name = schema or get_default_schema(SQLSERVER)
+            self._table_name = table
+            self._qualified_table = f'[{self._schema_name}].[{self._table_name}]'
+        elif handle.db_params:
+            # URI-based: use params
+            self.params: SqlServerConnectionParams = handle.db_params
+            self._schema_name = self.params.schema
+            self._table_name = self.params.table
+            self._qualified_table = f'[{self.params.schema}].[{self.params.table}]'
+        else:
+            raise ValueError("SQL Server handle missing db_params or external_conn")
+
+        self._io_debug_enabled = bool(os.getenv("KONTRA_IO_DEBUG"))
+        self._last_io_debug: Optional[Dict[str, Any]] = None
+
+    def schema(self) -> List[str]:
+        """Return column names without loading data."""
+        with _get_connection_ctx(self.handle) as conn:
+            cursor = conn.cursor()
+            # pymssql uses %s as placeholder (pyodbc uses ?)
+            cursor.execute(
+                """
+                SELECT column_name
+                FROM information_schema.columns
+                WHERE table_schema = %s AND table_name = %s
+                ORDER BY ordinal_position
+                """,
+                (self._schema_name, self._table_name),
+            )
+            return [row[0] for row in cursor.fetchall()]
+
+    def to_polars(self, columns: Optional[List[str]]) -> pl.DataFrame:
+        """
+        Load table data as a Polars DataFrame with optional column projection.
+
+        Supports both URI-based connections (handle.db_params) and
+        BYOC connections (handle.external_conn).
+
+        Args:
+            columns: List of columns to load. If None, loads all columns.
+
+        Returns:
+            Polars DataFrame with the requested columns.
+        """
+        t0 = time.perf_counter()
+
+        # Build column list for SELECT
+        if columns:
+            cols_sql = ", ".join(_esc_ident(c) for c in columns)
+        else:
+            cols_sql = "*"
+
+        query = f"SELECT {cols_sql} FROM {self._qualified_table}"
+
+        with _get_connection_ctx(self.handle) as conn:
+            cursor = conn.cursor()
+            cursor.execute(query)
+            # Fetch all rows - for large tables, consider chunked loading
+            rows = cursor.fetchall()
+            col_names = [desc[0] for desc in cursor.description] if cursor.description else []
+
+        t1 = time.perf_counter()
+
+        # Convert to Polars DataFrame
+        if rows:
+            df = pl.DataFrame(rows, schema=col_names, orient="row")
+        else:
+            # Empty DataFrame with correct schema
+            df = pl.DataFrame(schema={name: pl.Utf8 for name in col_names})
+
+        if self._io_debug_enabled:
+            self._last_io_debug = {
+                "materializer": "sqlserver",
+                "mode": "pymssql_fetch" if not self._is_byoc else "byoc_fetch",
+                "table": self._qualified_table,
+                "columns_requested": list(columns or []),
+                "column_count": len(columns or col_names),
+                "row_count": len(rows) if rows else 0,
+                "elapsed_ms": int((t1 - t0) * 1000),
+            }
+        else:
+            self._last_io_debug = None
+
+        return df
+
+    def io_debug(self) -> Optional[Dict[str, Any]]:
+        return self._last_io_debug
+
+
+def _esc_ident(name: str) -> str:
+    """Escape a SQL Server identifier (column/table name)."""
+    # SQL Server uses [brackets] for quoting identifiers
+    # Double any internal brackets
+    return "[" + name.replace("]", "]]") + "]"

kontra/engine/result.py

@@ -0,0 +1,15 @@
+# src/contra/engine/result.py
+from typing import List, Dict, Any
+from dataclasses import dataclass, asdict
+
+@dataclass
+class ValidationResult:
+    dataset: str
+    results: List[Dict[str, Any]]
+    summary: Dict[str, Any]
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+    def passed(self) -> bool:
+        return self.summary.get("passed", False)