kontra 0.5.2__py3-none-any.whl

kontra/connectors/sqlserver.py

@@ -0,0 +1,226 @@
+# src/kontra/connectors/sqlserver.py
+"""
+SQL Server connection utilities for Kontra.
+
+Supports multiple authentication methods:
+1. Full URI: mssql://user:pass@host:port/database/schema.table
+2. Environment variables: MSSQL_HOST, MSSQL_PORT, MSSQL_USER, MSSQL_PASSWORD, MSSQL_DATABASE
+3. SQLSERVER_URL (similar to DATABASE_URL pattern)
+
+Priority: URI values > SQLSERVER_URL > MSSQL_XXX env vars > defaults
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from .db_utils import (
+    DbConnectionConfig,
+    resolve_connection_params as _resolve_params,
+)
+
+
+# SQL Server-specific configuration for parameter resolution
+_MSSQL_CONFIG = DbConnectionConfig(
+    default_host="localhost",
+    default_port=1433,
+    default_user="sa",
+    default_schema="dbo",
+    env_host="MSSQL_HOST",
+    env_port="MSSQL_PORT",
+    env_user="MSSQL_USER",
+    env_password="MSSQL_PASSWORD",
+    env_database="MSSQL_DATABASE",
+    env_url="SQLSERVER_URL",
+    db_name="SQL Server",
+    uri_example="mssql://user:pass@host:1433/database/schema.table",
+    env_example="MSSQL_DATABASE",
+)
+
+
+@dataclass
+class SqlServerConnectionParams:
+    """Resolved SQL Server connection parameters."""
+
+    host: str
+    port: int
+    user: str
+    password: Optional[str]
+    database: str
+    schema: str
+    table: str
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Return connection kwargs for pymssql.connect()."""
+        return {
+            "server": self.host,
+            "port": self.port,
+            "user": self.user,
+            "password": self.password,
+            "database": self.database,
+        }
+
+    @property
+    def qualified_table(self) -> str:
+        """Return schema.table identifier."""
+        return f"{self.schema}.{self.table}"
+
+
+def resolve_connection_params(uri: str) -> SqlServerConnectionParams:
+    """
+    Resolve SQL Server connection parameters from URI + environment.
+
+    URI format:
+        mssql://user:pass@host:port/database/schema.table
+        mssql:///dbo.users  (uses env vars for connection)
+        sqlserver://...  (alias for mssql://)
+
+    Priority: URI values > SQLSERVER_URL > MSSQL_XXX env vars > defaults
+
+    Raises:
+        ValueError: If required parameters (database, table) cannot be resolved.
+    """
+    resolved = _resolve_params(uri, _MSSQL_CONFIG)
+
+    return SqlServerConnectionParams(
+        host=resolved.host,
+        port=resolved.port,
+        user=resolved.user,
+        password=resolved.password,
+        database=resolved.database,  # type: ignore (validated in _resolve_params)
+        schema=resolved.schema,
+        table=resolved.table,  # type: ignore (validated in _resolve_params)
+    )
+
+
+def get_connection(params: SqlServerConnectionParams):
+    """
+    Create a pymssql connection from resolved parameters.
+
+    Returns:
+        pymssql.Connection
+    """
+    try:
+        import pymssql
+    except ImportError as e:
+        raise ImportError(
+            "pymssql is required for SQL Server support.\n"
+            "Install with: pip install pymssql"
+        ) from e
+
+    try:
+        return pymssql.connect(**params.to_dict())
+    except pymssql.OperationalError as e:
+        raise ConnectionError(
+            f"SQL Server connection failed: {e}\n\n"
+            f"Connection details:\n"
+            f"  Host: {params.host}:{params.port}\n"
+            f"  Database: {params.database}\n"
+            f"  User: {params.user}\n\n"
+            "Check your connection settings or set environment variables:\n"
+            "  export MSSQL_HOST=localhost\n"
+            "  export MSSQL_PORT=1433\n"
+            "  export MSSQL_USER=your_user\n"
+            "  export MSSQL_PASSWORD=your_password\n"
+            "  export MSSQL_DATABASE=your_database"
+        ) from e
+
+
+def fetch_sqlserver_stats(params: SqlServerConnectionParams) -> Dict[str, Dict[str, Any]]:
+    """
+    Fetch SQL Server statistics from sys.dm_db_stats_properties and related DMVs.
+
+    Returns a dict keyed by column name with stats:
+        {
+            "column_name": {
+                "null_frac": 0.02,        # Estimated fraction of nulls
+                "n_distinct": 1000,       # Estimated distinct values (-1 = unique)
+                "rows": 10000,            # Rows when stats were computed
+            },
+            "__table__": {
+                "row_estimate": 10000,
+                "page_count": 100,
+            }
+        }
+    """
+    import pymssql
+
+    with get_connection(params) as conn:
+        with conn.cursor() as cursor:
+            # Table-level stats from sys.dm_db_partition_stats
+            cursor.execute(
+                """
+                SELECT SUM(row_count) AS row_estimate,
+                       SUM(used_page_count) AS page_count
+                FROM sys.dm_db_partition_stats ps
+                JOIN sys.objects o ON ps.object_id = o.object_id
+                JOIN sys.schemas s ON o.schema_id = s.schema_id
+                WHERE s.name = %s AND o.name = %s AND ps.index_id IN (0, 1)
+                """,
+                (params.schema, params.table),
+            )
+            row = cursor.fetchone()
+            table_stats = {
+                "row_estimate": row[0] if row and row[0] else 0,
+                "page_count": row[1] if row and row[1] else 0,
+            }
+
+            # Column-level stats from sys.dm_db_stats_properties + DBCC SHOW_STATISTICS
+            # We use density_vector from stats to estimate distinct values
+            cursor.execute(
+                """
+                SELECT
+                    c.name AS column_name,
+                    s.name AS stat_name,
+                    sp.rows,
+                    sp.modification_counter
+                FROM sys.stats s
+                JOIN sys.stats_columns sc ON s.stats_id = sc.stats_id AND s.object_id = sc.object_id
+                JOIN sys.columns c ON sc.column_id = c.column_id AND sc.object_id = c.object_id
+                CROSS APPLY sys.dm_db_stats_properties(s.object_id, s.stats_id) sp
+                WHERE s.object_id = OBJECT_ID(%s)
+                """,
+                (f"{params.schema}.{params.table}",),
+            )
+
+            result: Dict[str, Dict[str, Any]] = {"__table__": table_stats}
+
+            for row in cursor.fetchall():
+                col_name, stat_name, rows, mod_counter = row
+                if col_name not in result:
+                    result[col_name] = {
+                        "rows": rows,
+                        "modification_counter": mod_counter,
+                        "stat_name": stat_name,
+                    }
+
+            # For each column with stats, get density (1/distinct) from DBCC SHOW_STATISTICS
+            # This requires more complex parsing, so we'll do a simpler approach:
+            # Query actual distinct counts for key columns (more reliable for preplan)
+            for col_name in list(result.keys()):
+                if col_name == "__table__":
+                    continue
+                try:
+                    # Get null fraction
+                    cursor.execute(
+                        f"""
+                        SELECT
+                            CAST(SUM(CASE WHEN [{col_name}] IS NULL THEN 1 ELSE 0 END) AS FLOAT)
+                            / NULLIF(COUNT(*), 0) AS null_frac,
+                            COUNT(DISTINCT [{col_name}]) AS n_distinct
+                        FROM [{params.schema}].[{params.table}]
+                        """,
+                    )
+                    stats_row = cursor.fetchone()
+                    if stats_row:
+                        result[col_name]["null_frac"] = stats_row[0] or 0.0
+                        result[col_name]["n_distinct"] = stats_row[1] or 0
+                        # Mark as unique if distinct = row count
+                        if result[col_name]["n_distinct"] == table_stats["row_estimate"]:
+                            result[col_name]["n_distinct"] = -1  # Convention: -1 = all unique
+                except Exception:
+                    # Stats query failed, leave partial data
+                    pass
+
+            return result

