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 1c4dbf2f9..3b55cc1fb 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -18,6 +18,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 @@ -280,8 +281,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] @@ -293,8 +309,11 @@ class TestBuildWhere: "b1", ] - def test_nin_filter_translates_to_not_in_clause(self) -> None: - where, params = _build_where(_nin_filter([1, 2])) + def test_nin_filter_translates_to_not_in_clause( + self, + conn: sqlite3.Connection, + ) -> None: + where, params = _build_where(conn, _nin_filter([1, 2])) assert where == "(document_id NOT IN (?,?))" assert params == [1, 2] @@ -304,7 +323,10 @@ class TestBuildWhere: _query(store, [0.0] * DIM, top_k=5, filters=_nin_filter([1, 2])).ids, ) == ["c1"] - def test_empty_in_filter_excludes_everything(self) -> None: + def test_empty_in_filter_excludes_everything( + self, + conn: sqlite3.Connection, + ) -> None: """ GIVEN: - An IN filter with an empty value list @@ -314,11 +336,14 @@ class TestBuildWhere: - It excludes everything (the opposite of an empty NOT IN filter) -- an empty inclusion list must never widen results """ - where, params = _build_where(_in_filter([])) + where, params = _build_where(conn, _in_filter([])) assert where == "(1 = 0)" assert params == [] - def test_empty_nin_filter_excludes_nothing(self) -> None: + def test_empty_nin_filter_excludes_nothing( + self, + conn: sqlite3.Connection, + ) -> None: """ GIVEN: - A NOT IN filter with an empty value list -- e.g. an @@ -330,11 +355,14 @@ class TestBuildWhere: excludes everything) -- an empty exclusion list must never narrow results """ - where, params = _build_where(_nin_filter([])) + where, params = _build_where(conn, _nin_filter([])) assert where == "(1 = 1)" assert params == [] - 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. @@ -347,42 +375,88 @@ class TestBuildWhere: ), ], ) - where, params = _build_where(MetadataFilters(filters=[nested])) + where, params = _build_where(conn, MetadataFilters(filters=[nested])) assert where == "1 = 0" assert params == [] @pytest.mark.parametrize( - "build_filter", - [_in_filter, _nin_filter], + ("build_filter", "sql_op"), + [(_in_filter, "IN"), (_nin_filter, "NOT IN")], ids=["in", "nin"], ) - def test_fails_closed_when_filter_exceeds_max_values( + def test_filter_over_max_values_uses_permitted_ids_table( self, + conn: sqlite3.Connection, build_filter: Callable[[list[str]], MetadataFilters], - caplog: pytest.LogCaptureFixture, + sql_op: str, ) -> None: """ GIVEN: - An IN or NOT IN filter with more values than _MAX_IN_VALUES - (SQLite's own bound-parameter limit is 32766; this guard sits - below that with headroom for the query's other bound parameters) + (SQLite's 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 a - 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. Failing open on - a NOT IN would surface exactly the excluded rows + - It builds a subquery against PermittedIdsTable's TEMP TABLE, + loaded with every id, instead of a literal list SQLite would + reject past its own limit -- true for NOT IN too (e.g. an + install with an enormous trash), not just IN """ - oversized = build_filter([str(i) for i in range(_MAX_IN_VALUES + 1)]) + ids = list(range(_MAX_IN_VALUES + 1)) + oversized = build_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 {sql_op} (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 + + @pytest.mark.parametrize( + ("build_filter", "expected_ids"), + [(_in_filter, ["b1", "c1"]), (_nin_filter, ["a1"])], + ids=["in", "nin"], + ) + def test_query_and_get_nodes_scope_correctly_when_filter_exceeds_max_values( + self, + store: PaperlessSqliteVecVectorStore, + mocker: MockerFixture, + build_filter: Callable[[list[int]], MetadataFilters], + expected_ids: list[str], + ) -> None: + """ + GIVEN: + - _MAX_IN_VALUES lowered so a small IN/NOT IN filter exceeds it + WHEN: + - query() and get_nodes() are called with that filter + THEN: + - Both still correctly scope results -- the PermittedIdsTable + temp-table path behaves identically to the literal + IN(...)/NOT 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=build_filter([2, 3])) + nodes = store.get_nodes(filters=build_filter([2, 3])) + + assert sorted(result.ids) == expected_ids + assert sorted(n.node_id for n in nodes) == expected_ids 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 c7675261d..b45ce33c5 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,7 +105,10 @@ 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 / NIN / 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 @@ -114,6 +116,10 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: 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/NOT 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 "", [] @@ -136,20 +142,16 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: clauses.append("1 = 0" if is_in else "1 = 1") continue if len(values) > _MAX_IN_VALUES: - # Refuse rather than risk SQLite's own bound-parameter limit - # ("too many SQL variables"): a list this large must match no - # rows, never widen the scope to "everything" -- true for - # NOT IN too, where failing open would surface every - # excluded row. - logger.warning( - "Refusing to build a %s filter on %r with %d values " - "(over the %d-value safety limit); returning no rows.", - sql_op, - f.key, - len(values), - _MAX_IN_VALUES, + # A literal 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. Applies to NOT IN + # too (e.g. an install with an enormous trash), not just IN. + PermittedIdsTable.load(conn, values) + clauses.append( + f"{f.key} {sql_op} " + f"(SELECT id FROM {PermittedIdsTable.TABLE_NAME})", ) - clauses.append("1 = 0") continue placeholders = ",".join("?" for _ in values) clauses.append(f"{f.key} {sql_op} ({placeholders})") @@ -488,7 +490,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 @@ -504,7 +506,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