Merge pull request #1379 from Open-Source-Legal/claude/fix-issue-1377-gTvaU

Optimize corpus deletion checks with lazy subqueries

Author: JSv4

opencontractserver/llms/vector_stores/core_vector_stores.py

@@ -9,7 +9,7 @@
 from django.contrib.auth import get_user_model
 from django.db.models import Q, QuerySet
 
-from opencontractserver.annotations.models import Annotation
+from opencontractserver.annotations.models import Annotation, StructuralAnnotationSet
 from opencontractserver.constants.search import (
     FTS_CONFIG,
     HYBRID_SEARCH_OVERSAMPLE_FACTOR,
@@ -208,7 +208,7 @@ async def _build_base_queryset(self) -> QuerySet[Annotation]:
         # enumeration attacks.
         # -------------------------------------------------------------------------
         from opencontractserver.corpuses.models import Corpus
-        from opencontractserver.documents.models import Document
+        from opencontractserver.documents.models import Document, DocumentPath
 
         user = None
         if self.user_id:
@@ -289,20 +289,29 @@ async def _build_base_queryset(self) -> QuerySet[Annotation]:
 
         # Check for deleted documents in corpus
         if self.check_corpus_deletion and self.corpus_id and not self.document_id:
-            # Note: sync_to_async already imported at module level
-            from opencontractserver.annotations.models import StructuralAnnotationSet
-            from opencontractserver.documents.models import DocumentPath
-
-            # Get documents with active (non-deleted) paths in corpus
-            active_doc_ids = await sync_to_async(
-                lambda: list(
-                    DocumentPath.objects.filter(
-                        corpus_id=self.corpus_id, is_current=True, is_deleted=False
-                    ).values_list("document_id", flat=True)
+            # Lazy subquery — never round-trips through Python, so the
+            # generated SQL stays a single statement with a real subquery
+            # rather than a giant ``IN (val, val, ...)`` literal even for
+            # corpora with tens of thousands of documents.
+            active_doc_ids_qs = (
+                DocumentPath.objects.filter(
+                    corpus_id=self.corpus_id, is_current=True, is_deleted=False
                 )
-            )()
-
-            if active_doc_ids:
+                .values("document_id")
+                .distinct()
+            )
+            # Trade-off: this ``EXISTS`` adds one extra round-trip on the
+            # happy path, but lets us short-circuit the entire vector search
+            # for empty/all-deleted corpora (returning ``Annotation.none()``
+            # spares a downstream HNSW probe and keeps the existing
+            # operational warning). For corpora with at least one active
+            # document the cost is a single boolean SELECT and is dwarfed
+            # by the main query. Removing the check would also remove the
+            # debug log of the active-doc count that the materialised list
+            # used to provide.
+            has_active_docs = await sync_to_async(active_doc_ids_qs.exists)()
+
+            if has_active_docs:
                 # Two annotation shapes pass this filter:
                 #   1. Direct: Annotation.document_id is in the active set.
                 #   2. Structural: Annotation.document_id is NULL but the
@@ -311,13 +320,16 @@ async def _build_base_queryset(self) -> QuerySet[Annotation]:
                 #      structural annotations need an explicit OR clause —
                 #      otherwise every parser-produced structural row is
                 #      silently dropped on this corpus-wide path.
-                active_struct_set_ids = StructuralAnnotationSet.objects.filter(
-                    documents__in=active_doc_ids
-                ).values("id")
-                active_filters &= Q(document_id__in=active_doc_ids) | Q(
+                active_struct_set_ids = (
+                    StructuralAnnotationSet.objects.filter(
+                        documents__in=active_doc_ids_qs
+                    )
+                    .values("id")
+                    .distinct()
+                )
+                active_filters &= Q(document_id__in=active_doc_ids_qs) | Q(
                     structural=True, structural_set_id__in=active_struct_set_ids
                 )
-                _logger.debug(f"Found {len(active_doc_ids)} active documents in corpus")
             else:
                 _logger.warning(f"No active documents found in corpus {self.corpus_id}")
                 return Annotation.objects.none()
@@ -396,21 +408,23 @@ async def _build_base_queryset(self) -> QuerySet[Annotation]:
             #    ``Annotation.corpus_id``. Per-document visibility for these
             #    rows is enforced by the visibility filter further below
             #    plus the upfront IDOR check on ``corpus_id``.
-            # Document is imported earlier in this method (line 211); reusing
-            # the local binding avoids an F811 redefinition warning.
-            from opencontractserver.annotations.models import StructuralAnnotationSet
-
-            visible_corpus_doc_ids = await sync_to_async(
-                lambda: list(
-                    Document.objects.visible_to_user(user)
-                    .filter(path_records__corpus_id=self.corpus_id)
-                    .values_list("id", flat=True)
-                    .distinct()
+            # Both subqueries below stay lazy so the SQL planner sees a
+            # nested ``IN (SELECT ...)`` rather than a Python-materialised
+            # ``IN (val, val, ...)`` literal — important for corpora with
+            # tens of thousands of documents.
+            visible_corpus_doc_ids_qs = (
+                Document.objects.visible_to_user(user)
+                .filter(path_records__corpus_id=self.corpus_id)
+                .values("id")
+                .distinct()
+            )
+            visible_corpus_set_ids = (
+                StructuralAnnotationSet.objects.filter(
+                    documents__in=visible_corpus_doc_ids_qs
                 )
-            )()
-            visible_corpus_set_ids = StructuralAnnotationSet.objects.filter(
-                documents__in=visible_corpus_doc_ids
-            ).values("id")
+                .values("id")
+                .distinct()
+            )
             active_filters &= Q(
                 structural=True, structural_set_id__in=visible_corpus_set_ids
             ) | Q(structural=False, corpus_id=self.corpus_id)

opencontractserver/tests/test_corpus_isolation_vector_store.py

@@ -260,6 +260,120 @@ def test_corpus_wide_search_excludes_orphan_structural_set(self) -> None:
     # Document-scoped retrieval — must continue to work post-fix
     # ------------------------------------------------------------------ #
 
+    def test_deletion_aware_path_excludes_deleted_document_structural(self) -> None:
+        """Structural rows whose only ``DocumentPath`` is deleted are dropped.
+
+        Covers the deletion side of ``check_corpus_deletion=True``: the
+        ``DocumentPath`` row for ``doc_a`` is flipped to ``is_deleted=True``,
+        so ``active_doc_ids`` in
+        ``CoreAnnotationVectorStore._build_base_queryset`` becomes empty and
+        the corpus-only branch must short-circuit to no results — the
+        document's structural annotation must NOT come back, even though the
+        annotation row itself is otherwise unchanged.
+        """
+        DocumentPath.objects.filter(document=self.doc_a, corpus=self.corpus_a).update(
+            is_deleted=True
+        )
+
+        store = CoreAnnotationVectorStore(
+            user_id=self.user.id,
+            corpus_id=self.corpus_a.id,
+            document_id=None,
+            check_corpus_deletion=True,
+        )
+        query = VectorSearchQuery(
+            query_embedding=_constant_vector(384, 0.5), similarity_top_k=10
+        )
+        results = store.search(query)
+        returned_ids = {r.annotation.id for r in results}
+
+        self.assertNotIn(
+            self.struct_a.id,
+            returned_ids,
+            "Structural annotation whose only DocumentPath is deleted "
+            "leaked through the deletion-aware corpus-only path",
+        )
+
+    def test_deletion_aware_path_partial_deletion_keeps_active_drops_deleted(
+        self,
+    ) -> None:
+        """Partial deletion: with two docs in one corpus, only the active stays.
+
+        The previous test stresses the early-exit branch (every doc deleted
+        → ``has_active_docs`` is False → ``Annotation.none()``). This test
+        stresses the actual subquery filter:
+
+            active_filters &= Q(document_id__in=active_doc_ids_qs) | Q(
+                structural=True, structural_set_id__in=active_struct_set_ids
+            )
+
+        With one active document and one deleted document in the same
+        corpus, ``active_doc_ids_qs`` is non-empty but partial. The
+        deleted document's structural annotation must be excluded while
+        the active document's must come through. This is the scenario
+        most likely to regress silently if the OR clause is rebuilt or
+        the subquery is rewritten incorrectly.
+        """
+        # Add a second document to corpus_a with its own structural set,
+        # then mark its DocumentPath deleted while doc_a's stays active.
+        deleted_set = StructuralAnnotationSet.objects.create(
+            content_hash=hashlib.sha256(b"corpus-a-deleted-doc").hexdigest(),
+            creator=self.user,
+            parser_name="TestParser",
+            parser_version="1.0",
+        )
+        deleted_doc = Document.objects.create(
+            title="Doc A2 (deleted)",
+            creator=self.user,
+            is_public=True,
+            structural_annotation_set=deleted_set,
+        )
+        DocumentPath.objects.create(
+            document=deleted_doc,
+            corpus=self.corpus_a,
+            path="/doc_a2.pdf",
+            version_number=1,
+            is_current=True,
+            is_deleted=True,
+            creator=self.user,
+        )
+        deleted_struct = Annotation.objects.create(
+            structural_set=deleted_set,
+            annotation_label=self.label,
+            creator=self.user,
+            raw_text="Deleted-doc structural paragraph",
+            structural=True,
+            page=1,
+        )
+        deleted_struct.add_embedding(
+            get_default_embedder_path(), _constant_vector(384, 0.15)
+        )
+
+        store = CoreAnnotationVectorStore(
+            user_id=self.user.id,
+            corpus_id=self.corpus_a.id,
+            document_id=None,
+            check_corpus_deletion=True,
+        )
+        query = VectorSearchQuery(
+            query_embedding=_constant_vector(384, 0.5), similarity_top_k=10
+        )
+        results = store.search(query)
+        returned_ids = {r.annotation.id for r in results}
+
+        self.assertIn(
+            self.struct_a.id,
+            returned_ids,
+            "Active document's structural annotation was excluded by the "
+            "partial-deletion subquery filter",
+        )
+        self.assertNotIn(
+            deleted_struct.id,
+            returned_ids,
+            "Deleted document's structural annotation leaked through the "
+            "partial-deletion subquery filter",
+        )
+
     def test_document_scoped_search_still_returns_structural(self) -> None:
         """Document-scoped retrieval is the existing well-tested path.