From 53d208cc43297fb9c745b4f7691859ca317ab326 Mon Sep 17 00:00:00 2001 From: Laith Al-Saadoon Date: Sun, 13 Sep 2026 01:05:09 +0000 Subject: [PATCH 1/2] fix(query): size registration to the host, private spill dir, no root, no extension installs at query time Closes the five open findings of the MicroVM defensive review of main 4571e14 and pins findings 6 and 7 (closed by #8) with their exact reproductions. 1. Registration OOM under ~1.9 GiB RAM per core. The eager read_json readers reserve about 2x maximum_object_size per thread, so the 1 GiB constant cost 2 GiB a thread and `SET threads=4; SET memory_limit='6GB'` failed to register any corpus with two or more edges.jsonl files. The bound is now sized from the largest file the reader opens (plus a quarter and 1 MiB, floored at DuckDB's 16 MiB default, capped at 1 GiB), and query applies memory_limit, threads and temp_directory BEFORE register(). The cap derives from available RAM (MemAvailable, 80% ceilings, 8 GiB target, 512 MiB floor) and the thread count from the cap (one per 2 GiB, capped at sched_getaffinity). ATIF_SQL_QUERY_MEMORY_LIMIT and ATIF_SQL_QUERY_THREADS override both; malformed values exit 64. The old _query_memory_limit_bytes is the same function, re-derived. Panel on the 300-session corpus: byte-identical output unlimited and under 6GB/4 threads; unlimited wall 0.88-0.98 s -> 0.73-0.79 s. 2. Root could overwrite the corpus through a granted parquet with COPY ... (USE_TMP_FILE false). DuckDB 1.5.5 has no read-only grant (measured: no write switch in duckdb_settings(), and a read_parquet relation bound before enable_external_access=false is refused without a grant), so two layers: query, search and analyze refuse uid 0 with exit 77 root_refused unless ATIF_SQL_ALLOW_ROOT=1 (logged warning), and query refuses every file-facing statement kind before executing anything (COPY, COPY_DATABASE, EXPORT, ATTACH, DETACH, INSTALL, LOAD, PREPARE, EXECUTE) via DuckDB's own extract_statements on the hardened connection, exit 70 sandbox_refused, for any uid. Probe d2 now fails as any uid. 3. The spill directory was /.duckdb_tmp and writable by caller SQL, persisting files inside the corpus. It is now a per-process tempfile.mkdtemp (0700) under the system temp dir, the sole allowed_directories entry, removed on every exit path after the connection closes. Nothing under the corpus is writable. 4. register_vss ran INSTALL lance (a 242 MB download) during every registration. It now LOADs the extension only when duckdb_extensions() says it is installed, skips the load entirely when no store directory exists, and binds message_embeddings empty with a warning otherwise. query sets autoinstall_known_extensions=false and autoload_known_extensions=false before registration. Installing is an explicit act: `atif-sql embed --install-extension`, or a real embed run. `atif-sql status` reports vector_search ready | no_store | extension_missing; `search` exits 78 extension_missing rather than pretending the store is empty. 5. duckdb_settings() / current_setting() expose the grants and so every session id. DuckDB cannot hide a setting and the grants must be per file, so this is documented as accepted: the caller is the local user and query is not a privilege boundary. Findings 6 and 7: a transcript named `..jsonl` (stem `.`) and one named `*.jsonl`, materialized over three victims, are rejected and the tree is untouched; a session dir named `*` or `???` counts every session once on both read paths and grants no glob. Tests that fail on main 4571e14 and pass here (verified on a scratch checkout): registration under 4 threads/6GB and 16 threads/8GB (CLI and registry), COPY into the spill dir persists nothing, USE_TMP_FILE false over a granted parquet is refused, root is refused, and a corpus with a store and no extension installs nothing. --- AGENTS.md | 11 + README.md | 9 + SECURITY.md | 10 +- docs/CONTRACT.md | 12 + docs/architecture/data-flow.md | 21 +- docs/behavior/processes.md | 20 +- docs/reference/cli.md | 8 +- packages/atif-cli/src/atif_cli/app.py | 607 ++++++++++-- packages/atif-cli/src/atif_cli/errors.py | 9 + packages/atif-cli/tests/cli_fixtures.py | 15 + packages/atif-cli/tests/test_app.py | 871 ++++++++++++++++-- packages/atif-cli/tests/test_vss_commands.py | 73 ++ packages/atif-corpus/tests/test_session_id.py | 42 + .../src/atif_duck/infrastructure/registry.py | 144 ++- packages/atif-duck/tests/duck_fixtures.py | 15 + .../tests/test_registration_limits.py | 240 +++++ .../tests/test_sql_text_boundaries.py | 47 + 17 files changed, 1943 insertions(+), 211 deletions(-) create mode 100644 packages/atif-duck/tests/test_registration_limits.py diff --git a/AGENTS.md b/AGENTS.md index ab6571a..851f7bc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -141,6 +141,17 @@ For an LLM agent driving `atif-sql`, the discovery loop is three commands: prefer `atif-sql search 'query'` — it embeds the text first, then runs the same `semantic_search` kNN. + What the query sandbox will and won't do: `SELECT`, `EXPLAIN`, `SET + TimeZone`, and in-memory DDL/DML run; `COPY`, `EXPORT`, `ATTACH`, + `DETACH`, `INSTALL`, `LOAD`, `PREPARE` and `EXECUTE` exit 70 with kind + `sandbox_refused` before anything executes, so emit results on stdout + rather than writing files. Running as uid 0 exits 77 (`root_refused`) + unless `ATIF_SQL_ALLOW_ROOT=1`. The connection is sized to the host + before registration (`ATIF_SQL_QUERY_MEMORY_LIMIT` / `ATIF_SQL_QUERY_THREADS` + override), and no extension is ever installed at query time: `atif-sql + status` reports `vector_search` as `ready`, `no_store`, or + `extension_missing` (fix the last with `atif-sql embed --install-extension`). + Adding a view/macro? The drift tests force: a `DESCRIPTIONS` entry in `atif_duck/domain/catalog.py`, an `ARG_EXEMPLARS` entry for any new parameter name, `TABLE_MACRO_NAMES` membership if the DDL is `AS TABLE`, and diff --git a/README.md b/README.md index 6fb6966..d8ad495 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,15 @@ not in a public issue. Supported versions, the disclosure expectations, and what vulnerability in a tool that reads local transcripts are in **[SECURITY.md](SECURITY.md)**. +`atif-sql query` runs agent-composed SQL, so it runs it in a box: sized to the host before the +corpus is registered (override with `ATIF_SQL_QUERY_MEMORY_LIMIT` and `ATIF_SQL_QUERY_THREADS`), +a private spill directory outside the corpus that's removed on exit, no extension installs at +query time (`atif-sql embed --install-extension` is where the lance extension comes from), file +facing statements (`COPY`, `EXPORT`, `ATTACH`, `INSTALL`, `LOAD`, `PREPARE`, `EXECUTE`) refused +before they run, and a refusal to run as root unless `ATIF_SQL_ALLOW_ROOT=1` says so. The details +and the one accepted disclosure (`duckdb_settings()` lists the granted paths) are in +[docs/reference/cli.md](docs/reference/cli.md#query). + ## License [Apache License 2.0](LICENSE). Each of the seven module directories carries the same license diff --git a/SECURITY.md b/SECURITY.md index 08a2140..acc88e2 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -55,4 +55,12 @@ If your transcripts may not be sent to a third-party model provider, do not run DuckDB can read and write local files. The tool assumes the caller already owns the shell and the data; do not treat it as a sandbox, and do not wire it behind an interface that lets an untrusted party choose the SQL, the corpus root, or -the environment. +the environment. What the hardened connection does do is keep an injected +statement from writing the corpus or reaching the network: file-facing +statement kinds (`COPY`, `EXPORT`, `ATTACH`, `INSTALL`, `LOAD`, `PREPARE`, +`EXECUTE`) are refused before execution, the only granted directory is a +private per-process spill dir, no extension is ever installed at query time, +and `query`, `search` and `analyze` refuse to run as root (file modes don't +bind uid 0) unless `ATIF_SQL_ALLOW_ROOT=1` is set. Caller SQL can still read +`duckdb_settings()`, which lists the granted paths and so every session id; +that's the local user's own corpus listing, not a leak across a boundary. diff --git a/docs/CONTRACT.md b/docs/CONTRACT.md index e4e3dfe..adc437f 100644 --- a/docs/CONTRACT.md +++ b/docs/CONTRACT.md @@ -159,6 +159,15 @@ atif-sql materialize [--force] [--quiesce-seconds N] [--agent ...] [--workers N] # report adds rejected / rejected_session_ids atif-sql status [--agent ...] # corpus freshness, counts, watermark age atif-sql query 'SQL' [--format auto|json|csv] + # sized to the host before registration + # (ATIF_SQL_QUERY_MEMORY_LIMIT / _THREADS + # override); private mkdtemp spill dir; + # refuses uid 0 unless ATIF_SQL_ALLOW_ROOT=1; + # COPY/EXPORT/ATTACH/INSTALL/LOAD/PREPARE/EXECUTE + # exit 70 sandbox_refused before execution; + # never installs an extension +atif-sql embed --install-extension # the ONE place the lance DuckDB extension + # is downloaded (also done by a real embed run) atif-sql schema # static, <50ms, no duckdb bind ## Parity oracle (satisfied and retired) @@ -181,6 +190,9 @@ repo neither declares nor provides. source_root (default CLAUDE_CONFIG_DIR~/.claude /projects), corpus_root, quiesce_seconds=300, agent=claude-code (every --agent command reads it), materialize_workers=min(8, cpu_count) (materialize's pool size; 1 = single process). +ATIF_SQL_QUERY_MEMORY_LIMIT (DuckDB size literal, e.g. 6GB) and ATIF_SQL_QUERY_THREADS +override query's host-derived cap and thread count; ATIF_SQL_ALLOW_ROOT=1 lets +query/search/analyze run as uid 0 (a warning is logged). _default_*() factories read env at call time. With agent=codex the two roots re-derive to $CODEX_HOME (default ~/.codex)/sessions and ~/.atif-sql/corpus/codex; an explicitly set ATIF_SQL_SOURCE_ROOT or ATIF_SQL_CORPUS_ROOT always wins over that diff --git a/docs/architecture/data-flow.md b/docs/architecture/data-flow.md index f077c9a..f870706 100644 --- a/docs/architecture/data-flow.md +++ b/docs/architecture/data-flow.md @@ -103,12 +103,19 @@ sequenceDiagram written by a different provider or width raises instead of binding, because vectors from different models live in incompatible spaces and would return numerically valid but meaningless cosine scores — `packages/atif-duck/src/atif_duck/domain/embedding_guard.py:46`. -7. The fully-registered connection is then sandboxed: spill directory, memory cap, a directory - allowlist holding only the spill area, a path allowlist holding the individual analytics - parquets, `enable_external_access=false`, and `lock_configuration` last so caller SQL cannot - widen any of it — `packages/atif-cli/src/atif_cli/app.py:138`. -8. The caller's statement executes against the locked connection (`:606`) and the cursor drains in - batches to stdout — a JSON array of row objects on a pipe, a width-aligned table on a TTY — `packages/atif-cli/src/atif_cli/output.py:154`. +7. Before any of that, the connection was sized to the host (`_configure_query_resources`: a memory + cap from available RAM, a thread count from that cap, a private `mkdtemp` spill directory, and + extension auto-install and auto-load off), because registration is what needs the cap. The + fully-registered connection is then sandboxed: a directory allowlist holding only the private + spill area, a path allowlist holding the individual parquets the views read lazily, + `enable_external_access=false`, and `lock_configuration` last so caller SQL cannot widen any of + it (`packages/atif-cli/src/atif_cli/app.py`, `_harden_query_connection`). +8. The statement's kinds are checked with DuckDB's own parser on the locked connection; `COPY`, + `EXPORT`, `ATTACH`, `DETACH`, `INSTALL`, `LOAD`, `PREPARE` and `EXECUTE` exit 70 + (`sandbox_refused`) before anything runs. Then the caller's statement executes and the cursor + drains in batches to stdout: a JSON array of row objects on a pipe, a width-aligned table on a + TTY (`packages/atif-cli/src/atif_cli/output.py:154`). The spill directory is removed when the + process exits. ```mermaid sequenceDiagram @@ -119,6 +126,7 @@ sequenceDiagram participant Lance as Lance store CLI->>DB: duckdb.connect() + CLI->>DB: threads, memory_limit, private temp_directory, autoinstall off CLI->>Duck: register(con, corpus_root, expected model + dim) Duck->>DB: CREATE TEMP TABLE raw readers over corpus globs DB->>Disk: read_json sessions/*/meta.json then the rest @@ -130,6 +138,7 @@ sequenceDiagram Duck->>Duck: ensure_store_matches, then bind the view Duck-->>CLI: views and macros registered CLI->>DB: allowlists, external access off, lock_configuration + CLI->>DB: extract_statements(caller SQL), refuse file-facing kinds CLI->>DB: execute(caller SQL) DB-->>CLI: cursor CLI->>CLI: drain in batches to stdout diff --git a/docs/behavior/processes.md b/docs/behavior/processes.md index 7f7eb40..2b814eb 100644 --- a/docs/behavior/processes.md +++ b/docs/behavior/processes.md @@ -270,14 +270,18 @@ Entry point: `packages/atif-cli/src/atif_cli/app.py:529` 5. Create the macros, then the analytics views and analytics macros over the analytics parquets, which bind against both the parquets and the base views — `:1005`, `packages/atif-duck/src/atif_duck/infrastructure/analytics.py:124`. -6. Harden the connection in a fixed order: temp directory, memory cap, a - directory allowlist holding only the spill area, a file allowlist of the - analytics parquets, the config exemption list, then - `enable_external_access=false` and `lock_configuration=true` last — - `packages/atif-cli/src/atif_cli/app.py:138`. -7. Execute the caller's statement and stream the cursor: a plain table on a - TTY, a JSON array of row objects on a pipe — `:606`, - `packages/atif-cli/src/atif_cli/output.py:154`. +6. Harden the connection in a fixed order (the memory cap, thread count and + private spill directory were set before step 3, since registration is what + needs them): a directory allowlist holding only the spill area, a file + allowlist of the parquets the views read lazily, the config exemption + list, then `enable_external_access=false` and `lock_configuration=true` + last (`packages/atif-cli/src/atif_cli/app.py`, `_harden_query_connection`). +7. Check the statement's kinds with DuckDB's parser and refuse `COPY`, + `EXPORT`, `ATTACH`, `DETACH`, `INSTALL`, `LOAD`, `PREPARE` and `EXECUTE` + before anything runs (exit 70, `sandbox_refused`); then execute the + caller's statement and stream the cursor: a plain table on a TTY, a JSON + array of row objects on a pipe (`packages/atif-cli/src/atif_cli/output.py:154`). + The spill directory is removed on exit. 8. Classify any DuckDB failure into parse, catalog, or runtime — or an embedding-provider mismatch — and exit 64, 65, or 70 with a JSON error envelope — `packages/atif-cli/src/atif_cli/duck_errors.py:34`, `:66`. diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 5c40cdb..d484ecf 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -100,13 +100,17 @@ Flags: - `--corpus-root` — override the materialized corpus root. `:504` - `--format` — `table` on a TTY, a JSON array of row objects on a pipe. `:505` -The statement runs against a hardened connection: reads reach the registered views and nothing else, and the only writable path is the query engine's own spill directory `/.duckdb_tmp`. `:601` +The statement runs against a hardened connection: reads reach the registered views and nothing else, and nothing under the corpus root is writable. Before registration the connection is sized to the host: a memory cap derived from available RAM (half of physical RAM or 8 GiB, whichever is larger, never above 80% of what's available) and a thread count of one per 2 GiB of that cap, capped at the CPUs the process may use. `ATIF_SQL_QUERY_MEMORY_LIMIT` (a DuckDB size such as `6GB`) and `ATIF_SQL_QUERY_THREADS` override both; a malformed value exits 64. The spill directory is a private `mkdtemp` (mode 0700) under the system temp dir, the only directory the sandbox grants, and it's removed when the process exits. Extension auto-install and auto-load are off, and the lance extension is loaded only when it's already installed, so registration never reaches the network (`packages/atif-cli/src/atif_cli/app.py`, `_configure_query_resources`). + +Two layers keep caller SQL from writing the corpus. DuckDB's file grants are read-write and it has no read-only grant, so `COPY ... TO (USE_TMP_FILE false)` would overwrite one; the CLI therefore refuses every statement kind that names a file before executing anything, using DuckDB's own parser: `COPY`, `EXPORT`, `ATTACH`, `DETACH`, `INSTALL`, `LOAD`, `PREPARE` and `EXECUTE` exit 70 with kind `sandbox_refused`, for any uid, and a batch containing one of them runs nothing. And because a `0444` file mode doesn't bind root, `query` refuses to run as uid 0 (exit 77, kind `root_refused`) unless `ATIF_SQL_ALLOW_ROOT=1` is set, which logs a warning. `search` and `analyze` refuse root the same way. + +What caller SQL can still see: `duckdb_settings()` and `current_setting(...)` return the sandbox's own configuration, including the corpus root, the spill directory, the memory cap and every granted parquet path, which names every session id. DuckDB can't hide a setting from SQL and the grants have to be per file, so this is accepted: the caller is the local user, who can list the corpus and `SELECT session_id FROM sessions` anyway, and `query` isn't a privilege boundary (see SECURITY.md). The views themselves carry no corpus path as statement text. The registry hands its globs and file lists to `read_json(?)` as bound parameters and builds the parquet readers through DuckDB's relation API, so a corpus root such as `o'brien ?; --$1` and transcript content carrying SQL text both register as data (`packages/atif-duck/src/atif_duck/infrastructure/registry.py`). A session directory whose name fails the session id boundary (`packages/atif-duck/src/atif_duck/domain/session_id.py`) registers nothing and is logged once. Sessions that carry current columnar artifacts are served from their parquet files, so no JSON is parsed for them at query time; the rest are read from `trajectory.json`, and the views union the two. The per-session parquet files the registry bound are granted to the sandbox the same way the analytics parquets are (as individual `allowed_paths` entries, `packages/atif-cli/src/atif_cli/app.py:221`), and they're written read-only (`0444`), so a `COPY ... TO` at one of them fails at the filesystem even though DuckDB's grant is read-write. `atif-sql status` says which path a corpus takes. -Exit codes: `64` parse error, `65` catalog error, `65` embedding mismatch, `70` runtime error. `:550` +Exit codes: `64` parse error or a malformed `ATIF_SQL_QUERY_*` override, `65` catalog error, `65` embedding mismatch, `70` runtime error, `70` `sandbox_refused` (a statement kind the sandbox never runs), `77` `root_refused`. `:550` ## analyze diff --git a/packages/atif-cli/src/atif_cli/app.py b/packages/atif-cli/src/atif_cli/app.py index 79c99a4..457d135 100644 --- a/packages/atif-cli/src/atif_cli/app.py +++ b/packages/atif-cli/src/atif_cli/app.py @@ -54,7 +54,12 @@ import json import os +import re +import shutil import sys +import tempfile +from contextlib import contextmanager +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Annotated, Any @@ -72,7 +77,7 @@ ) if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Generator, Sequence from atif_corpus.application.materialize import MaterializationReport @@ -121,25 +126,269 @@ def _sql_str(value: str) -> str: return f"'{escaped}'" -#: Lower bound for ``query``'s DuckDB memory cap. Registration alone parses -#: the corpus through 1 GiB-per-object JSON buffers on every thread, and the -#: base views (``tool_calls``, ``tool_rank``) need several GiB more on a -#: multi-GB corpus — a tighter cap turns working queries into OutOfMemory. -_QUERY_MEMORY_FLOOR_BYTES = 8 * 1024**3 +#: Env override for the DuckDB memory cap ``query`` runs under: any size +#: literal DuckDB's ``SET memory_limit`` accepts (``6GB``, ``512MiB``, ``2000000000B``). +QUERY_MEMORY_LIMIT_ENV = "ATIF_SQL_QUERY_MEMORY_LIMIT" + +#: Env override for the DuckDB thread count ``query`` runs under. +QUERY_THREADS_ENV = "ATIF_SQL_QUERY_THREADS" + +#: Set to ``1`` to let ``query``, ``search`` and ``analyze`` run as uid 0. +ALLOW_ROOT_ENV = "ATIF_SQL_ALLOW_ROOT" + +#: Target for ``query``'s DuckDB memory cap when the host can afford it: the +#: base views (``tool_calls``, ``tool_rank``) need several GiB on a multi-GB +#: corpus, and a tighter cap turns working queries into OutOfMemory. A +#: target, not a floor: it never exceeds what the host has (the old 8 GiB +#: floor set a "limit" above physical RAM on an 8 GiB guest). +_QUERY_MEMORY_TARGET_BYTES = 8 * 1024**3 + +#: Never cap below this. DuckDB needs room to open the readers at all, and a +#: host with less available than this fails to register whatever the cap says. +_QUERY_MEMORY_MIN_BYTES = 512 * 1024**2 + +#: Cap budgeted per DuckDB thread when deriving the thread count. The JSON +#: readers reserve about twice their ``maximum_object_size`` per thread, and +#: that bound is now sized from the largest file present (the largest +#: trajectory.json seen is 436 MB), so 2 GiB a thread keeps a corpus on the +#: JSON path registering under the cap. +_QUERY_BYTES_PER_THREAD = 2 * 1024**3 + + +def _host_memory() -> tuple[int, int]: + """``(physical, available)`` bytes on this host. + + Physical comes from ``os.sysconf``. Available is Linux's ``MemAvailable`` + from ``/proc/meminfo`` (what a new allocation can really get, reclaimable + page cache included); where that file is absent (macOS) available is + taken as physical. + """ + physical = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") + available = physical + try: + with Path("/proc/meminfo").open(encoding="ascii") as handle: + for line in handle: + if line.startswith("MemAvailable:"): + available = int(line.split()[1]) * 1024 + break + except (OSError, ValueError, IndexError): + pass + return physical, min(physical, available) def _query_memory_limit_bytes() -> int: - """Bytes to cap ``query``'s DuckDB heap at. + """Bytes to cap ``query``'s DuckDB heap at, derived from the host. + + Half of physical RAM or :data:`_QUERY_MEMORY_TARGET_BYTES`, whichever is + larger, but never above 80% of physical RAM (DuckDB's own default) and + never above 80% of the memory available right now, floored at + :data:`_QUERY_MEMORY_MIN_BYTES`. On a 124 GiB host that is 62 GiB; on an + 8 GiB guest with 6 GiB free it is 4.8 GiB, where the previous rule said + 8 GiB and DuckDB's default said 6.4 GiB, neither of which the guest had. + """ + physical, available = _host_memory() + target = max(int(physical * 0.5), _QUERY_MEMORY_TARGET_BYTES) + ceiling = min(int(physical * 0.8), int(available * 0.8)) + return max(_QUERY_MEMORY_MIN_BYTES, min(target, ceiling)) + + +def _query_threads(memory_limit_bytes: int) -> int: + """DuckDB threads for ``query``: one per :data:`_QUERY_BYTES_PER_THREAD` of cap. - Half of physical RAM, but never below :data:`_QUERY_MEMORY_FLOOR_BYTES` - and never above DuckDB's own 80%-of-RAM default — on a small host the - default is already the tighter of the two, and raising it would make an - unbounded query worse rather than better. + Capped at the CPUs this process may run on (``sched_getaffinity``, which + sees a container's cpuset where ``cpu_count`` does not) and floored at + one. DuckDB's own default is the core count, which is what made a 4 vCPU + guest reserve four readers' worth of memory it did not have. """ - import os + affinity = getattr(os, "sched_getaffinity", None) + cpus = len(affinity(0)) if affinity is not None else (os.cpu_count() or 1) + return max(1, min(cpus, memory_limit_bytes // _QUERY_BYTES_PER_THREAD)) + + +_SIZE_UNITS: dict[str, int] = { + "B": 1, + "KB": 10**3, + "MB": 10**6, + "GB": 10**9, + "TB": 10**12, + "KIB": 1024, + "MIB": 1024**2, + "GIB": 1024**3, + "TIB": 1024**4, +} + +_SIZE_RE = re.compile(r"\s*(\d+(?:\.\d+)?)\s*([A-Za-z]*)\s*") + + +def _parse_size(text: str) -> int: + """Bytes for a DuckDB-style size literal (``6GB``, ``512 MiB``, ``123B``).""" + match = _SIZE_RE.fullmatch(text) + if match is None: + problem = f"{text!r} is not a size (expected e.g. 6GB, 512MiB, 2000000000B)" + raise ValueError(problem) + number, unit = match.group(1), (match.group(2) or "B").upper() + if unit not in _SIZE_UNITS: + problem = f"{text!r} has an unknown unit (use B, KB, MB, GB, TB, KiB, MiB, GiB, TiB)" + raise ValueError(problem) + return int(float(number) * _SIZE_UNITS[unit]) + + +@dataclass(frozen=True, slots=True) +class QueryResources: + """What ``query`` hands DuckDB before it registers anything.""" + + memory_limit_bytes: int + threads: int + + +def _query_resources() -> QueryResources: + """Resolve the memory cap and thread count for one ``query`` process. + + :data:`QUERY_MEMORY_LIMIT_ENV` and :data:`QUERY_THREADS_ENV` override the + host-derived values; the thread default follows whichever cap is in + force, so a caller who lowers the cap gets fewer threads for free. + Raises ``ValueError`` on a malformed override; the command turns that + into exit 64. + """ + memory_env = os.environ.get(QUERY_MEMORY_LIMIT_ENV, "").strip() + threads_env = os.environ.get(QUERY_THREADS_ENV, "").strip() + if memory_env: + memory = _parse_size(memory_env) + if memory <= 0: + problem = f"{QUERY_MEMORY_LIMIT_ENV} must be a positive size" + raise ValueError(problem) + else: + memory = _query_memory_limit_bytes() + if threads_env: + threads = int(threads_env) + if threads < 1: + problem = f"{QUERY_THREADS_ENV} must be a positive integer" + raise ValueError(problem) + else: + threads = _query_threads(memory) + return QueryResources(memory_limit_bytes=memory, threads=threads) - physical = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") - return min(int(physical * 0.8), max(int(physical * 0.5), _QUERY_MEMORY_FLOOR_BYTES)) + +def _configure_query_resources(con: Any, resources: QueryResources, spill_dir: Path) -> None: + """Size the connection to the host BEFORE registration. + + Registration is the heaviest thing ``query`` does (the eager JSON readers + reserve memory per thread), so the cap and thread count must be in force + before it runs, not after, which is where they used to be applied. The + spill directory is set here for the same reason: a registration that + spills must spill into the private directory. Extension auto-install and + auto-load are switched off so nothing during registration can reach the + network; the lance extension is LOADed only when it is already installed + (:func:`atif_duck.infrastructure.registry.load_lance_extension`). + """ + con.execute(f"SET threads={int(resources.threads)}") + con.execute(f"SET memory_limit='{int(resources.memory_limit_bytes)}B'") + con.execute(f"SET temp_directory={_sql_str(str(spill_dir))}") + con.execute("SET autoinstall_known_extensions=false") + con.execute("SET autoload_known_extensions=false") + + +@contextmanager +def _private_spill_dir() -> Generator[Path]: + """A per-process spill directory outside the corpus, gone when the process is. + + ``mkdtemp`` creates it mode 0700 under the system temp dir. It is the + ONLY directory the sandbox grants, so it is where DuckDB spills a query + that exceeds ``memory_limit``; it used to be ``/.duckdb_tmp``, + inside the tree ``materialize`` scans, and files written there by caller + SQL persisted between runs. Removed on every exit path, error included, + after the connection is closed. + """ + path = Path(tempfile.mkdtemp(prefix="atif-sql-query-")) + try: + yield path + finally: + shutil.rmtree(path, ignore_errors=True) + + +def _refuse_root(command: str, fmt: OutputFormat) -> None: + """Exit 77 when running as uid 0, unless :data:`ALLOW_ROOT_ENV` is ``1``. + + The corpus artifacts are written mode 0444, which is what stops a + ``COPY ... TO`` at a granted parquet from succeeding at the filesystem. + Root ignores file modes, so as uid 0 that protection is gone and caller + SQL could overwrite the corpus (measured: probe d2 of the MicroVM review + rewrote ``steps.parquet`` as a one-row file). The override exists for + containers that have no other user; it logs a warning so the run is on + record. + """ + geteuid = getattr(os, "geteuid", None) + if geteuid is None or geteuid() != 0: + return + if os.environ.get(ALLOW_ROOT_ENV, "").strip() == "1": + from loguru import logger + + logger.warning( + "atif-sql {} is running as root because {}=1; the 0444 file modes no longer " + "protect the corpus from the SQL this process runs", + command, + ALLOW_ROOT_ENV, + ) + return + err = ClassifiedError( + kind="root_refused", + exit_code=EXIT_CODES["root_refused"], + message=f"atif-sql {command} refuses to run as root (uid 0)", + hint=f"run it as an unprivileged user, or set {ALLOW_ROOT_ENV}=1 to override " + "(a warning is logged; root can then overwrite corpus files from SQL)", + ) + emit_error(err, fmt) + raise SystemExit(err.exit_code) + + +#: Statement kinds the query sandbox executes. Everything else is refused +#: before execution, by name, using DuckDB's own parser on the hardened +#: connection. The refused kinds are the ones that name a file or defer a +#: statement past this check: COPY and COPY_DATABASE (DuckDB's ``allowed_paths`` +#: grants are read-write, so ``COPY ... TO (USE_TMP_FILE +#: false)`` would overwrite corpus data, and root ignores the 0444 mode that +#: used to stop it), EXPORT, ATTACH, DETACH, LOAD, EXTENSION (INSTALL), +#: PREPARE and EXECUTE (a prepared COPY parses as PREPARE). An allowlist, so a +#: statement kind a future DuckDB adds is refused until someone reads what it +#: does. +_QUERY_STATEMENT_KINDS: frozenset[str] = frozenset( + { + "SELECT", + "EXPLAIN", + "SET", + "VARIABLE_SET", + "CREATE", + "CREATE_FUNC", + "DROP", + "ALTER", + "INSERT", + "UPDATE", + "DELETE", + "MERGE_INTO", + "CALL", + "PRAGMA", + "TRANSACTION", + "VACUUM", + "ANALYZE", + } +) + + +def _refused_statement_kinds(con: Any, sql: str) -> list[str]: + """Statement kinds in ``sql`` outside :data:`_QUERY_STATEMENT_KINDS`, in order, once each. + + ``extract_statements`` is the parser ``execute`` uses, so there is no + second grammar to disagree with it, and it runs on the hardened + connection so a PRAGMA that reads a file while parsing meets the same + allowlist the query would. A parse error propagates as DuckDB's own + ``ParserException`` (exit 64), exactly as ``execute`` would have raised it. + """ + refused: list[str] = [] + for statement in con.extract_statements(sql): + kind = str(statement.type.name) + if kind not in _QUERY_STATEMENT_KINDS and kind not in refused: + refused.append(kind) + return refused #: Config options caller SQL may still change after the sandbox locks. Only @@ -169,7 +418,7 @@ def _harden_query_connection( con: Any, *, corpus_root: Path, - temp_dir: Path, + spill_dir: Path, columnar_paths: Sequence[Path] = (), ) -> None: """Sandbox a fully-registered connection before it runs caller SQL. @@ -180,25 +429,35 @@ def _harden_query_connection( path, ``ATTACH`` of unrelated databases, and ``INSTALL httpfs`` for network egress. + The memory cap, thread count and spill directory are NOT set here any + more: :func:`_configure_query_resources` applies them before + ``register`` runs, because registration is what needs them. This + function arms the allowlists and freezes the configuration. + Reach is granted at two granularities, and the split is what keeps injected SQL from overwriting the corpus it was read from: - * ``allowed_directories`` holds ONLY ``temp_dir`` — the spill area, which - must accept writes for a query that exceeds ``memory_limit`` to - complete at all. It sits inside the corpus root but holds no corpus - data. + * ``allowed_directories`` holds ONLY ``spill_dir``, the per-process + private directory from :func:`_private_spill_dir`, which must accept + writes for a query that exceeds ``memory_limit`` to complete at all. + It lives outside the corpus and is removed when the process exits, so + nothing written there persists and nothing under the corpus root is + writable. (DuckDB grants a directory read-write, so a ``COPY`` into it + would succeed at this layer; the statement gate in :func:`query` + refuses COPY before it gets here.) * ``allowed_paths`` holds the individual files the views read lazily: the analytics parquets from :func:`_lazy_read_paths` and the per-session columnar parquets in ``columnar_paths`` (what ``register`` bound this connection to). A file grant lets - ``read_parquet`` open that exact path; because ``COPY`` writes through - a sibling ``tmp_`` file, which is a different path, the grant - does not carry a plain write. Nothing else under the corpus is named, - so ``COPY`` over a ``trajectory.json`` or into a new ``pwned.csv`` is - refused. The columnar parquets are additionally written read-only - (``0o444``) by their producer, which closes the ``USE_TMP_FILE false`` - residual hole documented on :func:`query` for them: the open for - writing fails at the filesystem for any non-root user. + ``read_parquet`` open that exact path. Nothing else under the corpus + is named, so ``read_text`` of a ``trajectory.json`` is refused. DuckDB + 1.5.5 has no read-only grant (measured: a ``read_parquet`` relation + bound before ``enable_external_access=false`` is refused at query time + without a grant, and ``duckdb_settings()`` lists no write switch), so + each grant is read-write and ``COPY ... TO + (USE_TMP_FILE false)`` would overwrite it. Two things close that: the + statement gate refuses COPY for any uid, and :func:`_refuse_root` + keeps the producer's 0444 modes meaningful by refusing uid 0. Four ordering constraints, each one required: @@ -214,9 +473,7 @@ def _harden_query_connection( including itself). Without it the memory cap is decorative: caller SQL can just ``SET memory_limit`` back up. """ - con.execute(f"SET temp_directory={_sql_str(str(temp_dir))}") - con.execute(f"SET memory_limit='{_query_memory_limit_bytes()}B'") - con.execute(f"SET allowed_directories=[{_sql_str(str(temp_dir))}]") + con.execute(f"SET allowed_directories=[{_sql_str(str(spill_dir))}]") lazy_paths = ", ".join( _sql_str(str(path)) for path in (*columnar_paths, *_lazy_read_paths(corpus_root)) ) @@ -612,6 +869,44 @@ def materialize( # --------------------------------------------------------------------------- +def _vector_surface(corpus_root: Path) -> dict[str, Any]: + """How ``query`` and ``search`` will see the embeddings store, without installing anything. + + The extension check reads ``duckdb_extensions()`` on a throwaway + connection (the local extension directory, a few milliseconds, no + network); the store check is the directory ``embed`` writes, resolved the + way ``query`` and ``search`` resolve it (``ATIF_SQL_LANCE_URI`` wins). + """ + import duckdb + + from atif_duck.infrastructure.registry import lance_extension_installed + from atif_embed.infrastructure.settings import EmbedSettings + + store = EmbedSettings().resolve_lance_uri(corpus_root) + con = duckdb.connect() + try: + installed = lance_extension_installed(con) + finally: + con.close() + store_present = store.is_dir() + if store_present and installed: + state, note = "ready", f"store at {store}" + elif store_present: + state = "extension_missing" + note = ( + "store present but the lance DuckDB extension is not installed; " + "run `atif-sql embed --install-extension`" + ) + else: + state, note = "no_store", "no embeddings store; run `atif-sql embed --all --no-dry-run`" + return { + "lance_extension_installed": installed, + "embeddings_store_present": store_present, + "vector_search": state, + "note": note, + } + + def _dir_bytes(root: Path) -> int: """Total bytes of every file under ``root`` (0 if absent).""" if not root.is_dir(): @@ -647,6 +942,14 @@ def status( ``materialize --force`` brings them over). It applies the same per-session predicate the registry does, so it cannot say ``columnar`` while ``query`` silently parses JSON. + + ``vector search`` says whether ``semantic_search`` and ``search`` can + reach an embeddings store: ``ready`` (store present, lance extension + installed), ``no_store`` (run ``atif-sql embed``), or + ``extension_missing`` (a store exists but the DuckDB lance extension is + not installed; ``query`` binds ``message_embeddings`` empty rather than + downloading it, so run ``atif-sql embed --install-extension``). Read + from DuckDB's local extension listing; never installs anything. """ import time @@ -687,6 +990,7 @@ def status( "live": len(plan.skipped_live), } coverage = columnar_coverage(settings.corpus_root) + vector = _vector_surface(settings.corpus_root) if resolve_format(fmt) is OutputFormat.TABLE: print(f"agent: {settings.agent.value}") print(f"source root: {settings.source_root}") @@ -706,6 +1010,7 @@ def status( f"({coverage.columnar_sessions} of {coverage.total_sessions} complete sessions " "carry typed columnar artifacts)" ) + print(f"vector search: {vector['vector_search']} ({vector['note']})") else: emit_json( { @@ -721,6 +1026,9 @@ def status( "columnar_sessions": coverage.columnar_sessions, "json_sessions": coverage.json_sessions, "query_path": coverage.query_path, + "lance_extension_installed": vector["lance_extension_installed"], + "embeddings_store_present": vector["embeddings_store_present"], + "vector_search": vector["vector_search"], }, fmt, ) @@ -755,28 +1063,55 @@ def query( Sandbox ------- - The statement runs against a hardened connection (see + Before anything is registered the connection is sized to the host + (:func:`_query_resources`: a memory cap derived from available RAM and a + thread count derived from that cap, overridable with + ``ATIF_SQL_QUERY_MEMORY_LIMIT`` and ``ATIF_SQL_QUERY_THREADS``), its spill + directory is a private ``mkdtemp`` (mode 0700) outside the corpus that is + removed when the process exits, and extension auto-install and auto-load + are off. The lance extension is loaded only when it is already installed; + when a store exists and it is not, the vector surface binds empty with a + warning and ``atif-sql status`` says so. Registration therefore reaches + neither the network nor the corpus for writing, and it fits an 8 GiB + guest, where it used to exit 70 out of memory before the cap applied. + + The statement then runs against the hardened connection (see :func:`_harden_query_connection`). Reads reach the registered views and - nothing else; the only writable path is the query engine's own spill - directory ``/.duckdb_tmp``. ``read_text`` outside the - corpus, ``COPY`` anywhere in it (including over a ``trajectory.json``), - ``ATTACH`` of unrelated databases, and extension installs all fail with - exit 70 rather than reading credentials, corrupting the corpus, or - reaching the network. The embedding-store guard binds here exactly as it - does for ``search``, so a store written by a different provider refuses - to bind (exit 65) instead of scoring garbage. - - Two residual holes remain, both over recomputable derived data rather - than the transcript artifacts: - - * ``COPY ... TO '' (USE_TMP_FILE false)`` - overwrites it. A plain ``COPY`` to the same path is refused, because - DuckDB stages it through a sibling ``tmp_`` that carries no - grant; ``USE_TMP_FILE false`` writes the granted path directly. Re-run - ``atif-sql analyze`` to rebuild. - * ``DELETE``/``INSERT`` against ``lance_store.main.embeddings`` reach the - ATTACHed store, which the filesystem allowlist does not cover. Re-run - ``atif-sql embed --all --no-dry-run`` to rebuild. + nothing else. ``read_text`` outside the corpus, ``ATTACH`` of unrelated + databases and extension installs fail with exit 70 rather than reading + credentials or reaching the network, and nothing under the corpus root + is writable. Two layers keep caller SQL from writing the corpus: + + * DuckDB's file grants are read-write, and ``COPY ... TO (USE_TMP_FILE false)`` used to overwrite one, so every + statement kind that names a file is refused BEFORE execution, for any + uid, by DuckDB's own parser: COPY, EXPORT, ATTACH, DETACH, INSTALL, + LOAD, and PREPARE/EXECUTE (which could defer one). Exit 70, kind + ``sandbox_refused``. + * A 0444 file mode does not bind root, so ``query`` refuses to run as + uid 0 (exit 77) unless ``ATIF_SQL_ALLOW_ROOT=1`` is set, which logs a + warning. + + The embedding-store guard binds here exactly as it does for ``search``, + so a store written by a different provider refuses to bind (exit 65) + instead of scoring garbage. + + One residual hole remains, over recomputable derived data rather than + the transcript artifacts: ``DELETE``/``INSERT`` against + ``lance_store.main.embeddings`` reach the ATTACHed store, which the + filesystem allowlist does not cover. Re-run + ``atif-sql embed --all --no-dry-run`` to rebuild. + + What caller SQL can still see, and why that is accepted: + ``duckdb_settings()`` and ``current_setting(...)`` return the sandbox's + own configuration, including the corpus root, the spill directory, the + memory cap and every granted parquet path, which names every session id. + DuckDB cannot hide a setting from SQL, and the grants have to be per-file + for the lazy reads to work at all, so the inventory is a property of the + design. The caller is the local user, who can list the corpus directory + and ``SELECT session_id FROM sessions`` anyway; ``query`` is not a + privilege boundary (SECURITY.md) and must not be exposed to a caller who + may not know the session ids. ``lock_configuration`` freezes the connection's settings, so ``SET`` is refused for everything except ``TimeZone``, which stays open for @@ -787,10 +1122,15 @@ def query( Exit codes ---------- * 64 parse_error malformed SQL (or no SQL and no --examples) + * 64 invalid_input a malformed ``ATIF_SQL_QUERY_MEMORY_LIMIT`` or + ``ATIF_SQL_QUERY_THREADS`` * 65 catalog_error unknown view/macro/column (try ``atif-sql schema``) * 65 embedding_mismatch the Lance store was written by another provider + * 70 sandbox_refused a statement kind the sandbox never runs (COPY, + EXPORT, ATTACH, INSTALL, LOAD, PREPARE, ...) * 70 runtime_error everything else (an unmaterialized corpus, or SQL the sandbox refused) + * 77 root_refused running as uid 0 without ``ATIF_SQL_ALLOW_ROOT=1`` Examples -------- @@ -821,6 +1161,20 @@ def query( ) raise SystemExit(EXIT_CODES["parse_error"]) + _refuse_root("query", fmt) + try: + resources = _query_resources() + except ValueError as exc: + err = ClassifiedError( + kind="invalid_input", + exit_code=EXIT_CODES["invalid_input"], + message=str(exc), + hint=f"{QUERY_MEMORY_LIMIT_ENV} takes a size such as 6GB or 512MiB; " + f"{QUERY_THREADS_ENV} takes a positive integer", + ) + emit_error(err, fmt) + raise SystemExit(err.exit_code) from exc + import duckdb from atif_cli.duck_errors import REGISTRATION_ERRORS, classify_registration_error @@ -832,35 +1186,49 @@ def query( expected_model, expected_dim = embed_settings.expected_embedding_identity() lance_uri = embed_settings.resolve_lance_uri(settings.corpus_root) - con = duckdb.connect() - try: - try: - sources = register( - con, - settings.corpus_root, - lance_uri=lance_uri, - expected_model=expected_model, - expected_dim=expected_dim, - ) - _harden_query_connection( - con, - corpus_root=settings.corpus_root, - temp_dir=settings.corpus_root / ".duckdb_tmp", - columnar_paths=sources.lazy_read_paths, - ) - cursor = con.execute(sql) - except REGISTRATION_ERRORS as exc: - err = classify_registration_error(exc) - emit_error(err, fmt) - raise SystemExit(err.exit_code) from exc + with _private_spill_dir() as spill_dir: + con = duckdb.connect() try: - emit_cursor(cursor, fmt) - except REGISTRATION_ERRORS as exc: - err = classify_registration_error(exc) - emit_error(err, fmt) - raise SystemExit(err.exit_code) from exc - finally: - con.close() + try: + _configure_query_resources(con, resources, spill_dir) + sources = register( + con, + settings.corpus_root, + lance_uri=lance_uri, + expected_model=expected_model, + expected_dim=expected_dim, + ) + _harden_query_connection( + con, + corpus_root=settings.corpus_root, + spill_dir=spill_dir, + columnar_paths=sources.lazy_read_paths, + ) + refused = _refused_statement_kinds(con, sql) + if refused: + err = ClassifiedError( + kind="sandbox_refused", + exit_code=EXIT_CODES["sandbox_refused"], + message=f"the query sandbox does not run {', '.join(refused)} statements", + hint="query runs SELECT and in-memory DDL/DML only; COPY, EXPORT, " + "ATTACH, DETACH, INSTALL, LOAD, PREPARE and EXECUTE are refused " + "because they reach files or defer a statement past this check", + ) + emit_error(err, fmt) + raise SystemExit(err.exit_code) + cursor = con.execute(sql) + except REGISTRATION_ERRORS as exc: + err = classify_registration_error(exc) + emit_error(err, fmt) + raise SystemExit(err.exit_code) from exc + try: + emit_cursor(cursor, fmt) + except REGISTRATION_ERRORS as exc: + err = classify_registration_error(exc) + emit_error(err, fmt) + raise SystemExit(err.exit_code) from exc + finally: + con.close() # --------------------------------------------------------------------------- @@ -930,6 +1298,8 @@ def analyze( fmt Summary format; ``auto`` = JSON on a pipe. """ + _refuse_root("analyze", fmt) + from atif_analytics.application.analyze import run_analyze from atif_analytics.infrastructure.settings import AnalyticsSettings @@ -975,17 +1345,67 @@ def analyze( # --------------------------------------------------------------------------- +def _install_lance_extension(fmt: OutputFormat, *, quiet: bool = False) -> None: + """``INSTALL lance; LOAD lance`` on a throwaway connection. + + The one network reach outside Bedrock, so it lives behind ``embed``. + With ``quiet`` (a real embed run) a failure is a warning rather than an + exit: the backfill itself still runs, and ``status`` will keep saying + the extension is missing until a later attempt succeeds. Without + ``quiet`` (``--install-extension``) the result is emitted as JSON and a + failure exits 70. + """ + import duckdb + + from atif_cli.duck_errors import classify_duckdb_error + from atif_duck.infrastructure.registry import install_lance_extension + + con = duckdb.connect() + try: + try: + install_path = install_lance_extension(con) + except duckdb.Error as exc: + if quiet: + from loguru import logger + + logger.warning( + "Could not install the lance DuckDB extension ({}); vector search stays " + "unavailable until `atif-sql embed --install-extension` succeeds", + str(exc).splitlines()[0], + ) + return + err = classify_duckdb_error(exc) + emit_error(err, fmt) + raise SystemExit(err.exit_code) from exc + finally: + con.close() + if not quiet: + emit_json({"extension": "lance", "installed": True, "install_path": install_path}, fmt) + + @app.command def embed( *, limit: int | None = None, all_steps: Annotated[bool, cyclopts.Parameter(name="--all")] = False, dry_run: bool = False, + install_extension: Annotated[bool, cyclopts.Parameter(name="--install-extension")] = False, corpus_root: Path | None = None, fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO, ) -> None: """Embed unembedded corpus steps with Cohere Embed v4 and append to LanceDB. + Extension + --------- + ``query`` and ``search`` read the store through DuckDB's lance extension + and never install it themselves (installing is a 242 MB download from + the extension repository, and ``query`` runs unattended). It is + installed here, where the network is already a deliberate act: every + REAL run installs it first, and ``--install-extension`` installs it and + exits without touching Bedrock or the store (``{"extension": "lance", + "installed": true, "install_path": ...}``). ``atif-sql status`` reports + whether it is present. + Cost ---- Calls Bedrock (``global.cohere.embed-v4:0``) on every unembedded step @@ -1001,6 +1421,7 @@ def embed( --limit N Cap the number of steps embedded this run. --all Explicitly embed EVERY unembedded step (full backfill). --dry-run Preview only; emit plan JSON, no embedding calls. + --install-extension Install the lance DuckDB extension and exit (no Bedrock). --corpus-root Override the materialized corpus root. Output @@ -1014,6 +1435,10 @@ def embed( state (the store or its config requires operator action; retrying without intervention cannot succeed, so unattended lanes suppress retries on 78). """ + if install_extension: + _install_lance_extension(fmt) + return + import asyncio from atif_embed.application.embed import run_backfill @@ -1034,6 +1459,8 @@ def embed( settings = _corpus_settings(None, corpus_root) embed_settings = EmbedSettings() + if not dry_run: + _install_lance_extension(fmt, quiet=True) try: result = asyncio.run( run_backfill( @@ -1107,8 +1534,13 @@ def search( Sorted by cosine distance ascending — highest sim first. Exit codes: 0 success, 2 no_embeddings, 65 embedding_mismatch (the store - was written by another provider), 70 runtime. + was written by another provider), 70 runtime, 77 root_refused (uid 0 + without ``ATIF_SQL_ALLOW_ROOT=1``), 78 extension_missing (a store exists + but the lance DuckDB extension is not installed; run + ``atif-sql embed --install-extension``). """ + _refuse_root("search", fmt) + import duckdb from atif_cli.duck_errors import ( @@ -1116,7 +1548,7 @@ def search( classify_duckdb_error, classify_registration_error, ) - from atif_duck.infrastructure.registry import register + from atif_duck.infrastructure.registry import lance_extension_installed, register from atif_embed.application.embed import embed_query from atif_embed.infrastructure.settings import EmbedSettings @@ -1140,6 +1572,19 @@ def search( emit_error(err, fmt) raise SystemExit(err.exit_code) from exc + if lance_uri.is_dir() and not lance_extension_installed(con): + emit_error( + ClassifiedError( + kind="extension_missing", + exit_code=EXIT_CODES["extension_missing"], + message="the lance DuckDB extension is not installed, so the embeddings " + f"store at {lance_uri} cannot be read", + hint="run: atif-sql embed --install-extension (a one-time download)", + ), + fmt, + ) + raise SystemExit(EXIT_CODES["extension_missing"]) + row = con.execute("SELECT count(*) FROM message_embeddings").fetchone() if not row or int(row[0]) == 0: emit_error( diff --git a/packages/atif-cli/src/atif_cli/errors.py b/packages/atif-cli/src/atif_cli/errors.py index ca9cac5..6686960 100644 --- a/packages/atif-cli/src/atif_cli/errors.py +++ b/packages/atif-cli/src/atif_cli/errors.py @@ -35,6 +35,15 @@ "validation_error": 65, # convert: trajectory failed TrajectoryValidator "embedding_mismatch": 65, # query/search: Lance store written by another provider "runtime_error": 70, # everything else duckdb (or the adapter) raises + "sandbox_refused": 70, # query: a statement kind the sandbox never runs (COPY, EXPORT, + # ATTACH, INSTALL, LOAD, PREPARE, ...). Same number as runtime_error, which is what + # every other sandbox refusal exits with; the distinct kind names the reason. + "root_refused": 77, # query/search/analyze: running as uid 0 without ATIF_SQL_ALLOW_ROOT=1. + # EX_NOPERM by convention: the process has too MUCH permission for the file modes + # that protect the corpus to mean anything. + "extension_missing": 78, # search: the Lance store exists but the lance DuckDB extension + # is not installed; `atif-sql embed --install-extension` is the operator action, so the + # same EX_CONFIG contract as terminal_state. "terminal_state": 78, # embed: store/config state requires OPERATOR action (EX_CONFIG); # retrying without intervention cannot succeed — unattended lanes suppress # retries on this code instead of burning identical ticks diff --git a/packages/atif-cli/tests/cli_fixtures.py b/packages/atif-cli/tests/cli_fixtures.py index f1bd4b6..75818cf 100644 --- a/packages/atif-cli/tests/cli_fixtures.py +++ b/packages/atif-cli/tests/cli_fixtures.py @@ -201,3 +201,18 @@ def write_analytics_parquets(corpus_root: Path, session_id: str) -> list[Path]: finally: scratch.close() return sorted(written) + + +# ``register_vss`` no longer installs the lance extension (that download at +# query time was finding 4 of the MicroVM review), so the suite installs it +# once up front. A no-op where it is already present; a one-time download on a +# fresh runner, exactly what every register() call used to do implicitly. +@pytest.fixture(scope="session", autouse=True) +def _lance_extension_present() -> None: # pyright: ignore[reportUnusedFunction] + import duckdb + + con = duckdb.connect() + try: + con.execute("INSTALL lance") + finally: + con.close() diff --git a/packages/atif-cli/tests/test_app.py b/packages/atif-cli/tests/test_app.py index 1de40a0..611108c 100644 --- a/packages/atif-cli/tests/test_app.py +++ b/packages/atif-cli/tests/test_app.py @@ -4,8 +4,12 @@ from __future__ import annotations +import hashlib import json +import os import re +import shutil +import subprocess import sys from datetime import UTC, datetime from pathlib import Path @@ -13,19 +17,26 @@ import pytest from cli_fixtures import write_analytics_parquets +from loguru import logger from atif_cli import app as app_module from atif_cli.app import ( + ALLOW_ROOT_ENV, DEFAULT_LOG_LEVEL, LOG_LEVEL_ENV, + QUERY_MEMORY_LIMIT_ENV, + QUERY_THREADS_ENV, _stderr_log_level, + analyze, app, convert, + embed, examples, materialize as materialize_cmd, query, schema, search, + status, ) from atif_cli.errors import EXIT_CODES from atif_cli.output import ( @@ -326,16 +337,107 @@ def duckdb_default_memory_limit() -> str: con.close() -def _expect_refusal(sql: str, corpus_root: Path, capsys: pytest.CaptureFixture[str]) -> str: - """Run ``sql`` expecting the sandbox to refuse it; return the message.""" +def _expect_refusal( + sql: str, + corpus_root: Path, + capsys: pytest.CaptureFixture[str], + *, + kind: str = "runtime_error", +) -> str: + """Run ``sql`` expecting the sandbox to refuse it; return the message. + + ``kind`` is ``runtime_error`` for a refusal DuckDB itself raises + (allowlist, locked configuration) and ``sandbox_refused`` for a + statement kind the CLI refuses by name before executing anything. + """ with pytest.raises(SystemExit) as excinfo: _run_query(sql, corpus_root) - assert excinfo.value.code == EXIT_CODES["runtime_error"] + assert excinfo.value.code == EXIT_CODES[kind] payload = json.loads(capsys.readouterr().err) - assert payload["error"]["kind"] == "runtime_error" + assert payload["error"]["kind"] == kind return str(payload["error"]["message"]) +def _tree_digest(root: Path) -> str: + """One hash over every path, mode and byte under ``root``.""" + digest = hashlib.sha256() + for path in sorted(p for p in root.rglob("*")): + rel = path.relative_to(root).as_posix() + digest.update(f"{rel}\0{oct(path.stat().st_mode)}\0".encode()) + if path.is_file(): + digest.update(path.read_bytes()) + digest.update(b"\n") + return digest.hexdigest() + + +def _capture_warnings() -> tuple[list[str], int]: + warnings: list[str] = [] + sink_id = logger.add(lambda message: warnings.append(str(message)), level="WARNING") + return warnings, sink_id + + +def _write_two_session_corpus(root: Path) -> Path: + """The contract corpus plus a second session: two ``edges.jsonl`` for the glob reader. + + One file registers under any limit because the reader uses one thread + per file; the per-thread reservation that finding 1 is about needs at + least two. + """ + _write_contract_corpus(root) + second = "22222222-2222-2222-2222-222222222222" + src = root / "sessions" / SESSION_ID + dst = root / "sessions" / second + shutil.copytree(src, dst) + for name in ("trajectory.json", "meta.json"): + path = dst / name + path.write_text(path.read_text().replace(SESSION_ID, second)) + return root + + +class _RecordingConnection: + """Wrap a DuckDB connection so every statement text is observable. + + ``execute`` is recorded; everything else (``read_parquet``, + ``extract_statements``, ``close``) delegates. Registration only ever + touches those, so the wrapper is transparent to the code under test. + """ + + def __init__(self, con: Any, statements: list[str]) -> None: + self._con = con + self._statements = statements + + def execute(self, sql: str, *args: Any, **kwargs: Any) -> Any: + self._statements.append(sql) + return self._con.execute(sql, *args, **kwargs) + + def __getattr__(self, name: str) -> Any: + return getattr(self._con, name) + + +@pytest.fixture +def empty_extension_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, list[str]]: + """Every ``duckdb.connect()`` in the command bodies sees NO installed extensions. + + The connection is pointed at an empty ``extension_directory`` and + wrapped so the statements it runs are recorded. Returns the directory + (still empty afterwards proves nothing was installed) and the recording. + """ + import duckdb + + ext_dir = tmp_path / "no-extensions" + ext_dir.mkdir() + statements: list[str] = [] + real_connect = duckdb.connect + + def connect(*args: Any, **kwargs: Any) -> Any: + con = real_connect(*args, **kwargs) + con.execute(f"SET extension_directory='{ext_dir}'") + return _RecordingConnection(con, statements) + + monkeypatch.setattr(duckdb, "connect", connect) + return ext_dir, statements + + class TestMaterializeReportOutput: """An unreadable session is counted nowhere else, so hiding it reads as success.""" @@ -403,19 +505,24 @@ def test_copy_out_of_tree_is_refused( ) -> None: target = tmp_path / "exfil.csv" - _expect_refusal(f"COPY (SELECT 42 AS x) TO '{target}'", query_corpus, capsys) + _expect_refusal( + f"COPY (SELECT 42 AS x) TO '{target}'", query_corpus, capsys, kind="sandbox_refused" + ) assert not target.exists() def test_attach_unrelated_database_is_refused( self, query_corpus: Path, tmp_path: Path, capsys: pytest.CaptureFixture[str] ) -> None: - _expect_refusal(f"ATTACH '{tmp_path / 'other.db'}' AS o", query_corpus, capsys) + _expect_refusal( + f"ATTACH '{tmp_path / 'other.db'}' AS o", query_corpus, capsys, kind="sandbox_refused" + ) - def test_extension_install_is_refused( + def test_extension_install_and_load_are_refused( self, query_corpus: Path, capsys: pytest.CaptureFixture[str] ) -> None: """No INSTALL/LOAD means no httpfs, which means no network egress.""" - _expect_refusal("INSTALL httpfs", query_corpus, capsys) + _expect_refusal("INSTALL httpfs", query_corpus, capsys, kind="sandbox_refused") + _expect_refusal("LOAD httpfs", query_corpus, capsys, kind="sandbox_refused") def test_injected_sql_cannot_reopen_the_sandbox( self, query_corpus: Path, capsys: pytest.CaptureFixture[str] @@ -426,6 +533,8 @@ def test_injected_sql_cannot_reopen_the_sandbox( "SET allowed_directories=['/']", "SET memory_limit='512GB'", "SET temp_directory='/tmp'", + "SET autoinstall_known_extensions=true", + "SET threads=64", ): _expect_refusal(statement, query_corpus, capsys) @@ -465,6 +574,180 @@ def test_corpus_root_containing_a_quote_still_binds( assert json.loads(capsys.readouterr().out) == [{"n": 1}] +class TestQueryStatementGate: + """Statement kinds that name a file are refused by name before anything executes. + + DuckDB's ``allowed_paths`` grants are read-write and it has no read-only + grant (pinned in :class:`TestHardenedConnectionLayer`), so this gate is + what stops ``COPY ... TO (USE_TMP_FILE false)`` for + any uid. + """ + + @pytest.mark.parametrize( + ("statement", "kind"), + [ + ("COPY (SELECT 1) TO '{tmp}/a.csv'", "COPY"), + ("COPY (SELECT 1) TO '{tmp}/a.csv' (USE_TMP_FILE false)", "COPY"), + ("EXPORT DATABASE '{tmp}/exported'", "EXPORT"), + ("PREPARE p AS COPY (SELECT 1) TO '{tmp}/a.csv'", "PREPARE"), + ("EXECUTE p", "EXECUTE"), + ("ATTACH '{tmp}/other.db' AS o", "ATTACH"), + ("DETACH o", "DETACH"), + # DuckDB parses INSTALL and LOAD to the same statement kind. + ("INSTALL httpfs", "LOAD"), + ("LOAD httpfs", "LOAD"), + ], + ) + def test_file_statement_kinds_are_refused_by_name( + self, + query_corpus: Path, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + statement: str, + kind: str, + ) -> None: + message = _expect_refusal( + statement.format(tmp=tmp_path), query_corpus, capsys, kind="sandbox_refused" + ) + assert kind in message + assert not (tmp_path / "a.csv").exists() + + def test_a_refused_statement_stops_the_whole_batch( + self, query_corpus: Path, tmp_path: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + """Nothing in a batch runs when any statement in it is refused.""" + target = tmp_path / "late.csv" + message = _expect_refusal( + f"SELECT 1; COPY (SELECT 1) TO '{target}'", query_corpus, capsys, kind="sandbox_refused" + ) + assert "COPY" in message + assert not target.exists() + + def test_select_batches_and_in_memory_ddl_still_run( + self, query_corpus: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + _run_query("SELECT 1 AS a; SELECT 2 AS b", query_corpus) + assert json.loads(capsys.readouterr().out) == [{"b": 2}] + _run_query("CREATE TEMP TABLE t AS SELECT 1 AS x; SELECT * FROM t", query_corpus) + assert json.loads(capsys.readouterr().out) == [{"x": 1}] + _run_query("EXPLAIN SELECT count(*) FROM sessions", query_corpus) + assert json.loads(capsys.readouterr().out) + + def test_parse_errors_still_exit_64( + self, query_corpus: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + with pytest.raises(SystemExit) as excinfo: + _run_query("SELEC 1", query_corpus) + assert excinfo.value.code == EXIT_CODES["parse_error"] + assert json.loads(capsys.readouterr().err)["error"]["kind"] == "parse_error" + + def test_allowlist_names_real_statement_kinds_and_omits_the_file_facing_ones(self) -> None: + import duckdb + + known = {member for member in dir(duckdb.StatementType) if member.isupper()} + assert known >= app_module._QUERY_STATEMENT_KINDS + for refused in ( + "COPY", + "COPY_DATABASE", + "EXPORT", + "ATTACH", + "DETACH", + "LOAD", + "EXTENSION", + "PREPARE", + "EXECUTE", + ): + assert refused in known + assert refused not in app_module._QUERY_STATEMENT_KINDS + + +class TestQueryRefusesRoot: + """A 0444 file mode does not bind uid 0, so the commands that run SQL refuse it.""" + + @staticmethod + def _as_root(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(os, "geteuid", lambda: 0) + monkeypatch.delenv(ALLOW_ROOT_ENV, raising=False) + + def test_query_exits_77_before_opening_duckdb( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + import duckdb + + self._as_root(monkeypatch) + + def refuse_connect(*args: object, **kwargs: object) -> Any: + del args, kwargs + pytest.fail("DuckDB was opened as root") + + monkeypatch.setattr(duckdb, "connect", refuse_connect) + with pytest.raises(SystemExit) as excinfo: + _run_query("SELECT 1", query_corpus) + assert excinfo.value.code == EXIT_CODES["root_refused"] == 77 + err = json.loads(capsys.readouterr().err)["error"] + assert err["kind"] == "root_refused" + assert "root" in err["message"] + assert ALLOW_ROOT_ENV in err["hint"] + + def test_search_and_analyze_refuse_root_too( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + self._as_root(monkeypatch) + with pytest.raises(SystemExit) as excinfo: + search("anything", corpus_root=query_corpus, fmt=OutputFormat.JSON) + assert excinfo.value.code == EXIT_CODES["root_refused"] + assert json.loads(capsys.readouterr().err)["error"]["kind"] == "root_refused" + with pytest.raises(SystemExit) as excinfo: + analyze(corpus_root=query_corpus, fmt=OutputFormat.JSON) + assert excinfo.value.code == EXIT_CODES["root_refused"] + assert json.loads(capsys.readouterr().err)["error"]["kind"] == "root_refused" + + def test_the_escape_hatch_runs_with_a_warning( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + self._as_root(monkeypatch) + monkeypatch.setenv(ALLOW_ROOT_ENV, "1") + warnings, sink_id = _capture_warnings() + try: + _run_query("SELECT count(*) AS n FROM sessions", query_corpus) + finally: + logger.remove(sink_id) + assert json.loads(capsys.readouterr().out) == [{"n": 1}] + assert any("root" in w and ALLOW_ROOT_ENV in w for w in warnings) + + def test_only_the_value_one_opens_the_hatch( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + self._as_root(monkeypatch) + monkeypatch.setenv(ALLOW_ROOT_ENV, "yes") + with pytest.raises(SystemExit) as excinfo: + _run_query("SELECT 1", query_corpus) + assert excinfo.value.code == EXIT_CODES["root_refused"] + capsys.readouterr() + + def test_an_unprivileged_uid_is_untouched( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + monkeypatch.setattr(os, "geteuid", lambda: 1000) + _run_query("SELECT count(*) AS n FROM sessions", query_corpus) + assert json.loads(capsys.readouterr().out) == [{"n": 1}] + + class TestQuerySandboxFileAllowlist: """The allowlist gates LAZY reads only — the surface a corpus-only test misses. @@ -521,10 +804,12 @@ def test_copy_over_a_corpus_artifact_is_refused( trajectory = query_corpus / "sessions" / SESSION_ID / "trajectory.json" before = trajectory.read_text() - message = _expect_refusal( - f"COPY (SELECT 'destroyed' AS x) TO '{trajectory}'", query_corpus, capsys + _expect_refusal( + f"COPY (SELECT 'destroyed' AS x) TO '{trajectory}'", + query_corpus, + capsys, + kind="sandbox_refused", ) - assert "Permission" in message assert trajectory.read_text() == before def test_copy_new_file_into_the_corpus_is_refused( @@ -532,7 +817,9 @@ def test_copy_new_file_into_the_corpus_is_refused( ) -> None: target = query_corpus / "pwned.csv" - _expect_refusal(f"COPY (SELECT 42) TO '{target}'", query_corpus, capsys) + _expect_refusal( + f"COPY (SELECT 42) TO '{target}'", query_corpus, capsys, kind="sandbox_refused" + ) assert not target.exists() def test_copy_into_the_lance_store_is_refused( @@ -541,62 +828,63 @@ def test_copy_into_the_lance_store_is_refused( _write_lance_store(query_corpus / "embeddings_lance") target = query_corpus / "embeddings_lance" / "exfil.csv" - _expect_refusal(f"COPY (SELECT 42) TO '{target}'", query_corpus, capsys) + _expect_refusal( + f"COPY (SELECT 42) TO '{target}'", query_corpus, capsys, kind="sandbox_refused" + ) assert not target.exists() - def test_traversal_out_of_the_spill_directory_is_refused( + def test_copy_into_the_legacy_spill_directory_is_refused_and_nothing_persists( self, query_corpus: Path, capsys: pytest.CaptureFixture[str] ) -> None: - """``.duckdb_tmp`` is writable and sits inside the corpus — ``..`` must not escape. + """Probe d3 of the review: ``/.duckdb_tmp`` used to be the one writable dir. - The spill dir is created up front on purpose: DuckDB reports a - missing parent directory as an IO error, which would make this pass - with no sandbox at all. + It is pre-created here, as the review did, so a success would be the + old grant and not a missing parent; it is no longer granted at all. """ - (query_corpus / ".duckdb_tmp").mkdir(exist_ok=True) - trajectory = query_corpus / "sessions" / SESSION_ID / "trajectory.json" - before = trajectory.read_text() - escape = query_corpus / ".duckdb_tmp" / ".." / "sessions" / SESSION_ID / "trajectory.json" + legacy = query_corpus / ".duckdb_tmp" + legacy.mkdir() + before = _tree_digest(query_corpus) - message = _expect_refusal( - f"COPY (SELECT 'destroyed' AS x) TO '{escape}'", query_corpus, capsys + _expect_refusal( + f"COPY (SELECT 1) TO '{legacy}/probe.csv'", query_corpus, capsys, kind="sandbox_refused" ) - assert "Permission" in message - assert trajectory.read_text() == before + _expect_refusal(f"SELECT * FROM read_text('{legacy}/probe.csv')", query_corpus, capsys) + assert _tree_digest(query_corpus) == before + assert list(legacy.iterdir()) == [] - def test_plain_copy_over_an_analytics_parquet_is_refused( - self, query_corpus: Path, capsys: pytest.CaptureFixture[str] + def test_use_tmp_file_false_over_a_granted_parquet_is_refused_for_any_uid( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], ) -> None: - """The read grant is per-file; DuckDB stages ``COPY`` through an ungranted sibling.""" + """Probe d2 of the review, the one that overwrote ``steps.parquet`` as root. + + The analytics parquet is a granted path and writable by this uid at + the filesystem, so a refusal here is the statement gate and not a + file mode. Repeated as uid 0 with the escape hatch open: still + refused, because the gate does not look at the uid. + """ parquet = query_corpus / "analytics" / "clusters.parquet" before = parquet.read_bytes() + statement = f"COPY (SELECT 'x' AS uuid) TO '{parquet}' (FORMAT PARQUET, USE_TMP_FILE false)" - _expect_refusal( - f"COPY (SELECT 'x' AS uuid) TO '{parquet}' (FORMAT PARQUET)", query_corpus, capsys - ) + _expect_refusal(statement, query_corpus, capsys, kind="sandbox_refused") + assert parquet.read_bytes() == before + + monkeypatch.setattr(os, "geteuid", lambda: 0) + monkeypatch.setenv(ALLOW_ROOT_ENV, "1") + _expect_refusal(statement, query_corpus, capsys, kind="sandbox_refused") assert parquet.read_bytes() == before - def test_documented_residual_holes_still_behave_as_documented( + def test_documented_residual_hole_still_behaves_as_documented( self, query_corpus: Path, capsys: pytest.CaptureFixture[str] ) -> None: - """Pins the two holes ``query``'s docstring admits, so the promise cannot rot. + """Pins the one hole ``query``'s docstring admits, so the promise cannot rot. - Both are recomputable derived data. If a future DuckDB closes either, - this test fails and the docstring's Sandbox section must lose the - corresponding paragraph. + Recomputable derived data. If a future DuckDB closes it, this test + fails and the docstring's Sandbox section must lose the paragraph. """ - parquet = query_corpus / "analytics" / "clusters.parquet" - before = parquet.read_bytes() - _run_query( - f"COPY (SELECT 'x' AS uuid) TO '{parquet}' (FORMAT PARQUET, USE_TMP_FILE false)", - query_corpus, - ) - capsys.readouterr() - assert parquet.read_bytes() != before, ( - "USE_TMP_FILE false no longer overwrites a granted parquet — " - "drop that paragraph from query's Sandbox docstring" - ) - _write_lance_store(query_corpus / "embeddings_lance") _run_query("DELETE FROM lance_store.main.embeddings", query_corpus) capsys.readouterr() @@ -607,67 +895,470 @@ def test_documented_residual_holes_still_behave_as_documented( ) -class TestQueryMemoryCap: - """``_query_memory_limit_bytes`` is pure arithmetic over ``os.sysconf``.""" +def _hardened_connection(corpus_root: Path, spill_dir: Path) -> Any: + """What ``query`` builds, minus the statement gate: DuckDB's own layer alone.""" + import duckdb - @staticmethod - def _limit_for(physical_bytes: int, monkeypatch: pytest.MonkeyPatch) -> int: - import os + from atif_duck.infrastructure.registry import register - from atif_cli import app as app_mod + con = duckdb.connect() + app_module._configure_query_resources(con, app_module._query_resources(), spill_dir) + sources = register(con, corpus_root) + app_module._harden_query_connection( + con, + corpus_root=corpus_root, + spill_dir=spill_dir, + columnar_paths=sources.lazy_read_paths, + ) + return con - page = 4096 - def _sysconf(name: str | int) -> int: - return page if name == "SC_PAGE_SIZE" else physical_bytes // page +class TestHardenedConnectionLayer: + """DuckDB's allowlist, exercised WITHOUT the statement gate. - monkeypatch.setattr(os, "sysconf", _sysconf) - return app_mod._query_memory_limit_bytes() + The CLI tests above hit the gate first, so this class is what proves + the second layer still holds on its own, and pins the two facts the + gate exists for: a directory grant and a file grant are both read-write. + """ - def test_ceiling_wins_on_a_small_host(self, monkeypatch: pytest.MonkeyPatch) -> None: - """4 GiB of RAM: 50% is 2 GiB, but the 80% ceiling (3.2 GiB) is tighter than the floor.""" - from atif_cli.app import _QUERY_MEMORY_FLOOR_BYTES + @pytest.fixture + def layer(self, query_corpus: Path, tmp_path: Path) -> Any: + spill = tmp_path / "spill" + spill.mkdir(mode=0o700) + con = _hardened_connection(query_corpus, spill) + yield con, spill + con.close() - physical = 4 * 1024**3 - assert self._limit_for(physical, monkeypatch) == int(physical * 0.8) - assert int(physical * 0.8) < _QUERY_MEMORY_FLOOR_BYTES + def test_duckdb_refuses_every_ungranted_path( + self, layer: Any, query_corpus: Path, tmp_path: Path + ) -> None: + import duckdb - def test_half_of_ram_wins_on_a_large_host(self, monkeypatch: pytest.MonkeyPatch) -> None: - physical = 128 * 1024**3 - assert self._limit_for(physical, monkeypatch) == int(physical * 0.5) + con, spill = layer + trajectory = query_corpus / "sessions" / SESSION_ID / "trajectory.json" + parquet = query_corpus / "analytics" / "clusters.parquet" + for statement in ( + f"COPY (SELECT 1) TO '{tmp_path / 'out.csv'}'", + f"COPY (SELECT 1) TO '{trajectory}'", + f"COPY (SELECT 1) TO '{trajectory}' (USE_TMP_FILE false)", + f"COPY (SELECT 1) TO '{query_corpus / 'pwned.csv'}'", + f"COPY (SELECT 'x' AS uuid) TO '{parquet}' (FORMAT PARQUET)", + f"COPY (SELECT 1) TO '{spill}/../escape.csv'", + f"SELECT * FROM read_text('{spill}/../../etc/passwd')", + "SELECT * FROM read_text('/etc/passwd')", + f"ATTACH '{tmp_path / 'other.db'}' AS o", + "INSTALL httpfs", + ): + with pytest.raises(duckdb.Error) as excinfo: + con.execute(statement) + assert "Permission" in str(excinfo.value) or "disabled" in str(excinfo.value), statement + assert not (tmp_path / "out.csv").exists() + assert not (tmp_path / "escape.csv").exists() + + def test_grants_are_read_write_which_is_why_the_gate_exists( + self, layer: Any, query_corpus: Path + ) -> None: + """DuckDB 1.5.5 has no read-only grant; if one appears, this fails and the design can simplify.""" + con, spill = layer + con.execute(f"COPY (SELECT 1 AS x) TO '{spill}/probe.csv'") + assert (spill / "probe.csv").is_file() - def test_floor_beats_half_between_the_two(self, monkeypatch: pytest.MonkeyPatch) -> None: - """12 GiB: half (6 GiB) is under the floor, and the floor is under the 80% ceiling.""" - from atif_cli.app import _QUERY_MEMORY_FLOOR_BYTES + parquet = query_corpus / "analytics" / "clusters.parquet" + before = parquet.read_bytes() + con.execute( + f"COPY (SELECT 'x' AS uuid) TO '{parquet}' (FORMAT PARQUET, USE_TMP_FILE false)" + ) + assert parquet.read_bytes() != before, ( + "DuckDB now refuses a write to a granted path: the statement gate is " + "belt-and-braces and query's Sandbox docstring can say so" + ) - physical = 12 * 1024**3 - assert self._limit_for(physical, monkeypatch) == _QUERY_MEMORY_FLOOR_BYTES - assert int(physical * 0.5) < _QUERY_MEMORY_FLOOR_BYTES < int(physical * 0.8) + def test_resources_and_extension_settings_are_in_force(self, layer: Any) -> None: + con, spill = layer + resources = app_module._query_resources() + settings = dict( + con.execute( + "SELECT name, value FROM duckdb_settings() WHERE name IN " + "('threads', 'temp_directory', 'autoinstall_known_extensions', " + "'autoload_known_extensions', 'enable_external_access', 'lock_configuration')" + ).fetchall() + ) + assert int(settings["threads"]) == resources.threads + assert settings["temp_directory"] == str(spill) + assert settings["autoinstall_known_extensions"] == "false" + assert settings["autoload_known_extensions"] == "false" + assert settings["enable_external_access"] == "false" + assert settings["lock_configuration"] == "true" - def test_never_exceeds_duckdbs_own_default(self, monkeypatch: pytest.MonkeyPatch) -> None: - for gib in (1, 4, 16, 64, 256): - physical = gib * 1024**3 - assert self._limit_for(physical, monkeypatch) <= int(physical * 0.8) - def test_connection_carries_the_cap_and_a_corpus_local_spill_dir( - self, query_corpus: Path, capsys: pytest.CaptureFixture[str] +class TestQuerySpillDirectory: + """The spill directory is private, outside the corpus, and gone when the process is.""" + + @pytest.fixture + def recorded_mkdtemp(self, monkeypatch: pytest.MonkeyPatch) -> list[tuple[Path, int]]: + import tempfile + + created: list[tuple[Path, int]] = [] + real = tempfile.mkdtemp + + def mkdtemp(*args: Any, **kwargs: Any) -> str: + path = str(real(*args, **kwargs)) + created.append((Path(path), Path(path).stat().st_mode & 0o777)) + return path + + # ``app`` reads ``tempfile.mkdtemp`` at call time, so the module attribute is the seam. + monkeypatch.setattr(tempfile, "mkdtemp", mkdtemp) + return created + + def test_private_dir_outside_the_corpus_removed_on_success( + self, + query_corpus: Path, + recorded_mkdtemp: list[tuple[Path, int]], + capsys: pytest.CaptureFixture[str], ) -> None: - """The settings are read back through caller SQL — the only honest witness.""" - from atif_cli.app import _query_memory_limit_bytes + _run_query( + "SELECT current_setting('temp_directory') AS tmp, " + "current_setting('allowed_directories') AS dirs", + query_corpus, + ) + row = json.loads(capsys.readouterr().out)[0] + [(spill, mode)] = recorded_mkdtemp + assert row["tmp"] == str(spill) + assert [Path(d) for d in row["dirs"]] == [spill] + assert spill.name.startswith("atif-sql-query-") + assert mode == 0o700 + assert not spill.is_relative_to(query_corpus) + assert not spill.exists() + assert not (query_corpus / ".duckdb_tmp").exists() + + def test_removed_on_error_too( + self, + query_corpus: Path, + recorded_mkdtemp: list[tuple[Path, int]], + capsys: pytest.CaptureFixture[str], + ) -> None: + with pytest.raises(SystemExit): + _run_query("SELEC 1", query_corpus) + capsys.readouterr() + [(spill, _)] = recorded_mkdtemp + assert not spill.exists() + + def test_a_query_process_leaves_the_corpus_tree_unchanged( + self, query_corpus: Path, tmp_path: Path + ) -> None: + """Through a real process: the review's probe d3, then a plain read. + + Before the fix the COPY exited 0 and ``probe.csv`` persisted inside + the corpus. Now the tree digest (paths, modes, bytes) is unchanged, + the pre-created legacy dir stays empty, no spill dir is left under + the process's TMPDIR, and a fresh corpus gains no ``.duckdb_tmp``. + """ + legacy = query_corpus / ".duckdb_tmp" + legacy.mkdir() + scratch = tmp_path / "process-tmp" + scratch.mkdir() + env: dict[str, str] = dict(os.environ) + env.update( + TMPDIR=str(scratch), + NO_COLOR="1", + PYTHONHASHSEED="0", + LITELLM_LOCAL_MODEL_COST_MAP="true", + ) + for var in ("ATIF_SQL_CORPUS_ROOT", "ATIF_SQL_LANCE_URI", "ATIF_SQL_EMBED_MODEL_ID"): + env.pop(var, None) + before = _tree_digest(query_corpus) + + def run(sql: str) -> subprocess.CompletedProcess[str]: + argv = [sys.executable, "-m", "atif_cli", "query", "--format", "json"] + argv += ["--corpus-root", str(query_corpus), sql] + return subprocess.run( # noqa: S603 - argv is built from constants and tmp paths + argv, capture_output=True, text=True, env=env, timeout=300, check=False + ) + + refused = run(f"COPY (SELECT 1) TO '{legacy}/probe.csv'") + assert refused.returncode == EXIT_CODES["sandbox_refused"], refused.stderr + assert json.loads(refused.stderr.strip().splitlines()[-1])["error"]["kind"] == ( + "sandbox_refused" + ) + read = run("SELECT count(*) AS n FROM sessions") + assert read.returncode == 0, read.stderr + assert json.loads(read.stdout) == [{"n": 1}] + + assert _tree_digest(query_corpus) == before + assert list(legacy.iterdir()) == [] + assert list(scratch.iterdir()) == [], "a spill dir outlived its process" + + legacy.rmdir() + fresh = run("SELECT count(*) AS n FROM steps") + assert fresh.returncode == 0, fresh.stderr + assert not (query_corpus / ".duckdb_tmp").exists() + + +class TestQueryResources: + """The cap and thread count are derived from the host and applied BEFORE registration.""" + + GIB = 1024**3 + + @staticmethod + def _host(monkeypatch: pytest.MonkeyPatch, physical: int, available: int) -> None: + monkeypatch.setattr(app_module, "_host_memory", lambda: (physical, available)) + def test_large_host_gets_half_of_ram(self, monkeypatch: pytest.MonkeyPatch) -> None: + self._host(monkeypatch, 124 * self.GIB, 89 * self.GIB) + assert app_module._query_memory_limit_bytes() == 62 * self.GIB + + def test_eight_gib_guest_is_capped_below_its_own_ram( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The old rule said 8 GiB here; the guest had 8 GiB in total and 6 free.""" + self._host(monkeypatch, 8 * self.GIB, 6 * self.GIB) + limit = app_module._query_memory_limit_bytes() + assert limit == int(6 * self.GIB * 0.8) + assert limit < 8 * self.GIB + + def test_target_wins_between_half_and_eighty_percent( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + self._host(monkeypatch, 12 * self.GIB, 12 * self.GIB) + assert app_module._query_memory_limit_bytes() == 8 * self.GIB + + def test_never_exceeds_eighty_percent_of_physical_or_available( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + for gib in (1, 4, 8, 16, 64, 256): + for free_fraction in (0.1, 0.5, 1.0): + physical = gib * self.GIB + available = int(physical * free_fraction) + self._host(monkeypatch, physical, available) + limit = app_module._query_memory_limit_bytes() + assert limit <= max(int(physical * 0.8), app_module._QUERY_MEMORY_MIN_BYTES) + assert limit <= max(int(available * 0.8), app_module._QUERY_MEMORY_MIN_BYTES) + + def test_floor_when_almost_nothing_is_free(self, monkeypatch: pytest.MonkeyPatch) -> None: + self._host(monkeypatch, 2 * self.GIB, 100 * 1024**2) + assert app_module._query_memory_limit_bytes() == app_module._QUERY_MEMORY_MIN_BYTES + + @staticmethod + def _sixteen_cpus(_pid: int) -> set[int]: + return set(range(16)) + + def test_threads_follow_the_cap(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(os, "sched_getaffinity", self._sixteen_cpus, raising=False) + assert app_module._query_threads(int(4.8 * self.GIB)) == 2 + assert app_module._query_threads(62 * self.GIB) == 16 + assert app_module._query_threads(1 * self.GIB) == 1 + + def test_host_memory_reads_this_host(self) -> None: + physical, available = app_module._host_memory() + assert physical > 0 + assert 0 < available <= physical + + @pytest.mark.parametrize( + ("text", "expected"), + [ + ("6GB", 6 * 10**9), + ("512MiB", 512 * 1024**2), + ("123B", 123), + ("1.5GiB", int(1.5 * 1024**3)), + (" 6 gb ", 6 * 10**9), + ("4096", 4096), + ], + ) + def test_parse_size(self, text: str, expected: int) -> None: + assert app_module._parse_size(text) == expected + + @pytest.mark.parametrize("text", ["", "abc", "6XB", "GB", "-1GB"]) + def test_parse_size_rejects_garbage(self, text: str) -> None: + with pytest.raises(ValueError, match=r"size|unit"): + app_module._parse_size(text) + + def test_env_overrides(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(os, "sched_getaffinity", self._sixteen_cpus, raising=False) + monkeypatch.setenv(QUERY_MEMORY_LIMIT_ENV, "6GB") + monkeypatch.setenv(QUERY_THREADS_ENV, "4") + assert app_module._query_resources() == app_module.QueryResources(6 * 10**9, 4) + monkeypatch.delenv(QUERY_THREADS_ENV) + assert app_module._query_resources().threads == 2, "threads follow the overridden cap" + + @pytest.mark.parametrize( + ("var", "value"), + [(QUERY_MEMORY_LIMIT_ENV, "lots"), (QUERY_THREADS_ENV, "0"), (QUERY_THREADS_ENV, "many")], + ) + def test_malformed_override_exits_64( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + var: str, + value: str, + ) -> None: + monkeypatch.setenv(var, value) + with pytest.raises(SystemExit) as excinfo: + _run_query("SELECT 1", query_corpus) + assert excinfo.value.code == EXIT_CODES["invalid_input"] + err = json.loads(capsys.readouterr().err)["error"] + assert err["kind"] == "invalid_input" + assert var in err["hint"] + + def test_connection_carries_the_cap_and_threads( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + """The settings are read back through caller SQL — the only honest witness.""" + monkeypatch.setenv(QUERY_MEMORY_LIMIT_ENV, "3GiB") + monkeypatch.setenv(QUERY_THREADS_ENV, "3") _run_query( "SELECT current_setting('memory_limit') AS mem, " - "current_setting('temp_directory') AS tmp", + "current_setting('threads') AS threads, " + "current_setting('autoinstall_known_extensions') AS autoinstall, " + "current_setting('autoload_known_extensions') AS autoload", query_corpus, ) row = json.loads(capsys.readouterr().out)[0] - assert row["tmp"] == str(query_corpus / ".duckdb_tmp") + assert row["mem"] == "3.0 GiB" + assert row["mem"] != duckdb_default_memory_limit() + assert int(row["threads"]) == 3 + assert row["autoinstall"] is False + assert row["autoload"] is False + + @pytest.mark.parametrize(("threads", "memory"), [(4, "6GB"), (16, "8GB")]) + def test_registration_fits_under_limits_that_used_to_oom( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + threads: int, + memory: str, + ) -> None: + """Finding 1 of the MicroVM review, as the review reproduced it on a workstation. + + ``SET threads=N; SET memory_limit=...`` on every connection before the + CLI touches it (the review's shim), plus the same values through the + env so the CLI applies them too. Before the fix the edges reader + reserved 2 GiB a thread and every query exited 70 out of memory once + the glob matched two or more files (measured on main: one file + registers, two do not); the readers are now sized from the files. + """ + import duckdb + + for var in ("ATIF_SQL_CORPUS_ROOT", "ATIF_SQL_LANCE_URI", "ATIF_SQL_EMBED_MODEL_ID"): + monkeypatch.delenv(var, raising=False) + corpus = _write_two_session_corpus(tmp_path / "corpus") + monkeypatch.setenv(QUERY_MEMORY_LIMIT_ENV, memory) + monkeypatch.setenv(QUERY_THREADS_ENV, str(threads)) + real_connect = duckdb.connect + + def limited(*args: Any, **kwargs: Any) -> Any: + con = real_connect(*args, **kwargs) + con.execute(f"SET threads={threads}; SET memory_limit='{memory}'") + return con + + monkeypatch.setattr(duckdb, "connect", limited) + _run_query("SELECT count(*) AS n FROM sessions", corpus) + assert json.loads(capsys.readouterr().out) == [{"n": 2}] + + +class TestQueryNeverInstallsExtensions: + """Registration must not reach the network, even when an embeddings store exists.""" + + def test_store_present_and_extension_absent_binds_empty_without_installing( + self, + query_corpus: Path, + empty_extension_dir: tuple[Path, list[str]], + capsys: pytest.CaptureFixture[str], + ) -> None: + ext_dir, statements = empty_extension_dir + _write_lance_store(query_corpus / "embeddings_lance") + warnings, sink_id = _capture_warnings() + try: + _run_query("SELECT count(*) AS n FROM message_embeddings", query_corpus) + finally: + logger.remove(sink_id) + assert json.loads(capsys.readouterr().out) == [{"n": 0}] + assert not any(re.match(r"\s*INSTALL\b", s, re.IGNORECASE) for s in statements) + assert list(ext_dir.rglob("*")) == [], "something was installed into the extension dir" + assert any("--install-extension" in w for w in warnings) + + def test_no_store_means_no_load_at_all( + self, + query_corpus: Path, + empty_extension_dir: tuple[Path, list[str]], + capsys: pytest.CaptureFixture[str], + ) -> None: + _, statements = empty_extension_dir + _run_query("SELECT count(*) AS n FROM message_embeddings", query_corpus) + assert json.loads(capsys.readouterr().out) == [{"n": 0}] + assert not any(re.match(r"\s*(INSTALL|LOAD)\b", s, re.IGNORECASE) for s in statements) + + def test_status_reports_the_missing_extension( + self, + query_corpus: Path, + tmp_path: Path, + empty_extension_dir: tuple[Path, list[str]], + capsys: pytest.CaptureFixture[str], + ) -> None: + _write_lance_store(query_corpus / "embeddings_lance") + source = tmp_path / "source" + source.mkdir() + status(source_root=source, corpus_root=query_corpus, fmt=OutputFormat.JSON) + payload = json.loads(capsys.readouterr().out) + assert payload["lance_extension_installed"] is False + assert payload["embeddings_store_present"] is True + assert payload["vector_search"] == "extension_missing" + + def test_search_exits_78_instead_of_pretending_the_store_is_empty( + self, + query_corpus: Path, + empty_extension_dir: tuple[Path, list[str]], + capsys: pytest.CaptureFixture[str], + ) -> None: + _write_lance_store(query_corpus / "embeddings_lance") + with pytest.raises(SystemExit) as excinfo: + search("anything", corpus_root=query_corpus, fmt=OutputFormat.JSON) + assert excinfo.value.code == EXIT_CODES["extension_missing"] == 78 + err = json.loads(capsys.readouterr().err)["error"] + assert err["kind"] == "extension_missing" + assert "--install-extension" in err["hint"] + + def test_status_reports_ready_and_no_store( + self, query_corpus: Path, tmp_path: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + source = tmp_path / "source" + source.mkdir() + status(source_root=source, corpus_root=query_corpus, fmt=OutputFormat.JSON) + payload = json.loads(capsys.readouterr().out) + assert payload["lance_extension_installed"] is True + assert payload["vector_search"] == "no_store" + _write_lance_store(query_corpus / "embeddings_lance") + status(source_root=source, corpus_root=query_corpus, fmt=OutputFormat.JSON) + assert json.loads(capsys.readouterr().out)["vector_search"] == "ready" + + def test_embed_install_extension_installs_and_exits( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + """The explicit install path: no scope flag needed, no Bedrock, JSON receipt.""" + from atif_duck.infrastructure import registry as registry_mod + + calls: list[object] = [] - expected_gib = _query_memory_limit_bytes() / 1024**3 - reported = row["mem"] - assert reported.endswith("GiB"), reported - assert float(reported.removesuffix(" GiB")) == pytest.approx(expected_gib, abs=0.1) - assert reported != duckdb_default_memory_limit() + def fake_install(con: object) -> str: + calls.append(con) + return "/ext/lance.duckdb_extension" + + monkeypatch.setattr(registry_mod, "install_lance_extension", fake_install) + embed(install_extension=True, corpus_root=query_corpus, fmt=OutputFormat.JSON) + assert json.loads(capsys.readouterr().out) == { + "extension": "lance", + "installed": True, + "install_path": "/ext/lance.duckdb_extension", + } + assert len(calls) == 1 + assert not (query_corpus / "embeddings_lance").exists() class TestQueryLockedConfiguration: diff --git a/packages/atif-cli/tests/test_vss_commands.py b/packages/atif-cli/tests/test_vss_commands.py index efdb051..eabbf95 100644 --- a/packages/atif-cli/tests/test_vss_commands.py +++ b/packages/atif-cli/tests/test_vss_commands.py @@ -161,6 +161,79 @@ def test_dry_run_needs_no_scope(self, corpus: Path, capsys: pytest.CaptureFixtur payload = json.loads(capsys.readouterr().out) assert payload["dry_run"] is True + @pytest.mark.usefixtures("_fake_cohere") + def test_real_run_installs_the_lance_extension_first( + self, corpus: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + """The network is already a deliberate act here, so this is where the download belongs.""" + from atif_duck.infrastructure import registry as registry_mod + + calls: list[object] = [] + + def recording_install(con: object) -> str: + calls.append(con) + return "/ext/lance" + + monkeypatch.setattr(registry_mod, "install_lance_extension", recording_install) + embed(all_steps=True, corpus_root=corpus, fmt=OutputFormat.JSON) + assert json.loads(capsys.readouterr().out)["rows_processed"] == 1 + assert len(calls) == 1 + + def test_dry_run_installs_nothing( + self, corpus: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + from atif_duck.infrastructure import registry as registry_mod + + def forbidden_install(con: object) -> str: + del con + pytest.fail("a dry run must not reach the extension repository") + + monkeypatch.setattr(registry_mod, "install_lance_extension", forbidden_install) + embed(dry_run=True, corpus_root=corpus, fmt=OutputFormat.JSON) + assert json.loads(capsys.readouterr().out)["dry_run"] is True + + @pytest.mark.usefixtures("_fake_cohere") + def test_a_failed_install_only_warns_and_the_backfill_still_runs( + self, corpus: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + import duckdb + from loguru import logger + + from atif_duck.infrastructure import registry as registry_mod + + def failing(con: object) -> str: + del con + problem = "Failed to download extension" + raise duckdb.IOException(problem) + + monkeypatch.setattr(registry_mod, "install_lance_extension", failing) + warnings: list[str] = [] + sink_id = logger.add(lambda message: warnings.append(str(message)), level="WARNING") + try: + embed(all_steps=True, corpus_root=corpus, fmt=OutputFormat.JSON) + finally: + logger.remove(sink_id) + assert json.loads(capsys.readouterr().out)["rows_processed"] == 1 + assert any("--install-extension" in w for w in warnings) + + def test_install_extension_flag_exits_70_when_the_download_fails( + self, corpus: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + import duckdb + + from atif_duck.infrastructure import registry as registry_mod + + def failing(con: object) -> str: + del con + problem = "Failed to download extension" + raise duckdb.IOException(problem) + + monkeypatch.setattr(registry_mod, "install_lance_extension", failing) + with pytest.raises(SystemExit) as excinfo: + embed(install_extension=True, corpus_root=corpus, fmt=OutputFormat.JSON) + assert excinfo.value.code == EXIT_CODES["runtime_error"] + assert json.loads(capsys.readouterr().err)["error"]["kind"] == "runtime_error" + class TestSearchCommand: def test_empty_store_exits_2(self, corpus: Path, capsys: pytest.CaptureFixture[str]) -> None: diff --git a/packages/atif-corpus/tests/test_session_id.py b/packages/atif-corpus/tests/test_session_id.py index 5fbba60..83f0fab 100644 --- a/packages/atif-corpus/tests/test_session_id.py +++ b/packages/atif-corpus/tests/test_session_id.py @@ -220,3 +220,45 @@ def test_a_clean_pass_reports_no_rejections(self, source_root: Path, corpus_root report = _materialize(source_root, corpus_root) assert report.rejected_session_ids == () assert report.rejected_count == 0 + + +class TestFilenamesThatOnceDestroyedTheCorpus: + """Findings 6 and 7 of the MicroVM review, exactly as reproduced there. + + A transcript named ``..jsonl`` has the stem ``.``; ``sessions_dir / "."`` + IS ``sessions_dir``, so the atomic swap once renamed the whole sessions + tree aside and deleted it, with exit 0 and a report that said nothing. A + transcript named ``*.jsonl`` materialized fine and then multiplied every + other session's rows in the readers. Both names now stop at the boundary. + """ + + VICTIMS = ("victim-1", "victim-2", "victim-3") + + def test_dot_jsonl_and_star_jsonl_leave_the_materialized_tree_alone( + self, source_root: Path, corpus_root: Path + ) -> None: + for sid in self.VICTIMS: + write_session(source_root, sid, mtime_ns=STALE_NS) + first = _materialize(source_root, corpus_root) + assert first.materialized_count == len(self.VICTIMS) + layout = CorpusLayout(corpus_root=corpus_root) + before = sorted(p.name for p in layout.sessions_dir.iterdir()) + assert before == sorted(self.VICTIMS) + + # ``write_session`` names the file ``.jsonl``: these are ``..jsonl`` and ``*.jsonl``. + for hostile in (".", "*"): + write_session(source_root, hostile, mtime_ns=STALE_NS) + assert (source_root / "-proj-a" / "..jsonl").is_file() + assert (source_root / "-proj-a" / "*.jsonl").is_file() + + second = _materialize(source_root, corpus_root) + assert second.rejected_session_ids == ("*", ".") + assert second.rejected_count == 2 + assert second.materialized_count == 0 + assert second.failed_count == 0 + assert second.sessions_removed == 0 + assert sorted(p.name for p in layout.sessions_dir.iterdir()) == before + for sid in self.VICTIMS: + assert (layout.sessions_dir / sid / "meta.json").is_file() + # No loose artifact files at the sessions root, which is what the swap left behind. + assert all(p.is_dir() for p in layout.sessions_dir.iterdir()) diff --git a/packages/atif-duck/src/atif_duck/infrastructure/registry.py b/packages/atif-duck/src/atif_duck/infrastructure/registry.py index 58e4a96..04987d2 100644 --- a/packages/atif-duck/src/atif_duck/infrastructure/registry.py +++ b/packages/atif-duck/src/atif_duck/infrastructure/registry.py @@ -100,7 +100,7 @@ ) if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Iterable, Sequence from pathlib import Path import duckdb @@ -166,13 +166,51 @@ def coverage(self) -> ColumnarCoverage: ) -#: Inlined ``read_json`` upper bound. Live trajectory.json files reach 436 MB -#: because harbor inlines subagent sidechains and tool outputs, so 1 GiB is -#: roughly 2.3x headroom over the largest observed document — not a knob to -#: trim. Lowering it does not reduce registration memory: measured peak RSS is -#: unchanged at 512 MiB and rises at 128 MiB, because this bounds the parse -#: buffer rather than preallocating per thread. -_MAX_OBJECT_SIZE: int = 1_073_741_824 +#: Ceiling for a ``read_json`` ``maximum_object_size``. Live trajectory.json +#: files reach 436 MB because harbor inlines subagent sidechains and tool +#: outputs; 1 GiB is about 2.3x headroom over the largest observed document. +_OBJECT_SIZE_CAP: int = 1_073_741_824 + +#: Floor for the same bound: DuckDB's own default. +_OBJECT_SIZE_FLOOR: int = 16_777_216 + +#: Headroom over the largest file when sizing the bound: a quarter of the file +#: plus one MiB, so a file that grows a little between the stat and the read +#: still parses. +_OBJECT_SIZE_HEADROOM_DIVISOR: int = 4 +_OBJECT_SIZE_HEADROOM_BYTES: int = 1_048_576 + +#: The DuckDB extension that reads the Lance embeddings store. +LANCE_EXTENSION: str = "lance" + + +def _object_size_bound(paths: Iterable[Path]) -> int: + """``maximum_object_size`` for a ``read_json`` over ``paths``, sized from the files. + + The bound is not free. DuckDB's eager JSON reader reserves about twice + this many bytes PER THREAD before it parses anything, so the 1 GiB + constant this replaced cost 2 GiB a thread: over 300 ``edges.jsonl`` + files of under 8 MB each it exhausted a 6 GB ``memory_limit`` at four + threads and a 25 GB one at sixteen, while 16 MiB read the same rows in + 0.10 s (measured on the 300-session panel corpus). An earlier note here + claimed lowering it changed nothing; that measurement held the thread + count at one. + + A newline-delimited file's objects are its lines and a one-document file + IS its object, so the largest file present bounds every object either + reader meets. That size plus headroom, floored at DuckDB's default and + capped at :data:`_OBJECT_SIZE_CAP`, is the bound. A path that vanished + between the listing and the stat counts as zero. + """ + largest = 0 + for path in paths: + try: + largest = max(largest, path.stat().st_size) + except OSError: + continue + with_headroom = largest + largest // _OBJECT_SIZE_HEADROOM_DIVISOR + _OBJECT_SIZE_HEADROOM_BYTES + return max(_OBJECT_SIZE_FLOOR, min(_OBJECT_SIZE_CAP, with_headroom)) + # Explicit projection for ``v_raw_trajectories``. # @@ -618,7 +656,9 @@ def register_raw(con: duckdb.DuckDBPyConnection, corpus_root: Path) -> RawSource # the columnar sessions' trajectory.json too, which is the cost # this whole arrangement exists to avoid. The list is the # statement's one parameter. - json_files = [str(sessions_dir / sid / "trajectory.json") for sid in json_ids] + trajectory_paths = [sessions_dir / sid / "trajectory.json" for sid in json_ids] + json_files = [str(path) for path in trajectory_paths] + trajectory_bound = _object_size_bound(trajectory_paths) con.execute( f""" CREATE OR REPLACE TEMP TABLE {_RAW_TRAJECTORIES_JSON_TABLE} AS @@ -631,9 +671,9 @@ def register_raw(con: duckdb.DuckDBPyConnection, corpus_root: Path) -> RawSource format='auto', filename=true, columns={{{_render_columns_clause(_TRAJECTORY_COLUMNS)}}}, - maximum_object_size={_MAX_OBJECT_SIZE} + maximum_object_size={int(trajectory_bound)} ); - """, # noqa: S608 # nosec B608 - file list is a bound parameter; table/columns/cap are constants + """, # noqa: S608 # nosec B608 - file list is a bound parameter; table/columns are constants; the bound is an int [json_files], ) logger.debug( @@ -663,10 +703,10 @@ def register_raw(con: duckdb.DuckDBPyConnection, corpus_root: Path) -> RawSource # edges.jsonl is one line per RAW record -> newline_delimited. # The record's own `source_file` (the raw transcript path) is kept; - # the edges.jsonl path itself is aliased to `edges_path`. - # edges.jsonl is one line per RAW record -> newline_delimited. - # The record's own `source_file` (the raw transcript path) is kept; - # the edges.jsonl path itself is aliased to `edges_path`. + # the edges.jsonl path itself is aliased to `edges_path`. The glob + # reads every matching file (the meta gate filters rows afterwards), + # so the bound is sized over every file the glob can reach. + edges_bound = _object_size_bound(sessions_dir.glob("*/edges.jsonl")) con.execute( f""" CREATE OR REPLACE TEMP TABLE {_RAW_EDGES_TABLE} AS @@ -680,10 +720,10 @@ def register_raw(con: duckdb.DuckDBPyConnection, corpus_root: Path) -> RawSource format='newline_delimited', filename=true, columns={{{_render_columns_clause(_EDGE_COLUMNS)}}}, - maximum_object_size={_MAX_OBJECT_SIZE} + maximum_object_size={int(edges_bound)} ) ) WHERE {meta_gate}; - """, # noqa: S608 # nosec B608 - glob is a bound parameter; table/columns/gate are constants + """, # noqa: S608 # nosec B608 - glob is a bound parameter; table/columns/gate are constants; the bound is an int [edges_glob], ) logger.debug("Registered {} from glob {}", _RAW_EDGES_TABLE, edges_glob) @@ -1186,6 +1226,50 @@ def _lance_table_present(con: duckdb.DuckDBPyConnection) -> bool: return row is not None and int(row[0]) > 0 +def lance_extension_installed(con: duckdb.DuckDBPyConnection) -> bool: + """True when the lance extension is present in this connection's extension directory. + + Read from ``duckdb_extensions()``, which lists the local directory and + never reaches the network, so the answer costs a few milliseconds and no + download. + """ + row = con.execute( + "SELECT installed FROM duckdb_extensions() WHERE extension_name = ?", [LANCE_EXTENSION] + ).fetchone() + return row is not None and bool(row[0]) + + +def load_lance_extension(con: duckdb.DuckDBPyConnection) -> bool: + """``LOAD`` the lance extension when it is installed; never install it. + + Returns ``False`` when the extension is absent, and executes no + ``INSTALL`` in that case whatever ``autoinstall_known_extensions`` says, + because the installed check runs first. Registration at query time must + never reach the network: the extension is 242 MB, fetched by name from + the default repository, and ``query`` runs unattended. + """ + if not lance_extension_installed(con): + return False + con.execute(f"LOAD {LANCE_EXTENSION};") + return True + + +def install_lance_extension(con: duckdb.DuckDBPyConnection) -> str: + """``INSTALL`` then ``LOAD`` the lance extension; return its install path. + + The one place atif-duck reaches the network, so it belongs to explicit + commands only (``atif-sql embed --install-extension``, or a real embed + run, which reaches Bedrock anyway). A query-time registration calls + :func:`load_lance_extension` instead. + """ + con.execute(f"INSTALL {LANCE_EXTENSION};") + con.execute(f"LOAD {LANCE_EXTENSION};") + row = con.execute( + "SELECT install_path FROM duckdb_extensions() WHERE extension_name = ?", [LANCE_EXTENSION] + ).fetchone() + return str(row[0]) if row and row[0] is not None else "" + + def register_vss( con: duckdb.DuckDBPyConnection, *, @@ -1198,7 +1282,14 @@ def register_vss( LanceDB stores embeddings + its IVF_HNSW_SQ index in one place (written by atif-embed's backfill); reads come back through DuckDB via the lance - core extension (``INSTALL lance; LOAD lance; ATTACH (TYPE LANCE)``). + core extension (``LOAD lance; ATTACH (TYPE LANCE)``). The extension is + only LOADed, never INSTALLed, here: :func:`load_lance_extension` checks + the local extension directory first and a missing extension degrades to + the empty fallback table with a warning, so registration never reaches + the network. Installing is an explicit act + (:func:`install_lance_extension`, behind ``atif-sql embed``). When no + store directory exists the extension is not loaded at all, which is the + state of every corpus that has not run ``embed``. The store probe runs through DuckDB itself. atif-duck declares no lancedb dependency: lancedb belongs to atif-embed, which writes the store, and the independence contract forbids an import edge between the two packages — so @@ -1231,18 +1322,25 @@ def register_vss( ------- bool ``True`` when the Lance table is reachable through the - ``message_embeddings`` view; ``False`` when no embeddings exist yet - (the name is created as an empty TABLE with the right schema so a + ``message_embeddings`` view; ``False`` when no embeddings exist yet, + or the store exists but the lance extension is not installed (the + name is created as an empty TABLE with the right schema so a downstream ``CREATE MACRO semantic_search`` can still bind). """ dim_i = int(dim) - con.execute("INSTALL lance;") - con.execute("LOAD lance;") import duckdb as _duckdb attached = False - if lance_uri.is_dir(): + if lance_uri.is_dir() and not load_lance_extension(con): + logger.warning( + "Lance store at {} cannot be read: the {} extension is not installed, so " + "message_embeddings binds empty. Run `atif-sql embed --install-extension` " + "(a one-time download) to enable vector search.", + lance_uri, + LANCE_EXTENSION, + ) + elif lance_uri.is_dir(): try: # ATTACH is one of the statement kinds DuckDB will not prepare, so # this is the one corpus path that still enters statement text; it diff --git a/packages/atif-duck/tests/duck_fixtures.py b/packages/atif-duck/tests/duck_fixtures.py index 001cd5c..9b4574a 100644 --- a/packages/atif-duck/tests/duck_fixtures.py +++ b/packages/atif-duck/tests/duck_fixtures.py @@ -707,3 +707,18 @@ def write_codex_session(root: Path) -> Path: ) ) return session_dir + + +# ``register_vss`` no longer installs the lance extension (that download at +# query time was finding 4 of the MicroVM review), so the suite installs it +# once up front. A no-op where it is already present; a one-time download on a +# fresh runner, exactly what every register() call used to do implicitly. +@pytest.fixture(scope="session", autouse=True) +def _lance_extension_present() -> None: # pyright: ignore[reportUnusedFunction] + import duckdb + + con = duckdb.connect() + try: + con.execute("INSTALL lance") + finally: + con.close() diff --git a/packages/atif-duck/tests/test_registration_limits.py b/packages/atif-duck/tests/test_registration_limits.py new file mode 100644 index 0000000..ee7bc3b --- /dev/null +++ b/packages/atif-duck/tests/test_registration_limits.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Registration fits the host, and never reaches the network. + +Two findings of the MicroVM review, pinned at the registry layer: + +* The eager ``read_json`` readers reserve about twice their + ``maximum_object_size`` per thread. At the old 1 GiB constant that was + 2 GiB a thread, so ``SET threads=4; SET memory_limit='6GB'`` (a 4 vCPU, + 8 GiB guest) failed to register a 131 MB corpus with an out-of-memory + error, and so did 16 threads under 25 GB. The bound is now sized from the + largest file the reader will open. +* ``register_vss`` used to ``INSTALL lance`` on every registration, a 242 MB + download from the extension repository. It now LOADs the extension only + when it is already installed and binds the vector surface empty otherwise. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import duckdb +import pytest +from duck_fixtures import SESSION_IDS +from loguru import logger +from test_vss import _write_lance + +from atif_duck.infrastructure import registry as registry_mod +from atif_duck.infrastructure.registry import ( + _OBJECT_SIZE_CAP, + _OBJECT_SIZE_FLOOR, + _object_size_bound, + install_lance_extension, + lance_extension_installed, + load_lance_extension, + register, + register_vss, +) + +MIB = 1024**2 + + +class TestRegistrationUnderHostLimits: + """The review's reproduction: limits set before ``register`` over the fixture corpus.""" + + @pytest.mark.parametrize(("threads", "memory"), [(4, "6GB"), (16, "8GB"), (16, "25GB")]) + def test_register_succeeds_under_limits_that_used_to_oom( + self, corpus_root: Path, threads: int, memory: str + ) -> None: + con = duckdb.connect() + try: + con.execute(f"SET threads={threads}; SET memory_limit='{memory}'") + sources = register(con, corpus_root) + assert set(sources.json_session_ids) == set(SESSION_IDS) + assert con.execute("SELECT count(*) FROM sessions").fetchone() == (2,) + messages = con.execute("SELECT count(*) FROM messages").fetchone() + assert messages is not None + assert int(messages[0]) > 0 + finally: + con.close() + + +class TestObjectSizeBound: + def test_no_files_gives_duckdbs_default(self) -> None: + assert _object_size_bound([]) == _OBJECT_SIZE_FLOOR == 16 * MIB + + def test_small_files_stay_at_the_floor(self, tmp_path: Path) -> None: + small = tmp_path / "edges.jsonl" + small.write_bytes(b"x" * 1000) + assert _object_size_bound([small]) == _OBJECT_SIZE_FLOOR + + def test_a_large_file_gets_a_quarter_plus_one_mib_of_headroom(self, tmp_path: Path) -> None: + big = tmp_path / "trajectory.json" + with big.open("wb") as handle: + handle.truncate(200 * MIB) # sparse: no bytes written + assert _object_size_bound([big]) == 200 * MIB + 50 * MIB + MIB + + def test_the_largest_file_wins(self, tmp_path: Path) -> None: + paths = [] + for i, size in enumerate((40 * MIB, 120 * MIB, 80 * MIB)): + path = tmp_path / f"{i}.json" + with path.open("wb") as handle: + handle.truncate(size) + paths.append(path) + assert _object_size_bound(paths) == 120 * MIB + 30 * MIB + MIB + + def test_capped_at_one_gib(self, tmp_path: Path) -> None: + huge = tmp_path / "huge.json" + with huge.open("wb") as handle: + handle.truncate(2 * 1024**3) + assert _object_size_bound([huge]) == _OBJECT_SIZE_CAP == 1024**3 + + def test_a_vanished_path_counts_as_zero(self, tmp_path: Path) -> None: + assert _object_size_bound([tmp_path / "gone.json"]) == _OBJECT_SIZE_FLOOR + + def test_the_bound_reaches_both_readers( + self, corpus_root: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The DDL carries the sized bound, not the old constant.""" + seen: list[int] = [] + real = registry_mod._object_size_bound + + def spy(paths: Any) -> int: + bound = real(paths) + seen.append(bound) + return bound + + monkeypatch.setattr(registry_mod, "_object_size_bound", spy) + con = duckdb.connect() + try: + register(con, corpus_root) + finally: + con.close() + assert len(seen) == 2, "one bound per eager JSON reader (trajectories, edges)" + assert all(bound == _OBJECT_SIZE_FLOOR for bound in seen) + + +class _Recording: + def __init__(self, con: duckdb.DuckDBPyConnection, statements: list[str]) -> None: + self._con = con + self._statements = statements + + def execute(self, sql: str, *args: Any, **kwargs: Any) -> Any: + self._statements.append(sql) + return self._con.execute(sql, *args, **kwargs) + + def __getattr__(self, name: str) -> Any: + return getattr(self._con, name) + + +def _connection_without_extensions(tmp_path: Path) -> tuple[Any, Path, list[str]]: + ext_dir = tmp_path / "no-extensions" + ext_dir.mkdir() + statements: list[str] = [] + con = duckdb.connect() + con.execute(f"SET extension_directory='{ext_dir}'") + return _Recording(con, statements), ext_dir, statements + + +def _installs(statements: list[str]) -> list[str]: + return [s for s in statements if s.lstrip().upper().startswith("INSTALL")] + + +def _loads(statements: list[str]) -> list[str]: + return [s for s in statements if s.lstrip().upper().startswith("LOAD")] + + +class TestVssNeverInstalls: + def test_store_present_extension_absent_binds_empty_and_installs_nothing( + self, tmp_path: Path + ) -> None: + store = _write_lance(tmp_path / "store") + con, ext_dir, statements = _connection_without_extensions(tmp_path) + warnings: list[str] = [] + sink_id = logger.add(lambda message: warnings.append(str(message)), level="WARNING") + try: + bound = register_vss(con, lance_uri=store, expected_model="test-embedder:1") + finally: + logger.remove(sink_id) + con.close() + assert bound is False + assert _installs(statements) == [] + assert _loads(statements) == [] + assert list(ext_dir.rglob("*")) == [] + assert any("--install-extension" in w for w in warnings) + + def test_autoinstall_left_on_still_downloads_nothing(self, tmp_path: Path) -> None: + """The installed check runs before LOAD, so DuckDB's default autoinstall never triggers.""" + store = _write_lance(tmp_path / "store") + con, ext_dir, _ = _connection_without_extensions(tmp_path) + assert con.execute("SELECT current_setting('autoinstall_known_extensions')").fetchone() == ( + True, + ) + try: + register_vss(con, lance_uri=store) + assert con.execute("SELECT count(*) FROM message_embeddings").fetchone() == (0,) + finally: + con.close() + assert list(ext_dir.rglob("*")) == [] + + def test_no_store_means_the_extension_is_not_even_loaded(self, tmp_path: Path) -> None: + con, _, statements = _connection_without_extensions(tmp_path) + try: + assert register_vss(con, lance_uri=tmp_path / "absent") is False + finally: + con.close() + assert _loads(statements) == [] + assert _installs(statements) == [] + + def test_installed_check_reads_the_extension_directory(self, tmp_path: Path) -> None: + con, _, _ = _connection_without_extensions(tmp_path) + try: + assert lance_extension_installed(con) is False + assert load_lance_extension(con) is False + finally: + con.close() + default = duckdb.connect() + try: + assert lance_extension_installed(default) is True + assert load_lance_extension(default) is True + finally: + default.close() + + def test_register_over_a_corpus_with_a_store_still_binds_every_other_view( + self, corpus_root: Path, tmp_path: Path + ) -> None: + _write_lance(corpus_root / "embeddings_lance") + con, _, statements = _connection_without_extensions(tmp_path) + try: + register(con, corpus_root) + assert con.execute("SELECT count(*) FROM sessions").fetchone() == (2,) + assert con.execute("SELECT count(*) FROM message_embeddings").fetchone() == (0,) + assert ( + con.execute( + "SELECT count(*) FROM duckdb_functions() WHERE function_name = 'semantic_search'" + ).fetchone()[0] + >= 1 + ) + finally: + con.close() + assert _installs(statements) == [] + + def test_explicit_install_returns_the_path(self) -> None: + """A no-op where the extension is present; the one INSTALL atif-duck still owns.""" + con = duckdb.connect() + try: + path = install_lance_extension(con) + finally: + con.close() + assert "lance" in path + + def test_registry_source_has_no_install_outside_the_explicit_helper(self) -> None: + import inspect + + source = inspect.getsource(registry_mod) + helper = inspect.getsource(install_lance_extension) + assert source.count('"INSTALL') + source.count("INSTALL {LANCE_EXTENSION}") == ( + helper.count('"INSTALL') + helper.count("INSTALL {LANCE_EXTENSION}") + ) diff --git a/packages/atif-duck/tests/test_sql_text_boundaries.py b/packages/atif-duck/tests/test_sql_text_boundaries.py index f97c091..9e5dc2d 100644 --- a/packages/atif-duck/tests/test_sql_text_boundaries.py +++ b/packages/atif-duck/tests/test_sql_text_boundaries.py @@ -335,6 +335,53 @@ def test_a_bad_dir_with_columnar_artifacts_is_still_rejected( assert all(bad not in str(path) for path in sources.lazy_read_paths) +class TestGlobShapedNamesDoNotMultiplyRows: + """Finding 7 of the MicroVM review, exactly as reproduced there. + + Before the boundary a session dir named ``*`` or ``???`` landed inside + the readers' path lists, DuckDB expanded it as a glob, and every other + session's rows were read twice: ``sessions`` had seven rows for four + ids and ``count(*) FROM steps`` was 14 where 8 was true. The name is now + rejected at the boundary, so a registration over such a corpus counts + each session exactly once, on both read paths, and grants no glob. + """ + + @staticmethod + def _per_session_counts(con: duckdb.DuckDBPyConnection) -> dict[str, tuple[int, int]]: + rows = _rows( + con, + "SELECT s.session_id, s.step_count, (SELECT count(*) FROM steps t " + "WHERE t.session_id = s.session_id) FROM sessions s ORDER BY 1", + ) + return {str(r[0]): (int(r[1]), int(r[2])) for r in rows} + + @pytest.mark.parametrize("name", ["*", "???"]) + def test_each_session_is_counted_once_on_both_paths(self, tmp_path: Path, name: str) -> None: + clean = duckdb.connect(":memory:") + register(clean, build_corpus(tmp_path / "clean")) + expected = self._per_session_counts(clean) + assert set(expected) == set(SESSION_IDS) + + root = build_corpus(tmp_path / "hostile") + _write_session( + root / "sessions" / name, name, _adversarial_trajectory() | {"session_id": name} + ) + for path_label in ("json", "columnar"): + if path_label == "columnar": + add_columnar(root, tuple(SESSION_IDS)) + con = duckdb.connect(":memory:") + sources = register(con, root) + assert sources.rejected_session_ids == (name,), path_label + assert self._per_session_counts(con) == expected, path_label + assert _rows(con, "SELECT count(*) FROM sessions") == [(len(SESSION_IDS),)] + assert _rows(con, "SELECT count(*) FROM steps") == [ + (sum(steps for steps, _ in expected.values()),) + ] + assert all(not any(ch in str(p) for ch in "*?[") for p in sources.lazy_read_paths), ( + "a glob metacharacter reached the sandbox's file allowlist" + ) + + class TestRemainingSqlLiteralSites: """The two statements DuckDB will not prepare; each fails without ``sql_literal``.""" From 8489c90bd2e1a59316c882fb3bfb06effa796cb4 Mon Sep 17 00:00:00 2001 From: Laith Al-Saadoon Date: Sun, 13 Sep 2026 01:16:17 +0000 Subject: [PATCH 2/2] test: register the lance install fixture under a public name conftest.py re-exports the fixture modules with `import *`, which skips underscore names, so `_lance_extension_present` never registered and a runner with no cached extension (CI) ran the store-binding tests without lance. Reproduced locally with an empty HOME: 6 failures, then 407 passed once the fixture had a public name and installed the extension itself. --- packages/atif-cli/tests/cli_fixtures.py | 5 ++++- packages/atif-duck/tests/duck_fixtures.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/packages/atif-cli/tests/cli_fixtures.py b/packages/atif-cli/tests/cli_fixtures.py index 75818cf..73f2373 100644 --- a/packages/atif-cli/tests/cli_fixtures.py +++ b/packages/atif-cli/tests/cli_fixtures.py @@ -207,8 +207,11 @@ def write_analytics_parquets(corpus_root: Path, session_id: str) -> list[Path]: # query time was finding 4 of the MicroVM review), so the suite installs it # once up front. A no-op where it is already present; a one-time download on a # fresh runner, exactly what every register() call used to do implicitly. +# A public name on purpose: ``conftest.py`` re-exports this module with +# ``import *``, which skips underscore names, so an underscore here would leave +# the fixture unregistered (it did, on a runner with no extension cached). @pytest.fixture(scope="session", autouse=True) -def _lance_extension_present() -> None: # pyright: ignore[reportUnusedFunction] +def lance_extension_present() -> None: import duckdb con = duckdb.connect() diff --git a/packages/atif-duck/tests/duck_fixtures.py b/packages/atif-duck/tests/duck_fixtures.py index 9b4574a..e8f11b5 100644 --- a/packages/atif-duck/tests/duck_fixtures.py +++ b/packages/atif-duck/tests/duck_fixtures.py @@ -713,8 +713,11 @@ def write_codex_session(root: Path) -> Path: # query time was finding 4 of the MicroVM review), so the suite installs it # once up front. A no-op where it is already present; a one-time download on a # fresh runner, exactly what every register() call used to do implicitly. +# A public name on purpose: ``conftest.py`` re-exports this module with +# ``import *``, which skips underscore names, so an underscore here would leave +# the fixture unregistered (it did, on a runner with no extension cached). @pytest.fixture(scope="session", autouse=True) -def _lance_extension_present() -> None: # pyright: ignore[reportUnusedFunction] +def lance_extension_present() -> None: import duckdb con = duckdb.connect()