From 469fc156bfc97a41b53dd8962dbf7fb43584d9b6 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 22:02:11 -1000 Subject: [PATCH 1/4] perf(store): the local SQLite document store answers indexed lookups from its indexes A filter that pins the leading fields of one of a collection's indexes by equality now reads only the rows that index finds, instead of parsing every document. find_one, query, count and the update, replace and delete paths all read through it, so VersionedEntityStore.next_version() and put_bundled stop scanning. The filter and sort still run in Python on those rows, in rowid order, so results are unchanged. Indexes are read from sqlite_master, so no migration: existing stores use the indexes they already have. The per-table index cache drops when PRAGMA schema_version changes, so a long-lived reader starts searching an index another process creates. A write that waits out the busy timeout now raises DatabaseLockedError naming the database file, which the CLI prints as one line. The first open's switch to WAL raises it too, so @local routing no longer turns a temporary lock into "move it aside". Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/cli/__init__.py | 11 +- .../document_store/sqlite_document_store.py | 92 ++++++-- tst/unit/cli/user_facing_errors_test.py | 4 +- tst/unit/store/sqlite_document_store_test.py | 215 +++++++++++++++++- 4 files changed, 299 insertions(+), 23 deletions(-) diff --git a/src/agent_env/cli/__init__.py b/src/agent_env/cli/__init__.py index e39688d4..e946f6ee 100644 --- a/src/agent_env/cli/__init__.py +++ b/src/agent_env/cli/__init__.py @@ -4,6 +4,7 @@ from agent_env.plugins._cli import load_cli_plugins, load_cli_root_options from agent_env.store.base import NotFoundError +from agent_env.store.document_store.sqlite_document_store import DatabaseLockedError from agent_env.store.routing import namespace_routing from agent_env.utils.docker_build import DockerBuildError @@ -18,11 +19,11 @@ from .up import up -# Errors that mean "not allowed", "not there" or "your build failed" rather than "agent-env is broken". -# Spelled out because there is no domain base class to catch: ConfigError subclasses ValueError (so does -# pydantic's ValidationError), while NotFoundError and DockerBuildError don't. A TypeError or a -# retry-exhausted DuplicateKeyError is a defect, and a defect keeps its traceback. -_USER_FACING_ERRORS = (ValueError, NotFoundError, DockerBuildError) +# Errors that mean "not allowed", "not there", "your build failed" or "locked" rather than "agent-env is +# broken". Spelled out because there is no domain base class to catch: ConfigError subclasses ValueError +# (so does pydantic's ValidationError), while NotFoundError, DockerBuildError and DatabaseLockedError don't. +# A TypeError or a retry-exhausted DuplicateKeyError is a defect, and a defect keeps its traceback. +_USER_FACING_ERRORS = (ValueError, NotFoundError, DockerBuildError, DatabaseLockedError) class _UserErrorsAreNotCrashes(click.Group): diff --git a/src/agent_env/store/document_store/sqlite_document_store.py b/src/agent_env/store/document_store/sqlite_document_store.py index a2e42f23..87a71c79 100644 --- a/src/agent_env/store/document_store/sqlite_document_store.py +++ b/src/agent_env/store/document_store/sqlite_document_store.py @@ -3,8 +3,10 @@ Stores each document as a JSON row in a per-collection table; SQLite provides uniqueness (via unique indexes on ``json_extract``), atomicity (transactions), and durability. ``Filter`` / ``UpdateSpec`` / ``Sort`` semantics are evaluated in -Python, mirroring ``mongo_document_store`` as the oracle. Datetimes are -serialized as ISO strings. +Python, mirroring ``mongo_document_store`` as the oracle; an equality on an +index's leading top-level fields only narrows the rows read, which takes those +fields to hold scalars (SQLite indexes an array whole, where Mongo indexes each +element). Datetimes are serialized as ISO strings. SQL safety note: every user-controlled *value* is bound as a ``?`` placeholder. The only interpolated tokens are SQL *identifiers* (table / index / column @@ -15,6 +17,7 @@ from __future__ import annotations +import itertools import json import logging import re @@ -31,6 +34,7 @@ from agent_env.store.document_store.document_store import ( DocumentStore, DuplicateKeyError, + Eq, Filter, Sort, UpdateSpec, @@ -63,6 +67,10 @@ def _enable_wal(conn: sqlite3.Connection) -> None: time.sleep(_WAL_SWITCH_RETRY_DELAY_SECONDS) +class DatabaseLockedError(Exception): + """Raised when another process holds the database's write lock for longer than the busy timeout.""" + + class LocalSqliteDocumentStore(DocumentStore): """DocumentStore backed by a single stdlib-``sqlite3`` database file.""" @@ -72,6 +80,8 @@ def __init__(self, path: str) -> None: self._connection: Optional[sqlite3.Connection] = None self._lock = threading.RLock() self._tables: set[str] = set() + self._indexes: dict[str, list[tuple[bool, list[str]]]] = {} + self._schema_version: Optional[int] = None @property def path(self) -> Path: @@ -93,7 +103,8 @@ def _conn(self) -> sqlite3.Connection: conn = sqlite3.connect( str(self._path), timeout=_BUSY_TIMEOUT_SECONDS, check_same_thread=False, isolation_level=None ) - _enable_wal(conn) + with self._lock_errors(): + _enable_wal(conn) self._connection = conn return self._connection @@ -101,7 +112,7 @@ def find_one( self, collection: str, filter: Filter, sort: Optional[Sort] = None ) -> Optional[dict]: with self._lock: - matches = [doc for _, doc in self._load(collection) if evaluation.matches(doc, filter)] + matches = [doc for _, doc in self._load(collection, filter) if evaluation.matches(doc, filter)] matches = evaluation.sort_docs(matches, sort) return matches[0] if matches else None @@ -114,7 +125,7 @@ def query( offset: Optional[int] = None, ) -> list[dict]: with self._lock: - docs = [doc for _, doc in self._load(collection) if evaluation.matches(doc, filter)] + docs = [doc for _, doc in self._load(collection, filter) if evaluation.matches(doc, filter)] docs = evaluation.sort_docs(docs, sort) if offset: docs = docs[offset:] @@ -124,10 +135,10 @@ def query( def count(self, collection: str, filter: Filter) -> int: with self._lock: - return sum(1 for _, doc in self._load(collection) if evaluation.matches(doc, filter)) + return sum(1 for _, doc in self._load(collection, filter) if evaluation.matches(doc, filter)) def insert(self, collection: str, doc: dict) -> None: - with self._lock: + with self._lock, self._lock_errors(): tbl = self._ensure_table(collection) self._insert_raw(tbl, dict(doc)) @@ -218,7 +229,7 @@ def ensure_index( "LocalSqliteDocumentStore has no TTL support; ttl_seconds=%s on %s%s ignored — rows will not auto-expire.", ttl_seconds, collection, fields, ) - with self._lock: + with self._lock, self._lock_errors(): tbl = self._ensure_table(collection) raw_name = tbl + "_" + "_".join(fields) + ("_unique" if unique else "_index") name = self._safe_ident(re.sub(r"\W", "_", raw_name)) @@ -228,6 +239,7 @@ def ensure_index( # nosemgrep: sqlalchemy-execute-raw-query -- name/tbl/cols are validated identifiers f'CREATE {uniq}INDEX IF NOT EXISTS "{name}" ON "{tbl}" ({cols})' ) + self._indexes.pop(tbl, None) # --- sqlite plumbing --- @@ -272,25 +284,31 @@ def _adopt_if_created(self, tbl: str) -> bool: return True return False - def _load(self, collection: str) -> list[tuple[int, dict]]: + def _load(self, collection: str, filter: Filter) -> list[tuple[int, dict]]: + """The rows that can match ``filter``, in rowid order: those an index finds, else all of them.""" tbl = self._table(collection) if tbl not in self._tables and not self._adopt_if_created(tbl): return [] - return [ - (row[0], json.loads(row[1])) - # nosemgrep: sqlalchemy-execute-raw-query -- tbl is a validated identifier - for row in self._conn.execute(f'SELECT rowid, doc FROM "{tbl}"') - ] + pinned = self._pinned(tbl, filter) + # The index's own expression, so SQLite searches the index; the value is JSON so it decodes as the field does. + terms = [f"json_extract(doc, '$.{self._safe_path(field)}') = json_extract(?, '$')" for field, _ in pinned] + where = f" WHERE {' AND '.join(terms)}" if terms else "" + rows = self._conn.execute( + # nosemgrep: sqlalchemy-execute-raw-query -- tbl and the fields are validated identifiers; values are bound + f'SELECT rowid, doc FROM "{tbl}"{where} ORDER BY rowid', [value for _, value in pinned] + ) + return [(rowid, json.loads(doc)) for rowid, doc in rows] def _first_match(self, collection: str, filter: Filter) -> Optional[tuple[int, dict]]: - for rowid, doc in self._load(collection): + for rowid, doc in self._load(collection, filter): if evaluation.matches(doc, filter): return rowid, doc return None @contextmanager def _immediate(self): - self._conn.execute("BEGIN IMMEDIATE") + with self._lock_errors(): + self._conn.execute("BEGIN IMMEDIATE") try: yield self._conn.execute("COMMIT") @@ -312,3 +330,45 @@ def _json_default(self, o): def _dumps(self, doc: dict) -> str: return json.dumps(doc, default=self._json_default) + + def _pinned(self, tbl: str, filter: Filter) -> list[tuple[str, str]]: + """``filter``'s equality pins on the leading fields of the index it pins most (a unique one on a tie).""" + eq = {} + for field, preds in filter.conditions.items(): + # Only str and int compare in SQL as in Python: Eq(None) also matches an absent field, Eq(list) an element. + values = [p.value for p in preds if isinstance(p, Eq) and isinstance(p.value, (str, int))] + # A dotted path can step through an array, which json_extract doesn't. + if values and "." not in field: + eq[field] = json.dumps(values[0]) + best, pinned = (0, False), [] + for unique, fields in self._indexes_of(tbl): + prefix = list(itertools.takewhile(eq.__contains__, fields)) + if (len(prefix), unique) > best: + best, pinned = (len(prefix), unique), [(f, eq[f]) for f in prefix] + return pinned + + def _indexes_of(self, tbl: str) -> list[tuple[bool, list[str]]]: + # Read from the database, not kept from ensure_index: routing reads the @local store before it ensures any. + version = self._conn.execute("PRAGMA schema_version").fetchone()[0] + if version != self._schema_version: # another connection changed the schema, maybe adding an index + self._indexes, self._schema_version = {}, version + if tbl not in self._indexes: + rows = self._conn.execute( + "SELECT sql FROM sqlite_master WHERE type='index' AND tbl_name=? AND sql IS NOT NULL", (tbl,) + ) + self._indexes[tbl] = [(sql.startswith("CREATE UNIQUE"), _INDEX_FIELD_RE.findall(sql)) for (sql,) in rows] + return self._indexes[tbl] + + @contextmanager + def _lock_errors(self): + try: + yield + except sqlite3.OperationalError as e: + if e.sqlite_errorcode != sqlite3.SQLITE_BUSY: + raise + raise DatabaseLockedError( + f"{self._path} is locked: another process has held its write lock for over {_BUSY_TIMEOUT_SECONDS:g}s" + ) from e + + +_INDEX_FIELD_RE = re.compile(r"json_extract\(doc, '\$\.([A-Za-z0-9_.]+)'\)") diff --git a/tst/unit/cli/user_facing_errors_test.py b/tst/unit/cli/user_facing_errors_test.py index 2770b771..489272de 100644 --- a/tst/unit/cli/user_facing_errors_test.py +++ b/tst/unit/cli/user_facing_errors_test.py @@ -13,6 +13,7 @@ from agent_env.config.errors import ConfigError from agent_env.store.base import NotFoundError from agent_env.store.document_store import DuplicateKeyError +from agent_env.store.document_store.sqlite_document_store import DatabaseLockedError def _group_raising(exc): @@ -32,7 +33,8 @@ def boom(): (ValueError("id '@ns/x' starts with the reserved '@' prefix"), "id '@ns/x' starts with the reserved '@' prefix"), (ConfigError("Unknown AGENT_ENV_DOCUMENT_STORE='bogus'"), "Unknown AGENT_ENV_DOCUMENT_STORE='bogus'"), (NotFoundError("Artifact definitely-missing not found"), "Artifact definitely-missing not found"), -], ids=["reserved-id", "config", "not-found"]) + (DatabaseLockedError("/s/documents.db is locked: another process"), "/s/documents.db is locked: another process"), +], ids=["reserved-id", "config", "not-found", "locked"]) def test_a_rule_the_caller_broke_prints_one_line(exc, message): result = CliRunner().invoke(_group_raising(exc), ["boom"]) diff --git a/tst/unit/store/sqlite_document_store_test.py b/tst/unit/store/sqlite_document_store_test.py index b96ed189..ea4e6b4a 100644 --- a/tst/unit/store/sqlite_document_store_test.py +++ b/tst/unit/store/sqlite_document_store_test.py @@ -4,6 +4,7 @@ assertions that prove Mongo parity also give quick backend-neutral coverage. """ +import random import sqlite3 import subprocess import sys @@ -11,7 +12,23 @@ import pytest -from agent_env.store import Filter, LocalSqliteDocumentStore, UpdateSpec, compare_and_swap +from agent_env.store import ( + AbsentOrNull, + Eq, + Exists, + Filter, + Gte, + In, + LocalSqliteDocumentStore, + Lte, + LteOrAbsent, + Ne, + Sort, + SortKey, + UpdateSpec, + VersionedEntityStore, + compare_and_swap, +) from agent_env.store.document_store import sqlite_document_store from tst.store import conformance @@ -78,6 +95,118 @@ def test_processes_opening_a_new_database_at_once_all_write(tmp_path): assert len(LocalSqliteDocumentStore(str(path)).query("coll", Filter.of())) == 200 +_KEY = "docs_coll_id_version_unique" +_INDEXED_READS = { + "find_one(id)": (lambda s: s.find_one("coll", Filter.of(id="a")), f"{_KEY} (=?)"), + "find_one(id, version)": ( + lambda s: s.find_one("coll", Filter.of(id="a", version=1)), f"{_KEY} (=? AND =?)", + ), + "query(id)": (lambda s: s.query("coll", Filter.of(id="a"), sort=Sort.by("version")), f"{_KEY} (=?)"), + "count(id)": (lambda s: s.count("coll", Filter.of(id="a")), f"{_KEY} (=?)"), + "next_version": (lambda s: VersionedEntityStore(s, "coll", dict, dict).next_version("a"), f"{_KEY} (=?)"), + "update": ( + lambda s: s.update("coll", Filter.of(id="a", version=1), UpdateSpec(set={"x": 2})), f"{_KEY} (=? AND", + ), + "update_one_and_get": ( + lambda s: s.update_one_and_get("coll", Filter.of(id="a").where("x", Ne(0)), UpdateSpec(inc={"x": 1})), + f"{_KEY} (=?)", + ), + "replace": ( + lambda s: s.replace("coll", Filter.of(id="a", version=1), {"id": "a", "version": 1}), f"{_KEY} (=? AND", + ), + "delete": (lambda s: s.delete("coll", Filter.of(id="b", version=1)), f"{_KEY} (=? AND =?)"), + "the unique index on a tie": (lambda s: s.find_one("coll", Filter.of(id="a", type="file")), f"{_KEY} (=?)"), + "a secondary index": (lambda s: s.query("coll", Filter.of(type="file")), "docs_coll_type_index (=?)"), + "the longest pinned prefix": ( + lambda s: s.query("coll", Filter.of(task_id="t", created_at_utc="c")), + "docs_coll_task_id_created_at_utc_index (=? AND =?)", + ), +} + + +@pytest.mark.parametrize("read", _INDEXED_READS.values(), ids=_INDEXED_READS.keys()) +def test_equality_on_an_index_prefix_searches_the_index(read, store_coll): + store, coll = store_coll + _indexed(store, coll) + call, plan = read + plans = _plans(store, lambda: call(store)) + assert plans and all(f"USING INDEX {plan}" in p for p in plans), plans + + +@pytest.mark.parametrize("filter", [ + Filter.of(x=1), # no index + Filter.of(version=1), # not an index's leading field + Filter().where("id", Eq(None)), # also matches an absent id + Filter().where("id", Eq(["a"])), # also matches an array holding ["a"] + Filter().where("id", Eq(1.5)), + Filter().where("id", In(["a"])), + Filter().where("a2a_tasks.a2a_task_id", Eq("t")), # a dotted path steps through arrays +], ids=repr) +def test_other_filters_scan(filter, store_coll): + store, coll = store_coll + _indexed(store, coll) + assert _plans(store, lambda: store.query(coll, filter)) == ["SCAN docs_coll"] + + +def test_a_reader_searches_an_index_another_connection_creates(tmp_path): + path = str(tmp_path / "shared.db") + reader, writer = LocalSqliteDocumentStore(path), LocalSqliteDocumentStore(path) + writer.insert("coll", {"id": "a", "version": 1}) + read = lambda: reader.find_one("coll", Filter.of(id="a")) + assert _plans(reader, read) == ["SCAN docs_coll"] + writer.ensure_index("coll", ["id", "version"], unique=True) + assert "USING INDEX docs_coll_id_version_unique (=?)" in _plans(reader, read)[0] + + +def test_indexed_reads_and_writes_match_a_full_scan(tmp_path): + """Random documents and operations give the same results, in the same order, as the same store with no index.""" + rng = random.Random(3205) + indexed = LocalSqliteDocumentStore(str(tmp_path / "indexed.db")) + scanned = LocalSqliteDocumentStore(str(tmp_path / "scanned.db")) + for fields in (["id", "version"], ["type"], ["k", "id"], ["a.b"]): + indexed.ensure_index("coll", fields) + for _ in range(300): + doc = _random_doc(rng) + indexed.insert("coll", doc) + scanned.insert("coll", doc) + statements = [] + indexed._conn.set_trace_callback(statements.append) + for _ in range(1500): + op, call = _random_op(rng) + assert _outcome(call, indexed) == _outcome(call, scanned), op + assert indexed.query("coll", Filter()) == scanned.query("coll", Filter()) + searched = [sql for sql in statements if sql.startswith("SELECT rowid, doc") and " WHERE " in sql] + assert len(searched) > 500 # the pushdown was exercised, not skipped + + +def test_a_write_the_busy_timeout_gives_up_on_names_the_store(tmp_path, monkeypatch): + monkeypatch.setattr(sqlite_document_store, "_BUSY_TIMEOUT_SECONDS", 0.05) + path = tmp_path / "locked.db" + store = LocalSqliteDocumentStore(str(path)) + store.insert("coll", {"id": "a"}) + holder = sqlite3.connect(str(path), isolation_level=None) + holder.execute("BEGIN IMMEDIATE") + insert = lambda: store.insert("coll", {"id": "b"}) + update = lambda: store.update("coll", Filter.of(id="a"), UpdateSpec(set={"x": 1})) + for write in (insert, update): + with pytest.raises(sqlite_document_store.DatabaseLockedError, match=f"{path} is locked: another process"): + write() + holder.execute("ROLLBACK") + assert update() == 1 + + +def test_a_first_open_the_busy_timeout_gives_up_on_names_the_store(tmp_path, monkeypatch): + monkeypatch.setattr(sqlite_document_store, "_BUSY_TIMEOUT_SECONDS", 0.05) + path = tmp_path / "locked.db" + holder = sqlite3.connect(str(path), isolation_level=None) + holder.execute("CREATE TABLE t (x)") + holder.execute("BEGIN EXCLUSIVE") # a rollback-journal database, so the switch to WAL waits on this lock + with pytest.raises(sqlite_document_store.DatabaseLockedError, match=f"{path} is locked: another process"): + LocalSqliteDocumentStore(str(path)).count("coll", Filter()) + holder.execute("ROLLBACK") + assert LocalSqliteDocumentStore(str(path)).count("coll", Filter()) == 0 + + class _Connection: """Fails its first ``failures`` statements with ``code``.""" @@ -107,3 +236,87 @@ def test_the_wal_switch_raises_other_errors_and_a_lock_held_too_long(monkeypatch monkeypatch.setattr(sqlite_document_store, "_BUSY_TIMEOUT_SECONDS", 0.05) with pytest.raises(sqlite3.OperationalError): sqlite_document_store._enable_wal(_Connection(failures=10**6, code=sqlite3.SQLITE_BUSY)) + + +_FIELDS = ["id", "version", "type", "k", "x", "tags", "a.b", "missing"] +_VALUES = [ + "a", "b", "A", "", "1", "[1]", '{"x": 1}', "\u00e9", "\U0001f600", "\ud800", "a\x00b", + 0, 1, -1, 2, 10**20, 1.0, 2.5, True, False, None, {"x": 1}, [1], ["a"], [], +] +# Indexed fields hold scalars or objects, as agent-env writes them; an array under one is out of scope. +_KEY_VALUES = [v for v in _VALUES if not isinstance(v, list)] + + +def _random_doc(rng): + doc = {} + for field in ("id", "version", "type", "k", "x"): + if rng.random() < 0.85: + doc[field] = rng.choice(["a", "b", "1", 1, 2, None] if field == "id" else _KEY_VALUES) + if rng.random() < 0.5: + doc["tags"] = rng.sample(_VALUES, rng.randint(0, 3)) + shape = rng.random() + if shape < 0.3: + doc["a"] = {"b": rng.choice(_VALUES)} + elif shape < 0.6: + doc["a"] = [{"b": rng.choice(_VALUES)} for _ in range(rng.randint(0, 3))] + return doc + + +def _random_filter(rng, values=_VALUES): + filter = Filter() + for field in rng.sample(["id", "version", "type", "k"], rng.randint(0, 3)): + filter = filter.where(field, Eq(rng.choice(values))) + for _ in range(rng.choice([0, 0, 1, 2])): + value = rng.choice(values) + pred = rng.choice([Eq(value), Ne(value), In([value, rng.choice(values)]), Gte(value), Lte(value), + Exists(rng.random() < 0.5), LteOrAbsent(value), AbsentOrNull()]) + filter = filter.where(rng.choice(_FIELDS), pred) + return filter + + +def _random_op(rng): + filter, key_filter, doc = _random_filter(rng), _random_filter(rng, _KEY_VALUES), _random_doc(rng) + sort = rng.choice([None, Sort.by(rng.choice(_FIELDS), rng.random() < 0.5), + Sort((SortKey("type", False), SortKey("version")))]) + limit, offset = rng.choice([None, 1, 3]), rng.choice([None, 1]) + update = UpdateSpec(set={rng.choice(["x", "type", "version"]): rng.choice(_KEY_VALUES)}, inc={"n": 1}) + name, call = rng.choice([ + ("find_one", lambda s: s.find_one("coll", filter, sort)), + ("query", lambda s: s.query("coll", filter, sort, limit, offset)), + ("count", lambda s: s.count("coll", filter)), + ("update", lambda s: s.update("coll", filter, update)), + ("upsert", lambda s: s.update("coll", key_filter, update, upsert=True)), # synthesizes from the filter + ("update_one_and_get", lambda s: s.update_one_and_get("coll", filter, update, return_after=bool(limit))), + ("replace", lambda s: s.replace("coll", filter, doc)), + ("delete", lambda s: s.delete("coll", filter)), + ]) + return f"{name} {key_filter if name == 'upsert' else filter} {sort}", call + + +def _outcome(call, store): + try: + return call(store) + except TypeError as e: # sort_docs can't order unlike types, with or without an index + return repr(e) + + +def _indexed(store, coll): + store.ensure_index(coll, ["id", "version"], unique=True) + store.ensure_index(coll, ["type"]) + store.ensure_index(coll, ["task_id"]) + store.ensure_index(coll, ["task_id", "created_at_utc"]) + store.ensure_index(coll, ["a2a_tasks.a2a_task_id"]) + store.insert(coll, {"id": "a", "version": 1, "type": "file", "x": 1}) + store.insert(coll, {"id": "a", "version": 2, "type": "file", "x": 1}) + + +def _plans(store, call) -> list[str]: + """The query plan of each row read ``call`` makes.""" + statements = [] + store._conn.set_trace_callback(statements.append) + try: + call() + finally: + store._conn.set_trace_callback(None) + reads = [sql for sql in statements if sql.startswith("SELECT rowid, doc")] + return [" ".join(row[3] for row in store._conn.execute(f"EXPLAIN QUERY PLAN {sql}")) for sql in reads] From 0c739bc1b2e46c56f122c7d8d3a191c1188e19ea Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 1 Oct 2026 10:07:13 -1000 Subject: [PATCH 2/4] refactor(store): one definition of the indexed field expression ensure_index and the pushdown query now build json_extract(doc, '$.') through _field_expr, so the index and the query that must match it can't drift apart. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/document_store/sqlite_document_store.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/agent_env/store/document_store/sqlite_document_store.py b/src/agent_env/store/document_store/sqlite_document_store.py index 87a71c79..4709a2f9 100644 --- a/src/agent_env/store/document_store/sqlite_document_store.py +++ b/src/agent_env/store/document_store/sqlite_document_store.py @@ -233,7 +233,7 @@ def ensure_index( tbl = self._ensure_table(collection) raw_name = tbl + "_" + "_".join(fields) + ("_unique" if unique else "_index") name = self._safe_ident(re.sub(r"\W", "_", raw_name)) - cols = ", ".join(f"json_extract(doc, '$.{self._safe_path(f)}')" for f in fields) + cols = ", ".join(_field_expr(self._safe_path(f)) for f in fields) uniq = "UNIQUE " if unique else "" self._conn.execute( # nosemgrep: sqlalchemy-execute-raw-query -- name/tbl/cols are validated identifiers @@ -291,7 +291,7 @@ def _load(self, collection: str, filter: Filter) -> list[tuple[int, dict]]: return [] pinned = self._pinned(tbl, filter) # The index's own expression, so SQLite searches the index; the value is JSON so it decodes as the field does. - terms = [f"json_extract(doc, '$.{self._safe_path(field)}') = json_extract(?, '$')" for field, _ in pinned] + terms = [f"{_field_expr(self._safe_path(field))} = json_extract(?, '$')" for field, _ in pinned] where = f" WHERE {' AND '.join(terms)}" if terms else "" rows = self._conn.execute( # nosemgrep: sqlalchemy-execute-raw-query -- tbl and the fields are validated identifiers; values are bound @@ -371,4 +371,10 @@ def _lock_errors(self): ) from e +def _field_expr(field: str) -> str: + """How an index stores a field, so a query that spells it the same way searches that index.""" + return f"json_extract(doc, '$.{field}')" + + +# The fields of an index, read back from the SQL `_field_expr` wrote; the two must agree. _INDEX_FIELD_RE = re.compile(r"json_extract\(doc, '\$\.([A-Za-z0-9_.]+)'\)") From 20aeb7dfdeebf9f06607a61384cfa3c736e4aa4c Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 1 Oct 2026 10:14:31 -1000 Subject: [PATCH 3/4] refactor(store): the index cache drops on a schema change only CREATE INDEX bumps PRAGMA schema_version on the connection that runs it, so the version check already covers an index this store adds; ensure_index no longer clears the cache by hand. The reader test now runs for the same store and another connection. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/store/document_store/sqlite_document_store.py | 3 +-- tst/unit/store/sqlite_document_store_test.py | 6 ++++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/agent_env/store/document_store/sqlite_document_store.py b/src/agent_env/store/document_store/sqlite_document_store.py index 4709a2f9..0f7904ce 100644 --- a/src/agent_env/store/document_store/sqlite_document_store.py +++ b/src/agent_env/store/document_store/sqlite_document_store.py @@ -239,7 +239,6 @@ def ensure_index( # nosemgrep: sqlalchemy-execute-raw-query -- name/tbl/cols are validated identifiers f'CREATE {uniq}INDEX IF NOT EXISTS "{name}" ON "{tbl}" ({cols})' ) - self._indexes.pop(tbl, None) # --- sqlite plumbing --- @@ -350,7 +349,7 @@ def _pinned(self, tbl: str, filter: Filter) -> list[tuple[str, str]]: def _indexes_of(self, tbl: str) -> list[tuple[bool, list[str]]]: # Read from the database, not kept from ensure_index: routing reads the @local store before it ensures any. version = self._conn.execute("PRAGMA schema_version").fetchone()[0] - if version != self._schema_version: # another connection changed the schema, maybe adding an index + if version != self._schema_version: # the schema changed, maybe by an index added here or by another connection self._indexes, self._schema_version = {}, version if tbl not in self._indexes: rows = self._conn.execute( diff --git a/tst/unit/store/sqlite_document_store_test.py b/tst/unit/store/sqlite_document_store_test.py index ea4e6b4a..eb4415cc 100644 --- a/tst/unit/store/sqlite_document_store_test.py +++ b/tst/unit/store/sqlite_document_store_test.py @@ -148,9 +148,11 @@ def test_other_filters_scan(filter, store_coll): assert _plans(store, lambda: store.query(coll, filter)) == ["SCAN docs_coll"] -def test_a_reader_searches_an_index_another_connection_creates(tmp_path): +@pytest.mark.parametrize("same_store", [False, True], ids=["another connection", "the same store"]) +def test_a_reader_searches_an_index_created_after_its_first_read(tmp_path, same_store): path = str(tmp_path / "shared.db") - reader, writer = LocalSqliteDocumentStore(path), LocalSqliteDocumentStore(path) + reader = LocalSqliteDocumentStore(path) + writer = reader if same_store else LocalSqliteDocumentStore(path) writer.insert("coll", {"id": "a", "version": 1}) read = lambda: reader.find_one("coll", Filter.of(id="a")) assert _plans(reader, read) == ["SCAN docs_coll"] From 40e8b0d2ac69405c489aef28adb9e4f1471a3c19 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 1 Oct 2026 10:48:40 -1000 Subject: [PATCH 4/4] revert(cli): keep the SQLite backend's lock error out of the CLI The CLI listed DatabaseLockedError among its one-line user errors, which made it import the SQLite backend. A lock timeout now prints like any other unexpected error, as a traceback that ends in the same message naming the database file. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/cli/__init__.py | 11 +++++------ tst/unit/cli/user_facing_errors_test.py | 4 +--- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/src/agent_env/cli/__init__.py b/src/agent_env/cli/__init__.py index e946f6ee..e39688d4 100644 --- a/src/agent_env/cli/__init__.py +++ b/src/agent_env/cli/__init__.py @@ -4,7 +4,6 @@ from agent_env.plugins._cli import load_cli_plugins, load_cli_root_options from agent_env.store.base import NotFoundError -from agent_env.store.document_store.sqlite_document_store import DatabaseLockedError from agent_env.store.routing import namespace_routing from agent_env.utils.docker_build import DockerBuildError @@ -19,11 +18,11 @@ from .up import up -# Errors that mean "not allowed", "not there", "your build failed" or "locked" rather than "agent-env is -# broken". Spelled out because there is no domain base class to catch: ConfigError subclasses ValueError -# (so does pydantic's ValidationError), while NotFoundError, DockerBuildError and DatabaseLockedError don't. -# A TypeError or a retry-exhausted DuplicateKeyError is a defect, and a defect keeps its traceback. -_USER_FACING_ERRORS = (ValueError, NotFoundError, DockerBuildError, DatabaseLockedError) +# Errors that mean "not allowed", "not there" or "your build failed" rather than "agent-env is broken". +# Spelled out because there is no domain base class to catch: ConfigError subclasses ValueError (so does +# pydantic's ValidationError), while NotFoundError and DockerBuildError don't. A TypeError or a +# retry-exhausted DuplicateKeyError is a defect, and a defect keeps its traceback. +_USER_FACING_ERRORS = (ValueError, NotFoundError, DockerBuildError) class _UserErrorsAreNotCrashes(click.Group): diff --git a/tst/unit/cli/user_facing_errors_test.py b/tst/unit/cli/user_facing_errors_test.py index 489272de..2770b771 100644 --- a/tst/unit/cli/user_facing_errors_test.py +++ b/tst/unit/cli/user_facing_errors_test.py @@ -13,7 +13,6 @@ from agent_env.config.errors import ConfigError from agent_env.store.base import NotFoundError from agent_env.store.document_store import DuplicateKeyError -from agent_env.store.document_store.sqlite_document_store import DatabaseLockedError def _group_raising(exc): @@ -33,8 +32,7 @@ def boom(): (ValueError("id '@ns/x' starts with the reserved '@' prefix"), "id '@ns/x' starts with the reserved '@' prefix"), (ConfigError("Unknown AGENT_ENV_DOCUMENT_STORE='bogus'"), "Unknown AGENT_ENV_DOCUMENT_STORE='bogus'"), (NotFoundError("Artifact definitely-missing not found"), "Artifact definitely-missing not found"), - (DatabaseLockedError("/s/documents.db is locked: another process"), "/s/documents.db is locked: another process"), -], ids=["reserved-id", "config", "not-found", "locked"]) +], ids=["reserved-id", "config", "not-found"]) def test_a_rule_the_caller_broke_prints_one_line(exc, message): result = CliRunner().invoke(_group_raising(exc), ["boom"])