dataface 0.1.2__py3-none-any.whl

dataface/core/execute/duckdb_cache.py

@@ -0,0 +1,712 @@
+"""DuckDB caching layer for persistent query result storage.
+
+Stage: EXECUTE (Cache)
+Purpose: Persist query results in DuckDB for fast dataface loads,
+cross-database queries, and snapshot history.
+
+Cache Key: (source_hash, query_hash, variables_hash)
+  - source_hash  — stable hash of the data source identity
+  - query_hash   — SHA-256 of the SQL/query definition
+  - variables_hash — SHA-256 of the variable values that affect the result
+
+face_slug and query_name are METADATA stored in _face_query_index, not key
+components. This lets two faces with identical SQL against the same source
+share one cache entry (cross-face deduplication).
+
+Storage:
+  - _query_results — canonical result store, keyed by (source_hash, query_hash, vars_hash)
+  - _face_query_index — maps (face_slug, query_name) → key triplet, for {{ results.X }}
+  - _query_failures — failure store, keyed by (source_hash, query_hash, vars_hash)
+
+Views: one per cached key, named by combined hash, for {{ results.X }} SQL execution.
+
+Dependencies:
+    - duckdb (optional at module load; required at DuckDBCache() construction)
+    - pyarrow (optional; falls back to row-by-row insert)
+
+See also:
+    - dataface/core/execute/cache_backend.py — QueryResultCache Protocol
+"""
+
+import contextlib
+import hashlib
+import json
+import logging
+import os
+import re
+import threading
+import traceback as tb_mod
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+try:
+    import duckdb
+
+    HAS_DUCKDB = True
+except ImportError:
+    duckdb = None  # type: ignore[assignment]
+    HAS_DUCKDB = False
+
+try:
+    import pyarrow as pa
+
+    HAS_PYARROW = True
+except ImportError:
+    pa = None
+    HAS_PYARROW = False
+
+from dataface._install_hint import install_hint
+from dataface.core.execute.cache_backend import CachedQueryFailure
+from dataface.core.execute.cache_keys import source_identity
+
+logger = logging.getLogger(__name__)
+
+# Metadata columns added to all cached result tables.
+# The key (source_hash, query_hash, variables_hash) is encoded in the table name;
+# we do NOT store it redundantly as per-row columns. Only snapshot sequencing,
+# timestamp, and the full variables JSON (for audit/debug) are stored per-row.
+_METADATA_COLS = {
+    "_snapshot_set_seq",
+    "_snapshot_run_seq",
+    "_run_timestamp",
+    "_variables",
+    "_source_hash",
+}
+
+# How long to keep a combined hash prefix for table names
+_HASH_PREFIX = 12
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Hash helpers (public — imported by executor.py)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def compute_query_hash(sql: str) -> str:
+    """Compute hash of query SQL."""
+    return hashlib.sha256(sql.encode()).hexdigest()[:16]
+
+
+def compute_variables_hash(variables: dict[str, Any]) -> str:
+    """Compute hash of variable values (only the relevant ones for a query)."""
+    if not variables:
+        return "0" * 16
+    serialized = json.dumps(sorted(variables.items()), default=str)
+    return hashlib.sha256(serialized.encode()).hexdigest()[:16]
+
+
+def compute_source_hash(
+    source: str | dict[str, Any] | None,
+    face_sources: dict[str, dict[str, Any]] | None = None,
+) -> str:
+    """Compute a stable hash for a data source identity.
+
+    The hash distinguishes connections with different identities (host,
+    user, project, keyfile, ...) so two users pointing the same source
+    name at different credentials never share a cache entry. Secrets
+    (password, token, api_key, ...) are stripped before hashing — see
+    `cache_keys.SECRET_FIELDS`.
+
+    Args:
+        source: Source identifier. May be:
+          - a connection string / name (str) like "warehouse_prod"
+          - a source config dict with arbitrary keys (inline YAML)
+          - None → treated as the default/in-process source
+        face_sources: Optional resolved-source map (`face.sources`) used
+          to expand a string name into its identity-bearing config dict.
+          When omitted or the name is absent, falls back to hashing the
+          name string as-is.
+
+    Returns:
+        16-char hex string that is stable across calls for the same input.
+    """
+    # Resolve named string references through the project sources map so
+    # different users with the same source name but different credentials
+    # produce different hashes.
+    if isinstance(source, str) and face_sources and source in face_sources:
+        source = face_sources[source]
+
+    identity = source_identity(source)
+
+    if identity is None:
+        # No explicit source — default adapter (in-memory DuckDB, CSV, etc.)
+        canonical = "__default__"
+    elif isinstance(identity, str):
+        canonical = identity.strip()
+    else:
+        # Dict source config: sort keys for stability
+        canonical = json.dumps(sorted(identity.items()), default=str, sort_keys=True)
+    return hashlib.sha256(canonical.encode()).hexdigest()[:16]
+
+
+def _result_table_name(source_hash: str, query_hash: str, variables_hash: str) -> str:
+    """Generate the table name for a cached result.
+
+    Format: _r_{combined_prefix} where combined is a short hash of all three key parts.
+    Short enough to stay under DuckDB's identifier limits.
+    """
+    combined = hashlib.sha256(
+        f"{source_hash}:{query_hash}:{variables_hash}".encode()
+    ).hexdigest()[:_HASH_PREFIX]
+    return f"_r_{combined}"
+
+
+def _view_name_for_table(table_name: str) -> str:
+    return f"{table_name}_v"
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# DuckDBCache
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+class DuckDBCache:
+    """DuckDB-based persistent cache for query results.
+
+    Implements QueryResultCache. Keys all operations on
+    (source_hash, query_hash, variables_hash). face_slug and query_name are
+    metadata stored in _face_query_index.
+
+    Attributes:
+        conn: DuckDB connection
+        db_path: Path to DuckDB file (or None for :memory:)
+        failure_ttl_seconds: Failure TTL; 0 disables failure caching.
+    """
+
+    def __init__(
+        self,
+        db_path: Path | None = None,
+        failure_ttl_seconds: int = 900,
+    ) -> None:
+        if not HAS_DUCKDB:
+            raise ImportError(
+                "DuckDB is required to use DuckDBCache. "
+                f"Install it with: {install_hint('duckdb')}"
+            )
+
+        self.db_path = db_path
+        self.failure_ttl_seconds = failure_ttl_seconds
+        self._lock = threading.RLock()
+
+        if db_path:
+            self.conn = duckdb.connect(str(db_path))
+        else:
+            self.conn = duckdb.connect(":memory:")
+
+        self.conn.execute("SET enable_object_cache = true")
+        self._drop_legacy_tables()
+        self._ensure_schema()
+
+    # ─────────────────────────────────────────────────────────────────────
+    # QueryResultCache Protocol implementation
+    # ─────────────────────────────────────────────────────────────────────
+
+    def get(
+        self, source_hash: str, query_hash: str, variables_hash: str
+    ) -> list[dict[str, Any]] | None:
+        """Return cached rows for this key, or None on miss."""
+        with self._lock:
+            tbl = _result_table_name(source_hash, query_hash, variables_hash)
+            if not self._table_exists(tbl):
+                return None
+
+            try:
+                quoted_tbl = _q(tbl)
+                rows = self.conn.execute(
+                    f"""
+                    SELECT * EXCLUDE ({', '.join(sorted(_METADATA_COLS))})
+                    FROM {quoted_tbl}
+                    WHERE _snapshot_set_seq = (SELECT MAX(_snapshot_set_seq) FROM {quoted_tbl})
+                    """,
+                ).fetchall()
+                if not rows:
+                    return None
+                columns = [d[0] for d in self.conn.description]
+                return [dict(zip(columns, row, strict=False)) for row in rows]
+            except duckdb.CatalogException:
+                return None
+
+    def put(
+        self,
+        source_hash: str,
+        query_hash: str,
+        variables_hash: str,
+        data: list[dict[str, Any]],
+        *,
+        face_slug: str,
+        query_name: str,
+        variables: dict[str, Any],
+        is_full_refresh: bool = False,
+        primary_key: str | None = None,
+    ) -> None:
+        """Store rows under (source_hash, query_hash, variables_hash)."""
+        if not data:
+            return
+
+        with self._lock:
+            tbl = _result_table_name(source_hash, query_hash, variables_hash)
+            set_seq, run_seq = self._next_sequences(tbl, is_full_refresh)
+            self._ensure_result_table(tbl, data)
+            self._insert_rows(
+                tbl,
+                data,
+                set_seq,
+                run_seq,
+                source_hash,
+                variables,
+            )
+            self._update_view(tbl, primary_key)
+            self._upsert_face_query_index(
+                face_slug, query_name, source_hash, query_hash, variables_hash
+            )
+
+        logger.info(
+            "Cached %d rows to %s (set_seq=%d, run_seq=%d)",
+            len(data),
+            tbl,
+            set_seq,
+            run_seq,
+        )
+
+    def get_failure(
+        self, source_hash: str, query_hash: str, variables_hash: str
+    ) -> CachedQueryFailure | None:
+        """Return a cached failure if within TTL, else None."""
+        if self.failure_ttl_seconds == 0:
+            return None
+        with self._lock:
+            row = self.conn.execute(
+                """
+                SELECT error_class, error_message, traceback, failed_at
+                FROM _query_failures
+                WHERE source_hash = ? AND query_hash = ? AND variables_hash = ?
+                """,
+                [source_hash, query_hash, variables_hash],
+            ).fetchone()
+        if row is None:
+            return None
+        error_class, error_message, traceback, failed_at = row
+        elapsed = (datetime.now() - failed_at).total_seconds()
+        if elapsed > self.failure_ttl_seconds:
+            return None
+        return CachedQueryFailure(
+            error_class=error_class,
+            error_message=error_message,
+            traceback=traceback,
+            failed_at=failed_at,
+        )
+
+    def put_failure(
+        self,
+        source_hash: str,
+        query_hash: str,
+        variables_hash: str,
+        exception: Exception,
+        *,
+        face_slug: str,
+        query_name: str,
+    ) -> None:
+        """Store a query failure keyed by (source_hash, query_hash, variables_hash)."""
+        with self._lock:
+            self.conn.execute(
+                """
+                INSERT OR REPLACE INTO _query_failures
+                    (source_hash, query_hash, variables_hash,
+                     face_slug, query_name,
+                     error_class, error_message, traceback, failed_at)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+                """,
+                [
+                    source_hash,
+                    query_hash,
+                    variables_hash,
+                    face_slug,
+                    query_name,
+                    type(exception).__name__,
+                    str(exception),
+                    "".join(tb_mod.format_exception(exception)),
+                    datetime.now(),
+                ],
+            )
+
+    def clear(self, source_hash: str, query_hash: str, variables_hash: str) -> None:
+        """Remove success and failure entries for this key."""
+        with self._lock:
+            self.conn.execute(
+                "DELETE FROM _query_failures WHERE source_hash = ? AND query_hash = ? AND variables_hash = ?",
+                [source_hash, query_hash, variables_hash],
+            )
+            tbl = _result_table_name(source_hash, query_hash, variables_hash)
+            if self._table_exists(tbl):
+                self.conn.execute(f"DROP TABLE {_q(tbl)}")
+                view = _view_name_for_table(tbl)
+                self.conn.execute(f"DROP VIEW IF EXISTS {_q(view)}")
+
+    def close(self) -> None:
+        with self._lock:
+            self.conn.close()
+
+    def supports_results_refs(self) -> bool:
+        return True
+
+    def rewrite_results_refs(self, face_slug: str, sql: str) -> str:
+        """Rewrite {{ results.X }} to view names for this face."""
+        pattern = r"\{\{\s*results\.(\w+)\s*\}\}"
+
+        def replace(match: re.Match) -> str:
+            qname = match.group(1)
+            with self._lock:
+                view = self._view_for_face_query(face_slug, qname)
+            if view is None:
+                raise RuntimeError(
+                    f"{{ results.{qname} }} referenced but not cached for face '{face_slug}'"
+                )
+            return view
+
+        return re.sub(pattern, replace, sql)
+
+    def execute_results_query(self, face_slug: str, sql: str) -> list[dict[str, Any]]:
+        """Execute SQL against cached result views."""
+        with self._lock:
+            rows = self.conn.execute(sql).fetchall()
+            columns = [d[0] for d in self.conn.description]
+            return [dict(zip(columns, row, strict=False)) for row in rows]
+
+    # ─────────────────────────────────────────────────────────────────────
+    # Private helpers
+    # ─────────────────────────────────────────────────────────────────────
+
+    def _drop_legacy_tables(self) -> None:
+        """Drop tables that match the old (face_slug_q_query_name) schema.
+
+        Legacy tables are dev artifacts; we drop them on open to avoid
+        shape mismatch. Cache re-warms naturally on next query.
+        """
+        try:
+            tables = self.conn.execute(
+                """
+                SELECT table_name FROM information_schema.tables
+                WHERE table_type = 'BASE TABLE'
+                  AND table_name NOT LIKE '_r_%'
+                  AND table_name NOT LIKE '_query_%'
+                  AND table_name NOT LIKE '_face_%'
+                """
+            ).fetchall()
+        except Exception:  # noqa: BLE001
+            return
+
+        dropped = 0
+        for (tbl,) in tables:
+            # Old pattern: <face_slug>_q_<query_name> and related
+            if "_q_" in tbl or tbl.endswith("_current") or "_archive_" in tbl:
+                with contextlib.suppress(Exception):
+                    self.conn.execute(f"DROP TABLE IF EXISTS {_q(tbl)}")
+                    dropped += 1
+        # Also drop old views
+        with contextlib.suppress(Exception):
+            views = self.conn.execute(
+                """
+                SELECT table_name FROM information_schema.tables
+                WHERE table_type = 'VIEW' AND table_name NOT LIKE '_r_%'
+                """
+            ).fetchall()
+            for (v,) in views:
+                with contextlib.suppress(Exception):
+                    self.conn.execute(f"DROP VIEW IF EXISTS {_q(v)}")
+
+        if dropped:
+            logger.info(
+                "legacy cache cleared on schema change (%d tables dropped)", dropped
+            )
+
+    def _ensure_schema(self) -> None:
+        """Create the system tables if they don't exist."""
+        self.conn.execute(
+            """
+            CREATE TABLE IF NOT EXISTS _query_failures (
+                source_hash TEXT NOT NULL,
+                query_hash TEXT NOT NULL,
+                variables_hash TEXT NOT NULL,
+                face_slug TEXT NOT NULL,
+                query_name TEXT NOT NULL,
+                error_class TEXT NOT NULL,
+                error_message TEXT NOT NULL,
+                traceback TEXT,
+                failed_at TIMESTAMP NOT NULL,
+                PRIMARY KEY (source_hash, query_hash, variables_hash)
+            )
+            """
+        )
+        self.conn.execute(
+            """
+            CREATE TABLE IF NOT EXISTS _face_query_index (
+                face_slug TEXT NOT NULL,
+                query_name TEXT NOT NULL,
+                source_hash TEXT NOT NULL,
+                query_hash TEXT NOT NULL,
+                variables_hash TEXT NOT NULL,
+                created_at TIMESTAMP NOT NULL,
+                PRIMARY KEY (face_slug, query_name, variables_hash)
+            )
+            """
+        )
+
+    def _table_exists(self, name: str) -> bool:
+        try:
+            result = self.conn.execute(
+                "SELECT COUNT(*) FROM information_schema.tables WHERE table_name = ?",
+                [name],
+            ).fetchone()
+            return bool(result and result[0] > 0)
+        except duckdb.CatalogException:
+            return False
+
+    def _next_sequences(
+        self, table_name: str, is_full_refresh: bool
+    ) -> tuple[int, int]:
+        if not self._table_exists(table_name):
+            return (1, 1)
+        try:
+            row = self.conn.execute(
+                f"""
+                SELECT COALESCE(MAX(_snapshot_set_seq), 0),
+                       COALESCE(MAX(_snapshot_run_seq), 0)
+                FROM {_q(table_name)}
+                """
+            ).fetchone()
+            if row is None:
+                return (1, 1)
+            set_seq, run_seq = (row[0] or 0), (row[1] or 0)
+            if is_full_refresh:
+                return (set_seq + 1, run_seq + 1)
+            return (set_seq or 1, run_seq + 1)
+        except duckdb.CatalogException:
+            return (1, 1)
+
+    def _ensure_result_table(self, table_name: str, data: list[dict[str, Any]]) -> None:
+        if self._table_exists(table_name):
+            return
+        sample = data[0]
+        cols = [
+            "_snapshot_set_seq INTEGER",
+            "_snapshot_run_seq INTEGER",
+            "_run_timestamp TIMESTAMP",
+            "_source_hash VARCHAR",
+            "_variables JSON",
+        ]
+        for col_name, value in sample.items():
+            col_type = _infer_type(value)
+            cols.append(f'"{col_name}" {col_type}')
+        self.conn.execute(f"CREATE TABLE {_q(table_name)} ({', '.join(cols)})")
+        self.conn.execute(
+            f"""
+            CREATE INDEX {_q(f'idx_{table_name}')}
+            ON {_q(table_name)} (_snapshot_set_seq, _snapshot_run_seq)
+            """
+        )
+
+    def _insert_rows(
+        self,
+        table_name: str,
+        data: list[dict[str, Any]],
+        set_seq: int,
+        run_seq: int,
+        source_hash: str,
+        variables: dict[str, Any],
+    ) -> None:
+        n = len(data)
+        now = datetime.now()
+
+        if HAS_PYARROW:
+            metadata = {
+                "_snapshot_set_seq": pa.array([set_seq] * n, type=pa.int64()),
+                "_snapshot_run_seq": pa.array([run_seq] * n, type=pa.int64()),
+                "_run_timestamp": pa.array([now] * n, type=pa.timestamp("us")),
+                "_source_hash": pa.array([source_hash] * n),
+                "_variables": pa.array([json.dumps(variables)] * n),
+            }
+            data_cols = {
+                k: _cache_safe_arrow_array([row.get(k) for row in data])
+                for k in data[0]
+            }
+            arrow_table = pa.table({**metadata, **data_cols})
+            self.conn.register("_df_insert_buf", arrow_table)
+            self.conn.execute(
+                f'INSERT INTO {_q(table_name)} SELECT * FROM "_df_insert_buf"'
+            )
+            self.conn.unregister("_df_insert_buf")
+        else:
+            for row in data:
+                values = [
+                    set_seq,
+                    run_seq,
+                    now,
+                    source_hash,
+                    json.dumps(variables),
+                ]
+                values.extend(row.get(k) for k in data[0])
+                placeholders = ", ".join("?" * len(values))
+                self.conn.execute(
+                    f"INSERT INTO {_q(table_name)} VALUES ({placeholders})", values
+                )
+
+    def _update_view(self, table_name: str, primary_key: str | None = None) -> None:
+        view = _view_name_for_table(table_name)
+        self.conn.execute(f"DROP VIEW IF EXISTS {_q(view)}")
+
+        all_cols = self._get_columns(table_name)
+        data_cols = sorted(all_cols - _METADATA_COLS)
+        if not data_cols:
+            return
+
+        exclude = ", ".join(sorted(_METADATA_COLS))
+        quoted_tbl = _q(table_name)
+
+        if primary_key and primary_key in data_cols:
+            quoted_pk = _q(primary_key)
+            sql = f"""
+                CREATE VIEW {_q(view)} AS
+                SELECT * EXCLUDE ({exclude})
+                FROM {quoted_tbl}
+                WHERE _snapshot_set_seq = (SELECT MAX(_snapshot_set_seq) FROM {quoted_tbl})
+                QUALIFY ROW_NUMBER() OVER (
+                    PARTITION BY {quoted_pk} ORDER BY _snapshot_run_seq DESC
+                ) = 1
+            """
+        else:
+            sql = f"""
+                CREATE VIEW {_q(view)} AS
+                SELECT * EXCLUDE ({exclude})
+                FROM {quoted_tbl}
+                WHERE _snapshot_set_seq = (SELECT MAX(_snapshot_set_seq) FROM {quoted_tbl})
+            """
+        self.conn.execute(sql)
+
+    def _get_columns(self, table_name: str) -> set[str]:
+        try:
+            rows = self.conn.execute(
+                "SELECT column_name FROM information_schema.columns WHERE table_name = ?",
+                [table_name],
+            ).fetchall()
+            return {r[0] for r in rows}
+        except duckdb.CatalogException:
+            return set()
+
+    def _upsert_face_query_index(
+        self,
+        face_slug: str,
+        query_name: str,
+        source_hash: str,
+        query_hash: str,
+        variables_hash: str,
+    ) -> None:
+        self.conn.execute(
+            """
+            INSERT OR REPLACE INTO _face_query_index
+                (face_slug, query_name, source_hash, query_hash, variables_hash, created_at)
+            VALUES (?, ?, ?, ?, ?, ?)
+            """,
+            [
+                face_slug,
+                query_name,
+                source_hash,
+                query_hash,
+                variables_hash,
+                datetime.now(),
+            ],
+        )
+
+    def _view_for_face_query(self, face_slug: str, query_name: str) -> str | None:
+        """Resolve (face_slug, query_name) → view name via _face_query_index."""
+        row = self.conn.execute(
+            """
+            SELECT source_hash, query_hash, variables_hash
+            FROM _face_query_index
+            WHERE face_slug = ? AND query_name = ?
+            ORDER BY created_at DESC LIMIT 1
+            """,
+            [face_slug, query_name],
+        ).fetchone()
+        if row is None:
+            return None
+        source_hash, query_hash, variables_hash = row
+        tbl = _result_table_name(source_hash, query_hash, variables_hash)
+        return _view_name_for_table(tbl)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Private module-level helpers
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def _q(identifier: str) -> str:
+    """Quote a SQL identifier (already assumed safe)."""
+    escaped = identifier.replace('"', '""')
+    return f'"{escaped}"'
+
+
+# Maximum DECIMAL precision DuckDB supports (HUGEINT-backed). BigQuery NUMERIC
+# can return values with higher precision (e.g. 47,38) which crash duckdb's
+# arrow-table register with `Unsupported Internal Arrow Type for Decimal`.
+_DUCKDB_MAX_DECIMAL_PRECISION = 38
+
+
+def _cache_safe_arrow_array(values: list[Any]) -> "pa.Array":
+    """Build a pyarrow array, downcasting too-precise decimals to float64.
+
+    BigQuery NUMERIC columns can have precision > 38 (e.g. NUMERIC(47, 38) on
+    fan-out-deduplicated SUM aggregates). pyarrow auto-infers ``decimal256``
+    for these, but DuckDB's DECIMAL is HUGEINT-backed and capped at precision
+    38; ``conn.register`` then raises ``NotImplementedException`` when it
+    encounters the over-precise type. Visualization downstream is fine with
+    float64 (15 significant digits is plenty for dollar amounts).
+    """
+    arr = pa.array(values)
+    arrow_type = arr.type
+    if (
+        pa.types.is_decimal(arrow_type)
+        and arrow_type.precision > _DUCKDB_MAX_DECIMAL_PRECISION
+    ):
+        return arr.cast(pa.float64())
+    return arr
+
+
+def _infer_type(value: Any) -> str:
+    if value is None:
+        return "VARCHAR"
+    if isinstance(value, bool):
+        return "BOOLEAN"
+    if isinstance(value, int):
+        return "BIGINT"
+    if isinstance(value, float):
+        return "DOUBLE"
+    if isinstance(value, datetime):
+        return "TIMESTAMP"
+    if isinstance(value, (list, dict)):
+        return "JSON"
+    return "VARCHAR"
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Entry-point helper
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def open_cache_from_env() -> DuckDBCache | None:
+    """Open a DuckDBCache from the DFT_CACHE_PATH environment variable.
+
+    Returns None when the variable is unset.
+    Raises ValueError when set but the path does not exist.
+    """
+    path_str = os.getenv("DFT_CACHE_PATH")
+    if not path_str:
+        return None
+    path = Path(path_str)
+    if not path.exists():
+        raise ValueError(
+            f"DFT_CACHE_PATH points to a non-existent file: {path}. "
+            "Create the cache file first (construct DuckDBCache(db_path=...) and populate it) "
+            "or unset the variable."
+        )
+    return DuckDBCache(db_path=path)