kontra/engine/backends/duckdb_session.py

@@ -0,0 +1,227 @@
+# src/kontra/backends/duckdb_session.py
+from __future__ import annotations
+
+import os
+from typing import Any, Dict
+from urllib.parse import urlparse
+
+import duckdb
+from kontra.connectors.handle import DatasetHandle
+
+# --- Public API ---
+
+
+def create_duckdb_connection(handle: DatasetHandle) -> duckdb.DuckDBPyConnection:
+    """
+    Create a DuckDB connection configured specifically for the given DatasetHandle.
+
+    This is the centralized factory for all DuckDB instances in Kontra.
+    It inspects the handle's scheme and fs_opts to load the correct
+    extensions (httpfs) and apply the necessary configuration
+    (e.g., S3 endpoints, credentials, region) for I/O.
+
+    Args:
+        handle: The DatasetHandle containing the URI and filesystem options.
+
+    Returns:
+        A configured duckdb.DuckDBPyConnection.
+    """
+    con = duckdb.connect()
+
+    # Apply performance/threading tweaks (reads env, but for runtime, not I/O)
+    _configure_threads(con)
+
+    # Apply I/O and credential configuration based on the data source
+    match handle.scheme:
+        case "s3":
+            _configure_s3(con, handle.fs_opts)
+        case "abfs" | "abfss" | "az":
+            _configure_azure(con, handle.fs_opts)  # Stubbed for future work
+        case "http" | "https":
+            _configure_http(con, handle.fs_opts)
+        case "file" | "":
+            # Local files need no special I/O config
+            pass
+        case _:
+            # Best-effort for unknown schemes: load httpfs just in case
+            try:
+                _configure_http(con, handle.fs_opts)
+            except Exception:
+                pass  # Ignore if httpfs fails to load
+
+    return con
+
+
+# --- Internal Helpers ---
+
+
+def _safe_set(con: duckdb.DuckDBPyConnection, key: str, value: Any) -> None:
+    """
+    Safely execute a DuckDB SET command, ignoring errors.
+    """
+    try:
+        con.execute(f"SET {key} = ?", [str(value)])
+    except Exception:
+        # Fails gracefully if the setting doesn't exist (e.g., wrong DuckDB version)
+        pass
+
+
+def _configure_threads(con: duckdb.DuckDBPyConnection) -> None:
+    """
+    Configure DuckDB thread count based on env vars or CPU count.
+    This is a performance tweak, not an I/O secret.
+    """
+    env_threads = os.getenv("DUCKDB_THREADS")
+    try:
+        nthreads = int(env_threads) if env_threads else (os.cpu_count() or 4)
+    except Exception:
+        nthreads = os.cpu_count() or 4
+
+    # Try both PRAGMA (older) and SET (newer) for compatibility
+    for sql in (f"PRAGMA threads={int(nthreads)};", f"SET threads = {int(nthreads)};"):
+        try:
+            con.execute(sql)
+            break
+        except Exception:
+            continue
+
+
+def _configure_http(
+    con: duckdb.DuckDBPyConnection, fs_opts: Dict[str, str]
+) -> None:
+    """
+    Install and load the httpfs extension for reading http(s):// files.
+    """
+    con.execute("INSTALL httpfs;")
+    con.execute("LOAD httpfs;")
+    _safe_set(con, "enable_object_cache", "true")
+
+
+def _configure_s3(con: duckdb.DuckDBPyConnection, fs_opts: Dict[str, str]) -> None:
+    """
+    Configure the httpfs extension for S3-compatible storage (AWS, MinIO, R2).
+
+    Expected fs_opts keys:
+    - s3_endpoint
+    - s3_region
+    - s3_url_style ('path' | 'host')
+    - s3_use_ssl ('true' | 'false')
+    - s3_access_key_id
+    - s3_secret_access_key
+    - s3_session_token
+    - s3_max_connections
+    """
+    _configure_http(con, fs_opts)  # S3 depends on httpfs
+
+    # Credentials
+    if ak := fs_opts.get("s3_access_key_id"):
+        _safe_set(con, "s3_access_key_id", ak)
+    if sk := fs_opts.get("s3_secret_access_key"):
+        _safe_set(con, "s3_secret_access_key", sk)
+    if st := fs_opts.get("s3_session_token"):
+        _safe_set(con, "s3_session_token", st)
+
+    # Region
+    if region := fs_opts.get("s3_region"):
+        _safe_set(con, "s3_region", region)
+
+    # Endpoint (MinIO/S3-compatible)
+    endpoint = fs_opts.get("s3_endpoint")
+    url_style = fs_opts.get("s3_url_style")
+    use_ssl = fs_opts.get("s3_use_ssl")
+
+    if endpoint:
+        # Parse "http://host:port" or just "host:port"
+        parsed = urlparse(endpoint)
+        hostport = parsed.netloc or parsed.path or endpoint
+        _safe_set(con, "s3_endpoint", hostport)
+
+        # Infer SSL from endpoint scheme if not explicitly set
+        if use_ssl is None:
+            use_ssl = "true" if parsed.scheme == "https" else "false"
+        _safe_set(con, "s3_use_ssl", use_ssl)
+
+        # Default to path-style for custom endpoints (MinIO-friendly)
+        if url_style is None:
+            url_style = "path"
+
+    if url_style:
+        _safe_set(con, "s3_url_style", url_style)
+
+    # Performance and reliability for large files over S3/HTTP
+    # http_timeout is in seconds (default 30s - increase for large files)
+    _safe_set(con, "http_timeout", "600")  # 10 minutes for large files
+    _safe_set(con, "http_retries", "5")    # More retries for reliability
+    _safe_set(con, "http_retry_wait_ms", "2000")  # 2s between retries
+    # Disable keep-alive for MinIO/S3-compatible - connection pooling can cause issues
+    _safe_set(con, "http_keep_alive", "false")
+
+
+def _configure_azure(
+    con: duckdb.DuckDBPyConnection, fs_opts: Dict[str, str]
+) -> None:
+    """
+    Configure the Azure extension for ADLS Gen2 (abfs://, abfss://) and Azure Blob (az://).
+
+    DuckDB 0.10+ has native Azure support via the 'azure' extension.
+    This handles authentication and endpoint configuration.
+
+    Expected fs_opts keys:
+    - azure_account_name: Storage account name
+    - azure_account_key: Storage account key
+    - azure_sas_token: SAS token (alternative to key)
+    - azure_connection_string: Full connection string (alternative)
+    - azure_tenant_id: For OAuth/service principal
+    - azure_client_id: For OAuth/service principal
+    - azure_client_secret: For OAuth/service principal
+    - azure_endpoint: Custom endpoint (Databricks, sovereign clouds, Azurite)
+
+    Raises:
+        RuntimeError: If Azure extension is not available (DuckDB < 0.10.0)
+    """
+    # Install and load the Azure extension
+    try:
+        con.execute("INSTALL azure;")
+        con.execute("LOAD azure;")
+    except Exception as e:
+        raise RuntimeError(
+            f"Azure extension not available. DuckDB >= 0.10.0 is required for Azure support. "
+            f"Error: {e}"
+        ) from e
+
+    # Account name (required for key/SAS auth)
+    if account_name := fs_opts.get("azure_account_name"):
+        _safe_set(con, "azure_storage_account_name", account_name)
+
+    # Account key auth
+    if account_key := fs_opts.get("azure_account_key"):
+        _safe_set(con, "azure_account_key", account_key)
+
+    # SAS token auth (alternative to account key)
+    # Note: DuckDB expects the token without leading '?'
+    if sas_token := fs_opts.get("azure_sas_token"):
+        # Strip leading '?' if present
+        if sas_token.startswith("?"):
+            sas_token = sas_token[1:]
+        _safe_set(con, "azure_sas_token", sas_token)
+
+    # Connection string auth
+    if conn_string := fs_opts.get("azure_connection_string"):
+        _safe_set(con, "azure_storage_connection_string", conn_string)
+
+    # OAuth / Service Principal auth
+    if tenant_id := fs_opts.get("azure_tenant_id"):
+        _safe_set(con, "azure_tenant_id", tenant_id)
+    if client_id := fs_opts.get("azure_client_id"):
+        _safe_set(con, "azure_client_id", client_id)
+    if client_secret := fs_opts.get("azure_client_secret"):
+        _safe_set(con, "azure_client_secret", client_secret)
+
+    # Custom endpoint (for Databricks, sovereign clouds, Azurite emulator)
+    if endpoint := fs_opts.get("azure_endpoint"):
+        _safe_set(con, "azure_endpoint", endpoint)
+
+    # Performance settings (same as S3)
+    _safe_set(con, "http_timeout", "600")  # 10 minutes for large files
+    _safe_set(con, "http_retries", "5")
+    _safe_set(con, "http_retry_wait_ms", "2000")

