diff --git a/src/paperless_ai/ai_classifier.py b/src/paperless_ai/ai_classifier.py index 41bf8e903..13a145709 100644 --- a/src/paperless_ai/ai_classifier.py +++ b/src/paperless_ai/ai_classifier.py @@ -96,19 +96,26 @@ def get_context_for_document( user: User | None = None, max_docs: int = 5, ) -> str: - visible_documents = ( - get_objects_for_user_owner_aware( - user, - "view_document", - Document, - ) - if user - else None - ) + # None means "no restriction" to query_similar_documents. A superuser + # (like no user at all) can see every document, so skip materializing + # every visible pk into a Python list and passing it through as a SQL + # IN filter: for a large library that is a wasted quadratic scan in the + # vector store at best, and past ~32,763 documents a hard + # sqlite3.OperationalError (SQLite's bound-parameter limit) at worst. + # get_objects_for_user_owner_aware() would return every Document for a + # superuser anyway (guardian's own with_superuser shortcut), so this + # changes nothing about which documents are considered -- only how we + # get there. visible_document_ids = ( - list(visible_documents.values_list("pk", flat=True)) - if visible_documents is not None - else None + None + if user is None or user.is_superuser + else list( + get_objects_for_user_owner_aware( + user, + "view_document", + Document, + ).values_list("pk", flat=True), + ) ) similar_docs = query_similar_documents( document=doc, diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 941625869..dd8f8c6fd 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -397,7 +397,7 @@ def _document_id_filters(doc_ids): MetadataFilter( key="document_id", operator=FilterOperator.IN, - value=sorted(doc_ids), + value=list(doc_ids), ), ], ) diff --git a/src/paperless_ai/tests/test_ai_classifier.py b/src/paperless_ai/tests/test_ai_classifier.py index 470c7fe07..0cc6d8231 100644 --- a/src/paperless_ai/tests/test_ai_classifier.py +++ b/src/paperless_ai/tests/test_ai_classifier.py @@ -3,6 +3,8 @@ from unittest.mock import MagicMock from unittest.mock import patch import pytest +import pytest_mock +from django.contrib.auth.models import User from django.test import override_settings from documents.models import Document @@ -278,3 +280,104 @@ def test_get_context_for_document_no_similar_docs(mock_document): with patch("paperless_ai.ai_classifier.query_similar_documents", return_value=[]): result = get_context_for_document(mock_document) assert result == "" + + +class TestGetContextForDocumentVisibility: + """get_context_for_document must not materialize every visible document + id for a user who can already see the whole library: a superuser (like + no user at all) gets document_ids=None (no restriction) straight + through to query_similar_documents(), instead of a full-library IN + filter that is wasteful at best and, past ~32,763 documents, a hard + sqlite3.OperationalError at worst (SQLite's bound-parameter limit). + """ + + def test_skips_permission_lookup_for_superuser( + self, + mock_document: MagicMock, + mock_similar_documents: list[MagicMock], + mocker: pytest_mock.MockerFixture, + ) -> None: + """ + GIVEN: + - A superuser + WHEN: + - get_context_for_document() is called + THEN: + - get_objects_for_user_owner_aware() is never called, and + query_similar_documents() is called with document_ids=None + """ + mock_query = mocker.patch( + "paperless_ai.ai_classifier.query_similar_documents", + return_value=mock_similar_documents, + ) + mock_get_objects = mocker.patch( + "paperless_ai.ai_classifier.get_objects_for_user_owner_aware", + ) + user = mocker.MagicMock(spec=User) + user.is_superuser = True + + get_context_for_document(mock_document, user, max_docs=2) + + mock_get_objects.assert_not_called() + assert mock_query.call_args.kwargs["document_ids"] is None + + def test_skips_permission_lookup_when_no_user( + self, + mock_document: MagicMock, + mock_similar_documents: list[MagicMock], + mocker: pytest_mock.MockerFixture, + ) -> None: + """ + GIVEN: + - No user (user=None) + WHEN: + - get_context_for_document() is called + THEN: + - get_objects_for_user_owner_aware() is never called, and + query_similar_documents() is called with document_ids=None + """ + mock_query = mocker.patch( + "paperless_ai.ai_classifier.query_similar_documents", + return_value=mock_similar_documents, + ) + mock_get_objects = mocker.patch( + "paperless_ai.ai_classifier.get_objects_for_user_owner_aware", + ) + + get_context_for_document(mock_document, None, max_docs=2) + + mock_get_objects.assert_not_called() + assert mock_query.call_args.kwargs["document_ids"] is None + + def test_restricts_to_visible_documents_for_non_superuser( + self, + mock_document: MagicMock, + mock_similar_documents: list[MagicMock], + mocker: pytest_mock.MockerFixture, + ) -> None: + """ + GIVEN: + - A non-superuser with a specific set of visible documents + WHEN: + - get_context_for_document() is called + THEN: + - query_similar_documents() is called with exactly that user's + visible document ids, unchanged from before this optimization + """ + mock_query = mocker.patch( + "paperless_ai.ai_classifier.query_similar_documents", + return_value=mock_similar_documents, + ) + mock_queryset = mocker.MagicMock() + mock_queryset.values_list.return_value = [1, 2, 3] + mock_get_objects = mocker.patch( + "paperless_ai.ai_classifier.get_objects_for_user_owner_aware", + return_value=mock_queryset, + ) + user = mocker.MagicMock(spec=User) + user.is_superuser = False + + get_context_for_document(mock_document, user, max_docs=2) + + mock_get_objects.assert_called_once_with(user, "view_document", Document) + assert mock_query.call_args.kwargs["document_ids"] == [1, 2, 3] diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 3d2dc13f5..3ed1b7e42 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.vector_store import _MAX_IN_VALUES from paperless_ai.vector_store import DB_FILENAME from paperless_ai.vector_store import DEFAULT_TABLE_NAME from paperless_ai.vector_store import SCHEMA_VERSION @@ -296,8 +297,47 @@ class TestBuildWhere: assert where == "1 = 0" assert params == [] - def test_query_with_untranslatable_filter_returns_no_rows(self, store) -> None: - store.add([make_node("a1", 1), make_node("b1", 2)]) + def test_fails_closed_when_in_filter_exceeds_max_values( + self, + caplog: pytest.LogCaptureFixture, + ) -> None: + """ + GIVEN: + - An 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) + 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 + """ + oversized = _in_filter([str(i) for i in range(_MAX_IN_VALUES + 1)]) + + with caplog.at_level("WARNING"): + where, params = _build_where(oversized) + + assert where == "(1 = 0)" + assert params == [] + assert "document_id" in caplog.text + + def test_query_with_untranslatable_filter_returns_no_rows( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: + """ + GIVEN: + - A store with one chunk each for documents 1 and 2 + WHEN: + - A query and a get_nodes() call are issued with a filter whose + only clause is an untranslatable nested MetadataFilters + THEN: + - Neither call raises, and both fail closed by returning no rows + (rather than emitting invalid or unfiltered SQL) + """ + store.add([make_node("a1", "1"), make_node("b1", "2")]) nested = MetadataFilters( filters=[ MetadataFilter( diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index 975b6c57b..b7de0936d 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -75,6 +75,17 @@ 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. +_MAX_IN_VALUES = 32700 + + def _pack(embedding: Sequence[float]) -> bytes: return struct.pack(f"{len(embedding)}f", *embedding) @@ -119,6 +130,21 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: if not values: # pragma: no cover 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, + ) + clauses.append("1 = 0") + continue placeholders = ",".join("?" for _ in values) clauses.append(f"{f.key} IN ({placeholders})") params.extend(values)