perf(store): the local SQLite document store answers indexed lookups from its indexes

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,15 +103,16 @@ 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
 
     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_.]+)'\)")