mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-09-02 16:07:15 +00:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7c0d3c2f86 | ||
|
|
ec7745e71f | ||
|
|
98e87d91ad |
@@ -314,7 +314,7 @@ def _consume_file(
|
||||
consumption_dir: Path,
|
||||
*,
|
||||
subdirs_as_tags: bool,
|
||||
) -> bool:
|
||||
) -> None:
|
||||
"""
|
||||
Queue a file for consumption.
|
||||
|
||||
@@ -322,20 +322,15 @@ def _consume_file(
|
||||
filepath: Path to the file to consume.
|
||||
consumption_dir: Base consumption directory.
|
||||
subdirs_as_tags: Whether to create tags from subdirectory names.
|
||||
|
||||
Returns:
|
||||
True if the file was successfully handed to Celery, False otherwise.
|
||||
Callers must not record the file as queued on failure, or the rescan
|
||||
will never retry it.
|
||||
"""
|
||||
# Verify file still exists and is accessible
|
||||
try:
|
||||
if not filepath.is_file():
|
||||
logger.debug(f"Not consuming {filepath}: not a file or doesn't exist")
|
||||
return False
|
||||
return
|
||||
except OSError as e:
|
||||
logger.warning(f"Not consuming {filepath}: {e}")
|
||||
return False
|
||||
return
|
||||
|
||||
# Get tags from path if configured
|
||||
tag_ids: list[int] | None = None
|
||||
@@ -360,9 +355,6 @@ def _consume_file(
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(f"Error while queuing document {filepath}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
@@ -500,12 +492,12 @@ class Command(BaseCommand):
|
||||
if not consumer_filter(Change.added, str(filepath)):
|
||||
continue
|
||||
|
||||
if _consume_file(
|
||||
_consume_file(
|
||||
filepath=filepath,
|
||||
consumption_dir=directory,
|
||||
subdirs_as_tags=subdirs_as_tags,
|
||||
):
|
||||
queued.add(filepath.resolve())
|
||||
)
|
||||
queued.add(filepath.resolve())
|
||||
|
||||
return queued
|
||||
|
||||
@@ -659,16 +651,14 @@ class Command(BaseCommand):
|
||||
|
||||
# Check for stable files
|
||||
for stable_path in tracker.get_stable_files():
|
||||
# Only remember files that were actually queued, so the
|
||||
# rescan does not re-queue them while the consume task
|
||||
# has yet to remove them from disk, but does retry a
|
||||
# failed publish instead of stranding it
|
||||
if _consume_file(
|
||||
_consume_file(
|
||||
filepath=stable_path,
|
||||
consumption_dir=directory,
|
||||
subdirs_as_tags=subdirs_as_tags,
|
||||
):
|
||||
queued.add(stable_path)
|
||||
)
|
||||
# Remember it so the rescan does not re-queue it while
|
||||
# the consume task has yet to remove it from disk
|
||||
queued.add(stable_path)
|
||||
|
||||
# Exit watch loop to reconfigure timeout
|
||||
break
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -445,13 +445,12 @@ class TestConsumeFile:
|
||||
target = consumption_dir / "document.pdf"
|
||||
shutil.copy(sample_pdf, target)
|
||||
|
||||
result = _consume_file(
|
||||
_consume_file(
|
||||
filepath=target,
|
||||
consumption_dir=consumption_dir,
|
||||
subdirs_as_tags=False,
|
||||
)
|
||||
|
||||
assert result is True
|
||||
mock_consume_file_delay.apply_async.assert_called_once()
|
||||
call_args = mock_consume_file_delay.apply_async.call_args
|
||||
consumable_doc = call_args.kwargs["kwargs"]["input_doc"]
|
||||
@@ -465,12 +464,11 @@ class TestConsumeFile:
|
||||
mock_consume_file_delay: MagicMock,
|
||||
) -> None:
|
||||
"""Test _consume_file handles nonexistent files gracefully."""
|
||||
result = _consume_file(
|
||||
_consume_file(
|
||||
filepath=consumption_dir / "nonexistent.pdf",
|
||||
consumption_dir=consumption_dir,
|
||||
subdirs_as_tags=False,
|
||||
)
|
||||
assert result is False
|
||||
mock_consume_file_delay.apply_async.assert_not_called()
|
||||
|
||||
def test_consume_directory(
|
||||
@@ -482,12 +480,11 @@ class TestConsumeFile:
|
||||
subdir = consumption_dir / "subdir"
|
||||
subdir.mkdir()
|
||||
|
||||
result = _consume_file(
|
||||
_consume_file(
|
||||
filepath=subdir,
|
||||
consumption_dir=consumption_dir,
|
||||
subdirs_as_tags=False,
|
||||
)
|
||||
assert result is False
|
||||
mock_consume_file_delay.apply_async.assert_not_called()
|
||||
|
||||
def test_consume_with_permission_error(
|
||||
@@ -502,33 +499,13 @@ class TestConsumeFile:
|
||||
shutil.copy(sample_pdf, target)
|
||||
|
||||
mocker.patch.object(Path, "is_file", side_effect=PermissionError("denied"))
|
||||
result = _consume_file(
|
||||
_consume_file(
|
||||
filepath=target,
|
||||
consumption_dir=consumption_dir,
|
||||
subdirs_as_tags=False,
|
||||
)
|
||||
assert result is False
|
||||
mock_consume_file_delay.apply_async.assert_not_called()
|
||||
|
||||
def test_consume_with_apply_async_failure(
|
||||
self,
|
||||
consumption_dir: Path,
|
||||
sample_pdf: Path,
|
||||
mock_consume_file_delay: MagicMock,
|
||||
) -> None:
|
||||
"""Test _consume_file reports failure when apply_async raises."""
|
||||
target = consumption_dir / "document.pdf"
|
||||
shutil.copy(sample_pdf, target)
|
||||
|
||||
mock_consume_file_delay.apply_async.side_effect = Exception("broker down")
|
||||
|
||||
result = _consume_file(
|
||||
filepath=target,
|
||||
consumption_dir=consumption_dir,
|
||||
subdirs_as_tags=False,
|
||||
)
|
||||
assert result is False
|
||||
|
||||
def test_consume_with_tags_error(
|
||||
self,
|
||||
consumption_dir: Path,
|
||||
@@ -545,12 +522,11 @@ class TestConsumeFile:
|
||||
side_effect=DatabaseError("Something happened"),
|
||||
)
|
||||
|
||||
result = _consume_file(
|
||||
_consume_file(
|
||||
filepath=target,
|
||||
consumption_dir=consumption_dir,
|
||||
subdirs_as_tags=True,
|
||||
)
|
||||
assert result is True
|
||||
mock_consume_file_delay.apply_async.assert_called_once()
|
||||
call_args = mock_consume_file_delay.apply_async.call_args
|
||||
overrides = call_args.kwargs["kwargs"]["overrides"]
|
||||
@@ -1273,52 +1249,6 @@ class TestProcessExistingFilesQueued:
|
||||
assert target.resolve() in queued
|
||||
|
||||
|
||||
@pytest.mark.management
|
||||
@pytest.mark.django_db
|
||||
class TestCommandRetryAfterQueueFailure:
|
||||
"""
|
||||
Regression test for GH #13923.
|
||||
|
||||
A file whose ``apply_async`` publish fails (e.g. broker briefly down)
|
||||
must not be marked as queued, so the periodic rescan retries it once
|
||||
the broker recovers, instead of stranding it until the consumer
|
||||
process is restarted.
|
||||
"""
|
||||
|
||||
def test_watch_loop_retries_failed_publish_on_rescan(
|
||||
self,
|
||||
consumption_dir: Path,
|
||||
sample_pdf: Path,
|
||||
mock_consume_file_delay: MagicMock,
|
||||
start_consumer: Callable[..., ConsumerThread],
|
||||
) -> None:
|
||||
"""A publish failure from the watch loop is retried by the rescan."""
|
||||
apply_async = mock_consume_file_delay.apply_async
|
||||
|
||||
def fail_first_call(*args: object, **kwargs: object) -> None:
|
||||
if apply_async.call_count == 1:
|
||||
raise Exception("broker down")
|
||||
|
||||
apply_async.side_effect = fail_first_call
|
||||
|
||||
thread = start_consumer(stability_delay=0.1, rescan_interval=0.3)
|
||||
|
||||
target = consumption_dir / "document.pdf"
|
||||
shutil.copy(sample_pdf, target)
|
||||
|
||||
deadline = monotonic() + 5.0
|
||||
while apply_async.call_count < 2 and monotonic() < deadline:
|
||||
sleep(0.1)
|
||||
|
||||
if thread.exception:
|
||||
raise thread.exception
|
||||
|
||||
assert apply_async.call_count >= 2, (
|
||||
"Expected the failed publish to be retried by the rescan, "
|
||||
f"but apply_async was only called {apply_async.call_count} time(s)"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.management
|
||||
@pytest.mark.django_db
|
||||
class TestCommandRescanRecovery:
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -705,12 +705,6 @@ CELERY_BROKER_TRANSPORT_OPTIONS = {
|
||||
CELERY_TASK_TRACK_STARTED = True
|
||||
CELERY_TASK_TIME_LIMIT: Final[int] = get_int_from_env("PAPERLESS_WORKER_TIMEOUT", 1800)
|
||||
|
||||
# https://docs.celeryq.dev/en/stable/userguide/configuration.html#std-setting-task_allow_error_cb_on_chord_header
|
||||
# Without this, a failing chord header never triggers the errback, so a mail
|
||||
# whose attachments all fail is never recorded and is re-fetched forever.
|
||||
# The errback runs once per failed header task, so it must be idempotent.
|
||||
CELERY_TASK_ALLOW_ERROR_CB_ON_CHORD_HEADER = True
|
||||
|
||||
CELERY_CACHE_BACKEND = "default"
|
||||
|
||||
# https://docs.celeryq.dev/en/stable/userguide/configuration.html#task-serializer
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -334,24 +334,18 @@ def error_callback(
|
||||
"""
|
||||
A shared task that is called whenever something goes wrong during
|
||||
consumption of a file. See queue_consumption_tasks.
|
||||
|
||||
With CELERY_TASK_ALLOW_ERROR_CB_ON_CHORD_HEADER enabled this runs once per
|
||||
failed header task, not once per chord, so it must be idempotent.
|
||||
"""
|
||||
rule = MailRule.objects.get(pk=rule_id)
|
||||
received = make_aware(message_date) if is_naive(message_date) else message_date
|
||||
|
||||
ProcessedMail.objects.get_or_create(
|
||||
ProcessedMail.objects.create(
|
||||
rule=rule,
|
||||
folder=rule.folder,
|
||||
uid=message_uid,
|
||||
uid_validity=uid_validity,
|
||||
defaults={
|
||||
"subject": message_subject,
|
||||
"received": received,
|
||||
"status": "FAILED",
|
||||
"error": traceback.format_exc(),
|
||||
},
|
||||
subject=message_subject,
|
||||
received=make_aware(message_date) if is_naive(message_date) else message_date,
|
||||
status="FAILED",
|
||||
error=traceback.format_exc(),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -36,7 +36,6 @@ from paperless_mail.mail import MailAccountHandler
|
||||
from paperless_mail.mail import MailError
|
||||
from paperless_mail.mail import TagMailAction
|
||||
from paperless_mail.mail import apply_mail_action
|
||||
from paperless_mail.mail import error_callback
|
||||
from paperless_mail.mail import get_mailbox
|
||||
from paperless_mail.models import MailAccount
|
||||
from paperless_mail.models import MailRule
|
||||
@@ -2046,44 +2045,6 @@ class TestPostConsumeAction(TestCase):
|
||||
self.assertIn("Test Exception", processed_mail.error)
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
class TestErrorCallback:
|
||||
def test_error_callback_is_idempotent_for_same_mail(self) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- A mail rule and a mail that failed to be consumed
|
||||
WHEN:
|
||||
- error_callback is invoked more than once for the same mail, as
|
||||
happens when task_allow_error_cb_on_chord_header fires the
|
||||
errback once per failed header task in a chord
|
||||
THEN:
|
||||
- Only one ProcessedMail row is created for that mail
|
||||
"""
|
||||
rule = MailRuleFactory()
|
||||
message_uid = "12345"
|
||||
|
||||
for _ in range(2):
|
||||
error_callback(
|
||||
None,
|
||||
Exception("Test Exception"),
|
||||
None,
|
||||
rule_id=rule.pk,
|
||||
message_uid=message_uid,
|
||||
message_subject="Test Subject",
|
||||
message_date=timezone.make_aware(
|
||||
timezone.datetime(2023, 1, 1, 12, 0, 0),
|
||||
),
|
||||
)
|
||||
|
||||
processed_mails = ProcessedMail.objects.filter(
|
||||
rule=rule,
|
||||
uid=message_uid,
|
||||
folder=rule.folder,
|
||||
)
|
||||
assert processed_mails.count() == 1
|
||||
assert processed_mails.get().status == "FAILED"
|
||||
|
||||
|
||||
class TestManagementCommand(TestCase):
|
||||
@mock.patch(
|
||||
"paperless_mail.management.commands.mail_fetcher.tasks.process_mail_accounts",
|
||||
|
||||
Reference in New Issue
Block a user