"""DuckDB demo database: built from bundled schema + seed SQL; row-capped execution.""" from __future__ import annotations import threading from dataclasses import dataclass from pathlib import Path import duckdb DATA_DIR = Path(__file__).resolve().parents[1] / "data" DEFAULT_MAX_ROWS = 200 DEFAULT_TIMEOUT_S = 5.0 @dataclass class QueryResult: columns: list[str] rows: list[tuple] truncated: bool = False def _statements(script: str): for raw in script.split(";"): s = raw.strip() if s and not s.startswith("--"): yield s def new_connection() -> "duckdb.DuckDBPyConnection": """A fresh in-memory DuckDB with bounded resources (memory + threads).""" con = duckdb.connect(":memory:") # Bound resource use so an accepted-but-pathological query can't exhaust the host. for pragma in ("SET memory_limit='512MB'", "SET threads TO 2"): try: con.execute(pragma) except Exception: # noqa: BLE001 - config name varies across DuckDB versions pass return con def build_connection(schema_path=None, seed_path=None) -> "duckdb.DuckDBPyConnection": """Build an in-memory DuckDB seeded from the bundled schema + seed SQL files.""" con = new_connection() schema_sql = Path(schema_path or DATA_DIR / "schema.sql").read_text(encoding="utf-8") seed_sql = Path(seed_path or DATA_DIR / "seed.sql").read_text(encoding="utf-8") for stmt in _statements(schema_sql): con.execute(stmt) for stmt in _statements(seed_sql): con.execute(stmt) return con def execute( con, sql: str, max_rows: int = DEFAULT_MAX_ROWS, timeout_s: float = DEFAULT_TIMEOUT_S, ) -> QueryResult: """Run a query and return at most ``max_rows`` rows. The cost cap has two parts: a wall-clock ``timeout_s`` (a watchdog calls ``con.interrupt()``) bounds execution time, and ``max_rows`` caps the result size. Resource limits are also set on the connection in ``build_connection``. """ timer = None if timeout_s and hasattr(con, "interrupt"): timer = threading.Timer(timeout_s, con.interrupt) timer.start() try: cur = con.execute(sql) columns = [d[0] for d in cur.description] if cur.description else [] fetched = cur.fetchmany(max_rows + 1) finally: if timer is not None: timer.cancel() truncated = len(fetched) > max_rows return QueryResult(columns=columns, rows=fetched[:max_rows], truncated=truncated)