Coverage for src/local_deep_research/research_library/services/collection_rebuild.py: 95%
225 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""Build a replacement collection index without changing its published state."""
3from __future__ import annotations
5import hashlib
6import json
7import uuid
8from collections.abc import Callable
9from contextlib import AbstractContextManager
10from datetime import UTC, datetime
11from pathlib import Path
12from typing import TYPE_CHECKING, Any
14import numpy as np
15from loguru import logger
16from sqlalchemy.orm import Session
18from ...config.paths import get_cache_directory
19from ...database.models.library import (
20 Collection,
21 Document,
22 DocumentChunk,
23 DocumentCollection,
24 EmbeddingProvider,
25 RagDocumentStatus,
26 RAGIndex,
27 RAGIndexStatus,
28)
29from ...database.session_context import safe_rollback
30from ...vector_stores.config import only_local_file_providers_registered
32if TYPE_CHECKING:
33 from .library_rag_service import LibraryRAGService, _PreparedDocument
36_CONFIG_FIELDS = (
37 "embedding_model",
38 "embedding_model_type",
39 "embedding_dimension",
40 "chunk_size",
41 "chunk_overlap",
42 "splitter_type",
43 "text_separators",
44 "distance_metric",
45 "normalize_vectors",
46 "index_type",
47)
50def document_fingerprint(document: Document) -> str:
51 """Identify the source text and metadata used to prepare a document."""
52 payload = {
53 key: getattr(document, key)
54 for key in (
55 "text_content",
56 "title",
57 "filename",
58 "original_url",
59 "authors",
60 "published_date",
61 "doi",
62 "arxiv_id",
63 "pmid",
64 "pmcid",
65 "extraction_method",
66 "word_count",
67 )
68 }
69 return hashlib.sha256(
70 json.dumps(payload, sort_keys=True, default=str).encode("utf-8")
71 ).hexdigest()
74def collection_configuration(collection: Collection) -> str:
75 """Snapshot the stored settings that must survive an unfinished rebuild."""
76 return json.dumps(
77 {key: getattr(collection, key) for key in _CONFIG_FIELDS},
78 sort_keys=True,
79 default=str,
80 )
83def _index_state(session: Session, collection_name: str) -> list[tuple]:
84 return sorted(
85 (
86 row.id,
87 row.index_hash,
88 row.is_current,
89 row.chunk_count,
90 row.total_documents,
91 row.last_updated_at,
92 )
93 for row in session.query(RAGIndex).filter_by(
94 collection_name=collection_name
95 )
96 )
99def _membership_state(session: Session, collection_id: str) -> list[tuple]:
100 return sorted(
101 (
102 row.document_id,
103 row.id,
104 row.indexed,
105 row.chunk_count,
106 row.last_indexed_at,
107 )
108 for row in session.query(DocumentCollection).filter_by(
109 collection_id=collection_id
110 )
111 )
114def rebuild_stage_prefix(collection_id: str) -> str:
115 """Return the name prefix shared by every rebuild stage of a collection."""
116 scope = hashlib.sha256(collection_id.encode()).hexdigest()[:16]
117 return f"collection_rebuild_{scope}_"
120class RebuildCancelled(RuntimeError):
121 """Cancellation observed before publishing a replacement."""
124class CollectionRebuild:
125 """A private generation whose only publication point is a DB commit.
127 Chunk rows and text-free vector files use a unique staging identity.
128 Publishing changes their collection identity, the current-index pointer,
129 collection configuration and document status in one transaction. Files
130 retain their generation paths, so no filesystem rename can race that
131 transaction. A process interruption can leave an unpublished generation,
132 but cannot remove the previously published one.
134 The caller holds the per-collection rebuild lock through preparation and
135 cleanup. Ordinary indexing/search need only the shorter collection lock
136 around publication. Document embedding never holds that lock or a DB
137 transaction. Interrupted stages stay unselected until the collection is
138 deleted or its index cleared (``remove_unfinished_rebuilds``); another
139 process's unfinished generation must not be removed speculatively.
140 """
142 def __init__(
143 self,
144 service: LibraryRAGService,
145 collection_id: str,
146 doc_info: list[tuple[str, str]],
147 *,
148 session_factory: Callable[[], AbstractContextManager[Session]],
149 ):
150 # The generation switch and post-commit retirement below are atomic
151 # only for local files. A remote drop cannot be rolled back, and the
152 # current RAGIndex rows do not record which backend owns their data.
153 # Refuse rebuilding if any remote provider is registered until those
154 # providers have a separate publication protocol.
155 if not only_local_file_providers_registered():
156 raise RuntimeError(
157 "Staged collection rebuild requires local-file vector stores"
158 )
159 self.service = service
160 self.collection_id = collection_id
161 # nosemgrep: semgrep.rules.sql-string-concatenation, reason: Collection identifier passed as an ORM value; never SQL statement text
162 self.collection_name = f"collection_{collection_id}"
163 prefix = rebuild_stage_prefix(collection_id)
164 # nosemgrep: semgrep.rules.sql-string-concatenation, reason: Collection identifier passed as an ORM value; never SQL statement text
165 self.stage_name = f"{prefix}{uuid.uuid4().hex}"
166 self.stage_id = self.stage_name.removeprefix("collection_")
167 self.vector_index = None
168 self.index_id = None
169 self.prepared = {}
170 self.published = False
171 self.target_configuration = None
172 self._session = session_factory
174 with self._session() as session:
175 collection = session.get(Collection, collection_id)
176 if collection is None: 176 ↛ 177line 176 didn't jump to line 177 because the condition on line 176 was never true
177 raise ValueError("Collection not found")
178 self.configuration = collection_configuration(collection)
179 self.original_indices = _index_state(session, self.collection_name)
180 self.original_memberships = _membership_state(
181 session, collection_id
182 )
183 self.document_ids = {row[0] for row in self.original_memberships}
184 if self.document_ids != {doc_id for doc_id, _title in doc_info}: 184 ↛ 185line 184 didn't jump to line 185 because the condition on line 184 was never true
185 raise ValueError("Collection changed before rebuilding; retry")
186 # Persist before preparation so cancellation or a process exit
187 # cannot silently erase the unfinished request. The user DB is
188 # encrypted; no document content is copied into this marker.
189 collection.rebuild_pending_since = datetime.now(UTC)
190 session.commit()
191 logger.info(
192 "Rebuilding collection index; the published index, if any, "
193 "remains in use until the replacement is complete. "
194 "Cancelling before publication keeps that index and does not "
195 "retry the rebuild."
196 )
198 def prepare(
199 self,
200 document_id: str,
201 is_cancelled: Callable[[], bool] | None = None,
202 ) -> dict[str, Any]:
203 from .library_rag_service import _PreparedDocument
205 result = self.service._prepare_document(
206 document_id, self.collection_id, True, is_cancelled
207 )
208 if result.get("status") == "needs_serial_write":
209 # Empty documents have no replacement chunks. Never dispatch the
210 # ordinary empty-document purge against the published generation.
211 return {
212 "status": "prepared",
213 "prepared": _PreparedDocument(
214 document_id=document_id,
215 collection_id=self.collection_id,
216 chunk_inputs=[],
217 vectors=np.empty((0, 0), dtype="float32"),
218 source_fingerprint=result["source_fingerprint"],
219 ),
220 }
221 return result
223 def write(self, prepared: _PreparedDocument) -> dict[str, Any]:
224 """Persist only private chunks; leave published flags and files alone."""
225 if prepared.document_id not in self.document_ids: 225 ↛ 226line 225 didn't jump to line 226 because the condition on line 225 was never true
226 raise ValueError("Document does not belong to this rebuild")
227 if prepared.chunk_inputs:
228 if self.vector_index is None:
229 self.vector_index = self.service._get_vector_index(
230 self.stage_id, self.stage_name
231 )
232 self.index_id = self.service.rag_index_record.id
233 with self._session() as session:
234 try:
235 self.vector_index.index_prepared(
236 source_type="document",
237 source_id=prepared.document_id,
238 chunks=prepared.chunk_inputs,
239 vectors=prepared.vectors,
240 replace=True,
241 session=session,
242 )
243 index = session.get(RAGIndex, self.index_id)
244 index.status = RAGIndexStatus.REBUILDING
245 session.commit()
246 except Exception:
247 safe_rollback(session, "collection_rebuild.write")
248 raise
249 self.prepared[prepared.document_id] = (
250 prepared.source_fingerprint,
251 len(prepared.chunk_inputs),
252 )
253 return {
254 "status": "success" if prepared.chunk_inputs else "cleared",
255 "chunk_count": len(prepared.chunk_inputs),
256 }
258 def _validate_sources(self, session, collection):
259 if collection_configuration(collection) != self.configuration:
260 raise ValueError("Collection configuration changed; retry rebuild")
261 if _index_state(session, self.collection_name) != self.original_indices:
262 raise ValueError("Collection index changed; retry rebuild")
263 # Removing a document from search changes its membership's indexing
264 # state without changing the collection's RAGIndex counters. Preserve
265 # that successful removal instead of restoring its staged vectors.
266 if (
267 _membership_state(session, self.collection_id)
268 != self.original_memberships
269 ):
270 raise ValueError("Collection indexing state changed; retry rebuild")
271 seen = set()
272 # Bound retained source text during the final validation. A writer
273 # transaction prevents edits between validation and publication.
274 documents = (
275 session.query(Document)
276 .join(DocumentCollection)
277 .filter(DocumentCollection.collection_id == self.collection_id)
278 .yield_per(1)
279 )
280 for document in documents:
281 seen.add(document.id)
282 expected = self.prepared.get(document.id)
283 if (
284 expected is None
285 or document_fingerprint(document) != expected[0]
286 ):
287 raise ValueError("Collection documents changed; retry rebuild")
288 if seen != self.document_ids or set(self.prepared) != self.document_ids: 288 ↛ 289line 288 didn't jump to line 289 because the condition on line 288 was never true
289 raise ValueError("Collection documents changed; retry rebuild")
291 def _validate_vectors(self):
292 with self._session() as session:
293 chunk_ids = {
294 row.id
295 for row in session.query(DocumentChunk.id).filter_by(
296 collection_name=self.stage_name
297 )
298 }
299 # Reopen the durable generation: the in-memory facade can
300 # still contain every vector after a file was lost or a
301 # persist failed. Publication requires the bytes on disk.
302 durable_index = (
303 self.service._get_vector_index(self.stage_id, self.stage_name)
304 if self.vector_index is not None
305 else None
306 )
307 live_ids = set(durable_index.live_ids()) if durable_index else set()
308 expected_chunks = sum(
309 n for _fingerprint, n in self.prepared.values()
310 )
311 if chunk_ids != live_ids or len(chunk_ids) != expected_chunks:
312 raise RuntimeError("Replacement index is incomplete")
314 return expected_chunks
316 def _publish_rows(self, session, expected_chunks):
317 from ..deletion.utils.cascade_helper import CascadeHelper
319 # Acquire SQLite's writer lock before reading the state
320 # that must remain unchanged through the publication.
321 found = (
322 session.query(Collection)
323 .filter_by(id=self.collection_id)
324 .update({Collection.id: Collection.id})
325 )
326 if not found: 326 ↛ 327line 326 didn't jump to line 327 because the condition on line 326 was never true
327 raise ValueError("Collection was deleted during rebuild")
328 # Validation, every publication write and the commit must share the
329 # transaction that took the lock. No callback runs before the commit:
330 # a callback that opened get_user_db_session on this thread would
331 # share this session, and its error handler would roll it back.
332 transaction = session.get_transaction()
333 session.expire_all()
334 collection = session.get(Collection, self.collection_id)
335 self._validate_sources(session, collection)
337 CascadeHelper.delete_collection_chunks(session, self.collection_name)
338 retired = CascadeHelper.delete_rag_indices_for_collection(
339 session, self.collection_name, unlink_files=False
340 )
342 # Explicit deletion also works with SQLite connections
343 # whose foreign-key enforcement was not configured.
344 session.query(RagDocumentStatus).filter_by(
345 collection_id=self.collection_id
346 ).delete(synchronize_session=False)
347 session.query(DocumentCollection).filter_by(
348 collection_id=self.collection_id
349 ).update({"indexed": False, "chunk_count": 0})
350 session.query(DocumentChunk).filter_by(
351 collection_name=self.stage_name
352 ).update({"collection_name": self.collection_name})
354 timestamp = datetime.now(UTC)
355 index = (
356 session.get(RAGIndex, self.index_id)
357 if self.index_id is not None
358 else None
359 )
360 if index is not None:
361 index.collection_name = self.collection_name
362 index.is_current = True
363 index.status = RAGIndexStatus.ACTIVE
364 index.chunk_count = expected_chunks
365 index.total_documents = sum(
366 n > 0 for _fingerprint, n in self.prepared.values()
367 )
368 index.last_updated_at = timestamp
369 for document_id, (
370 _fingerprint,
371 count,
372 ) in self.prepared.items():
373 if count:
374 session.add(
375 RagDocumentStatus(
376 document_id=document_id,
377 collection_id=self.collection_id,
378 rag_index_id=self.index_id,
379 chunk_count=count,
380 indexed_at=timestamp,
381 )
382 )
383 session.query(DocumentCollection).filter_by(
384 document_id=document_id,
385 collection_id=self.collection_id,
386 ).update(
387 {
388 "indexed": True,
389 "chunk_count": count,
390 "last_indexed_at": timestamp,
391 }
392 )
393 collection.embedding_model = self.service.embedding_model
394 collection.embedding_model_type = EmbeddingProvider(
395 self.service.embedding_provider
396 )
397 collection.embedding_dimension = (
398 index.embedding_dimension if index else None
399 )
400 collection.chunk_size = self.service.chunk_size
401 collection.chunk_overlap = self.service.chunk_overlap
402 collection.splitter_type = self.service.splitter_type
403 collection.text_separators = self.service.text_separators
404 collection.distance_metric = self.service.distance_metric
405 collection.normalize_vectors = self.service.normalize_vectors
406 collection.index_type = self.service.index_type
408 # Clear only with the generation switch, including an empty rebuild.
409 collection.rebuild_pending_since = None
411 # A rollback since the lock was taken means the writes above run in a
412 # newer transaction that validation never checked, and may hold only
413 # part of the publication. Refuse to commit them, before recording
414 # the target that _is_published() compares after a failure.
415 if session.get_transaction() is not transaction:
416 raise RuntimeError(
417 "Collection publication lost its transaction; retry rebuild"
418 )
419 self.target_configuration = collection_configuration(collection)
420 return retired["index_paths"]
422 def publish(self, is_cancelled: Callable[[], bool] | None = None) -> None:
423 """Publish a complete generation, then remove retired vector files.
425 Cancellation is decided once, before the publication transaction
426 opens. After that the publication commits or rolls back as a whole.
427 """
428 expected_chunks = self._validate_vectors()
429 old_paths = []
430 with self.service._collection_transaction_lock(self.collection_id):
431 # The background task's check reads the task row through
432 # get_user_db_session. On this thread it gets the session that
433 # publishes below. Run it while no scope is open, so its
434 # autoflush and its error rollback cannot reach the publication.
435 self._check_cancelled(is_cancelled)
436 try:
437 with self._session() as session:
438 try:
439 old_paths = self._publish_rows(session, expected_chunks)
440 session.commit()
441 self.published = True
442 except Exception:
443 safe_rollback(session, "collection_rebuild.publish")
444 raise
445 except Exception:
446 # A post-commit callback can raise after publication. Use a
447 # fresh session: the failed one may still be in COMMITTED
448 # state, where even rollback and SELECT are unavailable.
449 with self._session() as session:
450 self.published = self._is_published(session)
451 if not self.published:
452 raise
453 logger.warning(
454 "Collection replacement committed despite a commit callback failure"
455 )
457 # Clearing an index and indexing again can reuse a retired path.
458 # Keep collection ownership through removal so retirement cannot
459 # delete the new writer's file. The DB commit has already finished.
460 self._unlink(old_paths)
461 from ...web_search_engines.search_engines_config import (
462 invalidate_collection_engines_cache,
463 )
465 invalidate_collection_engines_cache(self.service.username)
467 @staticmethod
468 def _check_cancelled(is_cancelled):
469 if is_cancelled is not None and is_cancelled():
470 raise RebuildCancelled("Collection rebuild was cancelled")
472 def _is_published(self, session):
473 if self.index_id is None:
474 collection = session.get(Collection, self.collection_id)
475 return bool(
476 self.target_configuration is not None
477 and collection is not None
478 and collection.rebuild_pending_since is None
479 and collection_configuration(collection)
480 == self.target_configuration
481 and not _index_state(session, self.collection_name)
482 )
483 index = session.get(RAGIndex, self.index_id)
484 return bool(index and index.collection_name == self.collection_name)
486 @staticmethod
487 def _unlink(paths):
488 from ..deletion.utils.cascade_helper import CascadeHelper
490 for path in paths:
491 CascadeHelper.delete_faiss_index_files(
492 # nosemgrep: semgrep.rules.path-traversal-risk, reason: Deletion helper enforces containment beneath the explicit rag_indices root
493 Path(path),
494 allowed_root=get_cache_directory() / "rag_indices",
495 )
497 def close(self):
498 """Remove this unpublished generation after its workers have drained."""
499 from ..deletion.utils.cascade_helper import CascadeHelper
501 if self.published:
502 return
503 try:
504 with self._session() as session:
505 if self._is_published(session): 505 ↛ 506line 505 didn't jump to line 506 because the condition on line 505 was never true
506 self.published = True
507 return
508 logger.warning(
509 "Collection rebuild did not publish a replacement. The "
510 "previous index, if any, remains in use and may be "
511 "outdated. Automatic indexing does not retry full "
512 "rebuilds; run Re-Index All again."
513 )
514 CascadeHelper.delete_collection_chunks(session, self.stage_name)
515 retired = CascadeHelper.delete_rag_indices_for_collection(
516 session, self.stage_name, unlink_files=False
517 )
518 session.commit()
519 self._unlink(retired["index_paths"])
520 except Exception:
521 logger.exception(
522 "Could not remove an unpublished collection rebuild"
523 )
525 def __enter__(self):
526 return self
528 def __exit__(self, exc_type, exc_value, traceback):
529 self.close()
530 return False
533def remove_unfinished_rebuilds(
534 username: str,
535 collection_id: str,
536 session_factory: Callable[[], AbstractContextManager[Session]],
537) -> int:
538 """Remove the stages that interrupted rebuilds of a collection left.
540 Collection deletion and index clearing call this after their commit.
541 Nothing selects a stage again, and its chunk rows hold document text.
542 Each stage is removed the way ``CollectionRebuild.close`` removes its
543 own. A rebuild of the collection running in this process holds the
544 rebuild lock and removes its stage when it cannot publish, so while that
545 lock is held every stage is left alone. Returns the number removed.
546 """
547 from ..deletion.utils.cascade_helper import CascadeHelper
548 from .library_rag_service import _hold_collection_lock
550 with _hold_collection_lock(
551 username, collection_id, purpose="rebuild", blocking=False
552 ) as held:
553 if not held:
554 return 0
555 prefix = rebuild_stage_prefix(collection_id)
556 paths = []
557 with session_factory() as session:
558 names = {
559 name
560 for model in (DocumentChunk, RAGIndex)
561 for (name,) in session.query(model.collection_name)
562 .filter(
563 model.collection_name.startswith(prefix, autoescape=True)
564 )
565 .distinct()
566 }
567 for name in sorted(names):
568 CascadeHelper.delete_collection_chunks(session, name)
569 paths += CascadeHelper.delete_rag_indices_for_collection(
570 session, name, unlink_files=False
571 )["index_paths"]
572 session.commit()
573 CollectionRebuild._unlink(paths)
574 return len(names)