Compare commits

...
Author SHA1 Message Date
stumpylog 7c0d3c2f86 If the filter will exceed SQLite IN limits, load to a temporary table and use a subquery to filter instead 2026-09-02 08:33:34 -07:00
stumpylog ec7745e71f Minor improvements from a Claude review 2026-09-02 08:14:41 -07:00
stumpylog 98e87d91ad Fix: skip vector store document id filter for unrestricted chat users
ChatStreamingView built an IN filter from every permitted document id
for the "chat over all documents" case, which exceeds the vector
store's SQLite bound-parameter safety limit on installs with more
than ~32700 documents, silently returning no context. For a user who
can see every document (an active superuser), that filter never
narrows anything, so skip it and let the retriever search the whole
index instead.
2026-09-02 08:14:23 -07:00
8 changed files with 368 additions and 61 deletions
+74
View File
@@ -1,11 +1,18 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from unittest import mock
import pytest
from django.contrib.auth.models import Permission
from django.contrib.auth.models import User
from rest_framework import status
from rest_framework.test import APIClient
from rest_framework.test import APITestCase
if TYPE_CHECKING:
from pytest_mock import MockerFixture
class TestChatStreamingViewInputValidation(APITestCase):
def setUp(self) -> None:
@@ -42,3 +49,70 @@ class TestChatStreamingViewInputValidation(APITestCase):
format="json",
)
assert resp.status_code == status.HTTP_400_BAD_REQUEST
@pytest.mark.django_db
class TestChatStreamingViewUnrestrictedFlag:
"""The document id filter may only be skipped (``unrestricted=True``) for
a caller who can see every document, i.e. an active superuser.
"""
@pytest.fixture
def mocked_stream_chat(self, mocker: MockerFixture) -> mock.MagicMock:
"""AI enabled, with stream_chat_with_documents patched so the view
never touches the real vector store; returns the patched callable so
tests can inspect how it was called.
"""
mocker.patch("documents.views.AIConfig").return_value.ai_enabled = True
return mocker.patch(
"documents.views.stream_chat_with_documents",
return_value=iter(()),
)
@pytest.fixture
def viewer_client(self, user_client: APIClient, regular_user: User) -> APIClient:
"""The conftest regular-user client, additionally granted
view_document -- able to see every document without being a
superuser.
"""
regular_user.user_permissions.add(
*Permission.objects.filter(codename="view_document"),
)
return user_client
@pytest.mark.parametrize(
("client_fixture", "expected_unrestricted"),
[
pytest.param("admin_client", True, id="superuser_is_unrestricted"),
pytest.param("viewer_client", False, id="regular_user_is_restricted"),
],
)
def test_unrestricted_only_for_superuser(
self,
request: pytest.FixtureRequest,
mocked_stream_chat: mock.MagicMock,
client_fixture: str,
*,
expected_unrestricted: bool,
) -> None:
"""
GIVEN:
- A superuser, or a regular user holding view_document
WHEN:
- They post a chat question with no document_id
THEN:
- stream_chat_with_documents is called with unrestricted=True for
the superuser and unrestricted=False for the regular user, even
though that user can view every document
"""
client: APIClient = request.getfixturevalue(client_fixture)
client.post(
"/api/documents/chat/",
data={"q": "What's in these documents?"},
format="json",
)
assert (
mocked_stream_chat.call_args.kwargs["unrestricted"] is expected_unrestricted
)
+4
View File
@@ -180,6 +180,7 @@ from documents.permissions import has_system_status_permission
from documents.permissions import permitted_document_ids
from documents.permissions import permitted_object_ids
from documents.permissions import set_permissions_for_object
from documents.permissions import user_is_unrestricted
from documents.plugins.date_parsing import get_date_parser
from documents.schema import generate_object_with_permissions_schema
from documents.search import SearchHit
@@ -2329,10 +2330,12 @@ class ChatStreamingView(GenericAPIView[Any]):
return HttpResponseForbidden("Insufficient permissions")
documents = Document.objects.filter(pk=document.pk)
unrestricted = False
else:
documents = Document.objects.filter(
id__in=permitted_document_ids(request.user),
)
unrestricted = user_is_unrestricted(request.user)
output_language = get_llm_output_language(
ai_config=ai_config,
@@ -2343,6 +2346,7 @@ class ChatStreamingView(GenericAPIView[Any]):
stream_chat_with_documents(
query_str=question,
documents=documents,
unrestricted=unrestricted,
output_language=output_language,
),
content_type="text/event-stream",
+14 -3
View File
@@ -95,12 +95,15 @@ def _format_chat_metadata_trailer(references: list[dict[str, int | str]]) -> str
def stream_chat_with_documents(
query_str: str,
documents: QuerySet[Document],
*,
unrestricted: bool = False,
output_language: str | None = None,
):
try:
yield from _stream_chat_with_documents(
query_str,
documents,
unrestricted=unrestricted,
output_language=output_language,
)
except Exception as e:
@@ -111,6 +114,8 @@ def stream_chat_with_documents(
def _stream_chat_with_documents(
query_str: str,
documents: QuerySet[Document],
*,
unrestricted: bool = False,
output_language: str | None = None,
):
if not documents.exists():
@@ -123,9 +128,15 @@ def _stream_chat_with_documents(
from llama_index.core.retrievers import VectorIndexRetriever
config = AIConfig()
filters = _document_id_filters(
str(pk) for pk in documents.values_list("pk", flat=True)
)
if unrestricted:
# The caller can see every document, so an id filter would never narrow
# the search, only risk exceeding the vector store's bound parameter
# limit (_MAX_IN_VALUES in vector_store.py) on large installs.
filters = None
else:
filters = _document_id_filters(
str(pk) for pk in documents.values_list("pk", flat=True)
)
# Hold the shared read lock for the whole operation: the query engine
# retrieves from the vector store again during synthesis, so the connection
+29
View File
@@ -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
+66 -19
View File
@@ -1,4 +1,8 @@
from __future__ import annotations
import json
from typing import TYPE_CHECKING
from typing import Any
from unittest.mock import MagicMock
from unittest.mock import patch
@@ -18,6 +22,11 @@ from paperless_ai.chat import _build_chat_prompt
from paperless_ai.chat import _build_refine_prompt
from paperless_ai.chat import stream_chat_with_documents
if TYPE_CHECKING:
from pathlib import Path
import pytest_mock
@pytest.fixture(autouse=True)
def patch_embed_model():
@@ -312,6 +321,30 @@ def test_stream_chat_unexpected_failure_returns_generic_error(caplog) -> None:
@pytest.mark.django_db
class TestStreamChatRetrieval:
@pytest.fixture
def captured_filters(self, mocker: pytest_mock.MockerFixture) -> list[Any]:
"""Stub out the AI client and the retriever, capturing the ``filters``
kwarg of every VectorIndexRetriever construction.
VectorIndexRetriever is imported inside _stream_chat_with_documents,
so it is patched at the llama_index source for the lazy import to
pick it up.
"""
captured: list[Any] = []
retriever = mocker.MagicMock()
retriever.retrieve.return_value = []
def capture_retriever(*args, **kwargs) -> pytest_mock.MockType:
captured.append(kwargs.get("filters"))
return retriever
mocker.patch("paperless_ai.chat.AIClient")
mocker.patch(
"llama_index.core.retrievers.VectorIndexRetriever",
side_effect=capture_retriever,
)
return captured
def test_no_nodes_yields_no_content_message(
self,
temp_llm_index_dir,
@@ -329,9 +362,9 @@ class TestStreamChatRetrieval:
def test_chat_filter_contains_only_requested_document_ids(
self,
temp_llm_index_dir,
mock_embed_model,
mocker,
temp_llm_index_dir: Path,
mock_embed_model: pytest_mock.MockType,
captured_filters: list[Any],
) -> None:
"""The MetadataFilter passed to the retriever must be scoped to the
requested documents only content from other indexed documents must
@@ -342,22 +375,6 @@ class TestStreamChatRetrieval:
indexing.llm_index_add_or_update_document(included)
indexing.llm_index_add_or_update_document(excluded)
# VectorIndexRetriever is imported inside _stream_chat_with_documents;
# patch it at the llama_index source so the lazy import picks it up.
captured_filters = []
mock_retriever = mocker.MagicMock()
mock_retriever.retrieve.return_value = []
def capture_retriever(*args, **kwargs):
captured_filters.append(kwargs.get("filters"))
return mock_retriever
mocker.patch("paperless_ai.chat.AIClient")
mocker.patch(
"llama_index.core.retrievers.VectorIndexRetriever",
side_effect=capture_retriever,
)
list(
chat.stream_chat_with_documents(
"question?",
@@ -372,6 +389,36 @@ class TestStreamChatRetrieval:
assert str(included.pk) in filter_values
assert str(excluded.pk) not in filter_values
def test_unrestricted_chat_skips_document_id_filter(
self,
temp_llm_index_dir: Path,
mock_embed_model: pytest_mock.MockType,
captured_filters: list[Any],
) -> None:
"""
GIVEN:
- A document indexed in the vector store
WHEN:
- stream_chat_with_documents is called with unrestricted=True
THEN:
- The retriever receives no document id filter (filters=None), so
the whole index is searched instead of an IN-list that risks the
vector store's safety limit on large installs
"""
document = DocumentFactory.create(content="indexed document content")
indexing.llm_index_add_or_update_document(document)
list(
chat.stream_chat_with_documents(
"question?",
Document.objects.filter(pk=document.pk),
unrestricted=True,
),
)
assert captured_filters, "VectorIndexRetriever was never constructed"
assert captured_filters[0] is None
@pytest.mark.django_db
def test_get_document_references_only_queries_referenced_documents(
self,
+83
View File
@@ -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
+74 -16
View File
@@ -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,
+24 -23
View File
@@ -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