kontra/engine/backends/duckdb_utils.py

@@ -0,0 +1,18 @@
+# src/kontra/engine/backends/duckdb_utils.py
+from __future__ import annotations
+
+
+def esc_ident(name: str) -> str:
+    """
+    Quote an identifier for DuckDB (double quotes, escape internal quotes).
+    This is a centralized helper used by executors and materializers.
+    """
+    return '"' + name.replace('"', '""') + '"'
+
+
+def lit_str(s: str) -> str:
+    """
+    Return a single-quoted SQL string literal with internal quotes escaped.
+    This is a centralized helper used by executors and materializers.
+    """
+    return "'" + s.replace("'", "''") + "'"

kontra/engine/backends/polars_backend.py

@@ -0,0 +1,47 @@
+# src/kontra/engine/backends/polars_backend.py
+from __future__ import annotations
+
+"""
+Polars Backend (Adapter)
+
+Thin adapter that defers execution to the RuleExecutionPlan's compiled executor.
+Keeps the backend boundary explicit and behavior deterministic.
+"""
+
+from typing import Any, Callable, Dict, List
+
+import polars as pl
+
+
+class PolarsBackend:
+    name = "polars"
+
+    def __init__(self, executor: Callable[[pl.DataFrame, Any], List[Dict[str, Any]]]):
+        """
+        Parameters
+        ----------
+        executor : callable
+            Function that evaluates the compiled plan against a materialized
+            Polars DataFrame (typically RuleExecutionPlan.execute_compiled).
+        """
+        self._executor = executor
+
+    def supports(self, connector_caps: int) -> bool:
+        """Capability hook reserved for future; always True for local DataFrames."""
+        return True
+
+    def compile(self, compiled_plan: Any) -> Any:
+        """No-op for Polars: pass through the compiled plan."""
+        return compiled_plan
+
+    def execute(self, df: pl.DataFrame, compiled_artifact: Any) -> Dict[str, Any]:
+        """Execute the compiled artifact against `df` and wrap results."""
+        results = self._executor(df, compiled_artifact)
+        return {"results": results}
+
+    def introspect(self, df: pl.DataFrame) -> Dict[str, Any]:
+        """Basic observability: row count and available columns."""
+        return {
+            "row_count": int(df.height),
+            "available_cols": list(df.columns),
+        }