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..0f7904ce 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,11 +229,11 @@ 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)) - 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 @@ -272,25 +283,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"{_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 + 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 +329,51 @@ 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: # 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( + "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 + + +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_.]+)'\)") diff --git a/tst/unit/store/sqlite_document_store_test.py b/tst/unit/store/sqlite_document_store_test.py index b96ed189..eb4415cc 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,120 @@ 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"] + + +@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 = 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"] + 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 +238,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]