diff --git a/src/paperless_ai/tables.py b/src/paperless_ai/tables.py index c8d6d838c..f13b662f2 100644 --- a/src/paperless_ai/tables.py +++ b/src/paperless_ai/tables.py @@ -154,6 +154,35 @@ class DocumentMetaTable: } +class PermittedIdsTable: + """Per-connection scratch space for an oversized IN-filter id list. + + A literal ``IN (?,?,...)`` list binds one SQL parameter per id, capped by + SQLite's own SQLITE_MAX_VARIABLE_NUMBER (see _MAX_IN_VALUES in + vector_store.py). Loading the ids into a TEMP TABLE and filtering via a + subquery instead has no such limit. TEMP tables live in a + connection-private namespace -- never visible to another connection, + even under this identical name -- so this is safe under the vector + store's one-connection-per-request model without any extra locking or + per-call naming scheme. + """ + + TABLE_NAME = "permitted_document_ids" + + @staticmethod + def load(conn: sqlite3.Connection, ids: Iterable[int]) -> None: + """Replace this connection's scratch table with ``ids``.""" + conn.execute(f"DROP TABLE IF EXISTS temp.{PermittedIdsTable.TABLE_NAME}") + conn.execute( + f"CREATE TEMP TABLE {PermittedIdsTable.TABLE_NAME} " + "(id INTEGER PRIMARY KEY)", + ) + conn.executemany( + f"INSERT INTO {PermittedIdsTable.TABLE_NAME} (id) VALUES (?)", + ((i,) for i in ids), + ) + + class IndexMetaTable: """Typed accessors over index_meta's key/value rows -- replaces PaperlessSqliteVecVectorStore._meta_get_on/_meta_set_on, which returned diff --git a/src/paperless_ai/tests/test_tables.py b/src/paperless_ai/tests/test_tables.py index 1feda711c..547f63f60 100644 --- a/src/paperless_ai/tests/test_tables.py +++ b/src/paperless_ai/tests/test_tables.py @@ -9,6 +9,7 @@ from paperless_ai.tables import DocumentChunksTable from paperless_ai.tables import DocumentMetaRow from paperless_ai.tables import DocumentMetaTable from paperless_ai.tables import IndexMetaTable +from paperless_ai.tables import PermittedIdsTable @pytest.fixture @@ -338,3 +339,85 @@ class TestIndexMetaTable: IndexMetaTable.increment_total_inserts(conn, 100) IndexMetaTable.reset_total_inserts(conn, 7) assert IndexMetaTable.get_total_inserts(conn) == 7 + + +class TestPermittedIdsTable: + def _loaded_ids(self, conn: sqlite3.Connection) -> list[int]: + return [ + row["id"] + for row in conn.execute( + f"SELECT id FROM {PermittedIdsTable.TABLE_NAME} ORDER BY id", + ) + ] + + def test_load_then_read_back_all_ids(self, conn: sqlite3.Connection) -> None: + """ + GIVEN: + - A bare sqlite3 connection + WHEN: + - load() is called with a set of ids + THEN: + - Every id is present in the TEMP TABLE, and only those ids + """ + PermittedIdsTable.load(conn, [3, 1, 2]) + assert self._loaded_ids(conn) == [1, 2, 3] + + def test_load_replaces_previous_contents(self, conn: sqlite3.Connection) -> None: + """ + GIVEN: + - A connection whose PermittedIdsTable already holds one id set + WHEN: + - load() is called again with a different id set + THEN: + - Only the new ids are present -- a connection reused across + multiple queries in one request never leaks a stale filter + """ + PermittedIdsTable.load(conn, [1, 2, 3]) + PermittedIdsTable.load(conn, [4, 5]) + assert self._loaded_ids(conn) == [4, 5] + + def test_load_is_connection_private(self) -> None: + """ + GIVEN: + - Two separate connections + WHEN: + - Each loads PermittedIdsTable with a different id set, under + the identical TABLE_NAME + THEN: + - Each connection sees only its own ids -- TEMP TABLE is + connection-private, so concurrent requests never collide or + cross-contaminate despite sharing the same table name (the + vector store opens one connection per request; see + PaperlessSqliteVecVectorStore) + """ + conn_a = sqlite3.connect(":memory:") + conn_a.row_factory = sqlite3.Row + conn_b = sqlite3.connect(":memory:") + conn_b.row_factory = sqlite3.Row + try: + PermittedIdsTable.load(conn_a, [1, 2, 3]) + PermittedIdsTable.load(conn_b, [4, 5, 6]) + assert self._loaded_ids(conn_a) == [1, 2, 3] + assert self._loaded_ids(conn_b) == [4, 5, 6] + finally: + conn_a.close() + conn_b.close() + + def test_load_handles_more_ids_than_a_bound_parameter_list_could( + self, + conn: sqlite3.Connection, + ) -> None: + """ + GIVEN: + - An id count over SQLite's own bound-parameter limit + (SQLITE_MAX_VARIABLE_NUMBER, 32766 by default) -- more than a + literal IN(?,?,...) list could ever bind in one statement + WHEN: + - load() is called with that many ids + THEN: + - Every id is loaded without error, since executemany() binds + one row at a time rather than one statement with N parameters + """ + ids = list(range(40_000)) + PermittedIdsTable.load(conn, ids) + assert self._loaded_ids(conn) == ids diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 3ed1b7e42..95649a759 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -17,6 +17,7 @@ from paperless_ai.migrations import Migration from paperless_ai.migrations import m0001_v1_to_v2 from paperless_ai.tables import DocumentChunksTable from paperless_ai.tables import DocumentMetaTable +from paperless_ai.tables import PermittedIdsTable from paperless_ai.vector_store import _MAX_IN_VALUES from paperless_ai.vector_store import DB_FILENAME from paperless_ai.vector_store import DEFAULT_TABLE_NAME @@ -267,8 +268,23 @@ class TestCrud: class TestBuildWhere: - def test_ne_filter_translates_to_not_equal_clause(self) -> None: - where, params = _build_where(_ne_filter(1)) + @pytest.fixture + def conn(self) -> Generator[sqlite3.Connection, None, None]: + """A bare connection, sufficient for _build_where(): it only ever + touches the connection via PermittedIdsTable, which needs no vec0 + extension loaded. + """ + connection = sqlite3.connect(":memory:") + try: + yield connection + finally: + connection.close() + + def test_ne_filter_translates_to_not_equal_clause( + self, + conn: sqlite3.Connection, + ) -> None: + where, params = _build_where(conn, _ne_filter(1)) assert where == "(document_id != ?)" assert params == [1] @@ -280,7 +296,10 @@ class TestBuildWhere: "b1", ] - def test_fails_closed_when_no_filter_is_translatable(self) -> None: + def test_fails_closed_when_no_filter_is_translatable( + self, + conn: sqlite3.Connection, + ) -> None: # A nested MetadataFilters is not a MetadataFilter, so it is skipped. # With no translatable clauses, the function must fail closed rather # than emit "()" (invalid SQL) and never widen document access. @@ -293,35 +312,74 @@ class TestBuildWhere: ), ], ) - where, params = _build_where(MetadataFilters(filters=[nested])) + where, params = _build_where(conn, MetadataFilters(filters=[nested])) assert where == "1 = 0" assert params == [] - def test_fails_closed_when_in_filter_exceeds_max_values( + def test_in_filter_over_max_values_uses_permitted_ids_table( self, - caplog: pytest.LogCaptureFixture, + conn: sqlite3.Connection, ) -> None: """ GIVEN: - An IN filter with more values than _MAX_IN_VALUES (SQLite's - own bound-parameter limit is 32766; this guard sits below + own bound-parameter limit is 32766; this threshold sits below that with headroom for the query's other bound parameters) WHEN: - _build_where() translates it to SQL THEN: - - It fails closed ("1 = 0", no params) instead of building an - IN clause SQLite would reject, and logs a warning -- this - filter scopes document access, so refusing to build it must - never widen the scope to "everything" by accident + - It builds a subquery against PermittedIdsTable's TEMP TABLE, + loaded with every id, instead of a literal IN(...) list that + SQLite would reject past its own limit -- the filter still + scopes document access to exactly the requested ids, never + widening the scope to "everything" """ - oversized = _in_filter([str(i) for i in range(_MAX_IN_VALUES + 1)]) + ids = list(range(_MAX_IN_VALUES + 1)) + oversized = _in_filter([str(i) for i in ids]) - with caplog.at_level("WARNING"): - where, params = _build_where(oversized) + where, params = _build_where(conn, oversized) - assert where == "(1 = 0)" + assert where == ( + f"(document_id IN (SELECT id FROM {PermittedIdsTable.TABLE_NAME}))" + ) assert params == [] - assert "document_id" in caplog.text + loaded = [ + row[0] + for row in conn.execute( + f"SELECT id FROM {PermittedIdsTable.TABLE_NAME} ORDER BY id", + ) + ] + assert loaded == ids + + def test_query_and_get_nodes_scope_correctly_when_in_filter_exceeds_max_values( + self, + store: PaperlessSqliteVecVectorStore, + mocker: MockerFixture, + ) -> None: + """ + GIVEN: + - _MAX_IN_VALUES lowered so a small IN filter exceeds it + WHEN: + - query() and get_nodes() are called with that filter + THEN: + - Both still correctly scope results to the permitted ids -- the + PermittedIdsTable temp-table path behaves identically to the + literal IN(...) path it replaces above the threshold + """ + mocker.patch("paperless_ai.vector_store._MAX_IN_VALUES", 1) + store.add( + [ + make_node("a1", 1, seed=0.0), + make_node("b1", 2, seed=1.0), + make_node("c1", 3, seed=2.0), + ], + ) + + result = _query(store, [0.0] * DIM, top_k=10, filters=_in_filter([2, 3])) + nodes = store.get_nodes(filters=_in_filter([2, 3])) + + assert sorted(result.ids) == ["b1", "c1"] + assert sorted(n.node_id for n in nodes) == ["b1", "c1"] def test_query_with_untranslatable_filter_returns_no_rows( self, diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index b7de0936d..e12fa3634 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -30,6 +30,7 @@ from paperless_ai.tables import DocumentChunksTable from paperless_ai.tables import DocumentMetaRow from paperless_ai.tables import DocumentMetaTable from paperless_ai.tables import IndexMetaTable +from paperless_ai.tables import PermittedIdsTable logger = logging.getLogger("paperless_ai.vector_store") @@ -75,14 +76,12 @@ class _Row(NamedTuple): embedding: bytes -# _build_where(): the largest IN value list translated into bound SQL -# parameters. SQLite's own hard limit (SQLITE_MAX_VARIABLE_NUMBER) is 32766 -# by default; this leaves headroom below that for the query's other bound -# parameters (the embedding blob, k, and any NE clause) and for the limit -# itself to move. An IN filter this large should not happen in practice -- -# callers are expected to pass None (no filter) rather than every id when -# the filter would not actually narrow anything -- so this is a guard -# against a future regression, not a normal code path. +# _build_where(): the largest IN value list translated into a literal +# IN (?,?,...) clause. SQLite's own hard limit (SQLITE_MAX_VARIABLE_NUMBER) +# is 32766 by default; this leaves headroom below that for the query's other +# bound parameters (the embedding blob, k, and any NE clause) and for the +# limit itself to move. Above this threshold _build_where() switches to +# PermittedIdsTable instead of failing closed -- see its docstring. _MAX_IN_VALUES = 32700 @@ -106,13 +105,20 @@ def _vec0_params(rows: list[_Row]) -> list[tuple[str, int, str, bytes]]: return [(r.chunk_id, r.document_id, r.node_content, r.embedding) for r in rows] -def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: +def _build_where( + conn: sqlite3.Connection, + filters: MetadataFilters | None, +) -> tuple[str, list[int]]: """Translate the EQ / IN / NE filters we use into a parameterized SQL clause on vec0 metadata columns. Returns ("", []) when there is nothing to filter. document_id is vec0's only filterable column and is INTEGER; every value is coerced via int() here so callers (which today still pass strings in places, e.g. indexing.py's MetadataFilter construction) don't have to be individually correct -- vec0 doesn't coerce types itself. + + ``conn`` is only used for an IN filter over _MAX_IN_VALUES: it loads the + ids into PermittedIdsTable's TEMP TABLE on that connection rather than + binding them as SQL parameters. """ if filters is None or not filters.filters: return "", [] @@ -131,19 +137,14 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: clauses.append("1 = 0") continue if len(values) > _MAX_IN_VALUES: - # Fail closed (see the empty-clauses case below) rather than - # let SQLite raise "too many SQL variables" past its own - # limit: this filter scopes document access, so an IN list - # too large to safely bind must match no rows, never widen - # the scope to "everything" by accident. - logger.warning( - "Refusing to build an IN filter on %r with %d values " - "(over the %d-value safety limit); returning no rows.", - f.key, - len(values), - _MAX_IN_VALUES, + # A literal IN(...) list this large would exceed SQLite's own + # bound-parameter limit. Load the ids into a TEMP TABLE on + # this connection instead and filter via subquery, which has + # no such limit -- see PermittedIdsTable. + PermittedIdsTable.load(conn, values) + clauses.append( + f"{f.key} IN (SELECT id FROM {PermittedIdsTable.TABLE_NAME})", ) - clauses.append("1 = 0") continue placeholders = ",".join("?" for _ in values) clauses.append(f"{f.key} IN ({placeholders})") @@ -482,7 +483,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): ) if not self.table_exists(): return [] - where, params = _build_where(filters) + where, params = _build_where(self._conn, filters) sql = "SELECT node_content, embedding FROM " + DEFAULT_TABLE_NAME if where: sql += " WHERE " + where @@ -498,7 +499,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): if query.query_embedding is None: # pragma: no cover return VectorStoreQueryResult(nodes=[], similarities=[], ids=[]) top_k = query.similarity_top_k if query.similarity_top_k is not None else 10 - where, params = _build_where(query.filters) + where, params = _build_where(self._conn, query.filters) sql = ( "SELECT id, node_content, embedding, distance FROM " + DEFAULT_TABLE_NAME