Address PR #1376 follow-up: lazy subqueries, top-level import, test gap

Follow-up to issue #1377:

- Move ``StructuralAnnotationSet`` to the existing top-level
  ``opencontractserver.annotations.models`` import. ``Annotation`` was
  already imported at module scope from the same module, so the inline
  imports inside ``_build_base_queryset`` ("avoid circular imports") had
  no rationale.

- Convert ``active_doc_ids`` and ``visible_corpus_doc_ids`` from
  Python-materialised lists to fully lazy querysets so the SQL planner
  receives ``IN (SELECT ...)`` subqueries rather than
  ``IN (val, val, ...)`` literals — important for corpora with tens of
  thousands of documents. The deletion-aware empty-corpus short-circuit
  now uses ``.exists()`` instead of list truthiness.

- Add ``.distinct()`` to ``visible_corpus_set_ids`` and
  ``active_struct_set_ids`` for consistency with
  ``visible_corpus_doc_ids`` and to avoid duplicate IDs when a structural
  set is shared across multiple documents in the corpus.

- Drop stale ``(line 211)`` line-number anchor from the corpus-only
  branch comment.

- Add ``test_deletion_aware_path_excludes_deleted_document_structural``
  to close the previously uncovered branch where
  ``check_corpus_deletion=True`` AND a document's only ``DocumentPath``
  has ``is_deleted=True`` — the structural annotation must be excluded.

Author: claude

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,
@@ -289,20 +289,22 @@ 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
                 )
-            )()
+                .values("document_id")
+                .distinct()
+            )
+            has_active_docs = await sync_to_async(active_doc_ids_qs.exists)()
 
-            if active_doc_ids:
+            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 +313,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 +401,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,40 @@ 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_document_scoped_search_still_returns_structural(self) -> None:
         """Document-scoped retrieval is the existing well-tested path.