Coverage for src/local_deep_research/vector_stores/legacy_cleanup.py: 81%
237 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 07:11 +0000
1"""One-time migration of legacy RAG index docstore files off disk.
3The previous RAG index format persisted a companion ``<hash>.pkl`` docstore next
4to each ``<hash>.faiss``. That docstore held the chunk **text** (plaintext at
5rest) plus the ``index_to_docstore_id`` map (FAISS position -> uuid) that the
6index needs. The current format stores no text in the index — text lives only in
7the encrypted DB and is rehydrated by id.
9This runs at startup (pure filesystem, no login / DB access). For each legacy
10``.pkl`` it:
12 1. extracts ONLY the ``index_to_docstore_id`` map (position -> uuid) into a
13 small, **text-free** ``<hash>.idmap.json`` sidecar (plain JSON — no pickle,
14 no text, no metadata), then
15 2. deletes the ``.pkl`` — removing all plaintext from disk.
17The per-user, login-time re-key (elsewhere) later consumes the sidecar to
18rebuild the index keyed by ``DocumentChunk.id`` (no re-embedding), then deletes
19the sidecar. If a ``.pkl`` can't be read (corrupt), it is deleted anyway
20(plaintext removed) and that collection falls back to a normal reindex.
22IMPORTANT: this MUST ship in the same release as the vector-store cutover.
23Before the cutover the old code still reads the ``.pkl``'s text for snippets, so
24removing it early would break search — do NOT wire this into startup until the
25old FAISS path is gone. It is exposed as a plain function so it can be called
26from tests and from the cutover's startup hook.
27"""
29import json
30import os
31import re
32import stat
33import tempfile
34import time
35from pathlib import Path
36from typing import Dict, List, Mapping, Tuple
38import numpy as np
39from faiss import (
40 IndexFlatIP,
41 IndexFlatL2,
42 IndexHNSWFlat,
43 IndexIDMap,
44 IndexIDMap2,
45 read_index,
46 vector_to_array,
47)
48from loguru import logger
50from ..config.paths import get_cache_directory
51from ..research_library.services.faiss_safe_load import (
52 load_index_to_docstore_id,
53)
55_RAG_CACHE_SUBDIR = "rag_indices"
56_SIDECAR_SUFFIX = ".idmap.json"
57# Only sweep sidecar ``.tmp`` files older than this — long enough that no single
58# write is still in flight, so a concurrent (multi-process) phase-1 run's temp is
59# never mistaken for an orphan. Mirrors faiss_store._STALE_TMP_AGE_SECONDS.
60_STALE_TMP_AGE_SECONDS = 3600
62# A legitimate docstore key is a ``uuid.uuid4().hex`` (32 hex chars, what
63# _store_chunks_to_db writes); accept the dashed uuid form too. Anything else in
64# a (tampered) ``.pkl`` — arbitrary text, whitespace — is rejected so it can
65# never be smuggled into the "text-free" sidecar.
66_ID_RE = re.compile(
67 r"[0-9a-f]{32}"
68 r"|[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}",
69 re.IGNORECASE,
70)
71# Anchored so a crafted name like ``x.pkl.corrupt-1.idmap.json`` is NOT matched.
72_QUARANTINE_RE = re.compile(r"\.pkl\.corrupt-[0-9-]+$")
75def _rag_cache_root() -> Path:
76 return get_cache_directory() / _RAG_CACHE_SUBDIR
79def _is_live_pkl(path: Path) -> bool:
80 """A live legacy docstore ``<hash>.pkl`` (has a re-keyable sibling index)."""
81 return path.is_file() and path.name.endswith(".pkl")
84def _is_quarantined_pkl(path: Path) -> bool:
85 """A quarantined docstore ``<hash>.pkl.corrupt-<ns>`` (dead, no sidecar)."""
86 return path.is_file() and _QUARANTINE_RE.search(path.name) is not None
89def _sidecar_for(pkl_path: Path) -> Path:
90 # <hash>.pkl -> <hash>.idmap.json
91 return pkl_path.with_name(pkl_path.stem + _SIDECAR_SUFFIX)
94def _find_legacy_docstores(root: Path) -> Tuple[List[Path], List[str]]:
95 """Return ``(found_docstores, unscannable_dirs)``.
97 A non-empty ``unscannable_dirs`` means ``os.walk`` could not read a
98 directory, so a plaintext ``.pkl`` inside it may survive UNDETECTED — the
99 caller must treat that as an INCOMPLETE migration, not clean success (a
100 re-scan is blind to the very same unreadable directory)."""
101 if not root.exists(): 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true
102 return [], []
104 scan_errors: List[str] = []
106 def _onerror(exc: OSError) -> None:
107 # os.walk (like rglob) SILENTLY skips a directory it can't scan. This
108 # migration's whole guarantee is "no plaintext .pkl left on disk", so a
109 # scan error must be LOUD and must fail the completeness check —
110 # otherwise an unreadable subdir could hide an un-deleted .pkl while the
111 # run reports clean success.
112 target = getattr(exc, "filename", "?")
113 scan_errors.append(str(target))
114 logger.error(
115 "RAG plaintext migration could not scan "
116 f"{target} ({exc}); a legacy .pkl there may "
117 "NOT have been removed — fix permissions and re-run, or delete it "
118 "manually."
119 )
121 found: List[Path] = []
122 for dirpath, dirs, files in os.walk(str(root), onerror=_onerror):
123 # os.walk does NOT descend into symlinked subdirectories
124 # (followlinks=False) and does NOT call onerror for them — they are
125 # silently skipped. A legacy .pkl inside one would survive undetected,
126 # so flag each symlinked subdir as unscanned (same as a permission
127 # error) rather than miss it. followlinks=True is avoided deliberately:
128 # a symlink cycle would hang the migration.
129 for d in dirs:
130 sub = Path(dirpath) / d
131 # A stat error (PermissionError, etc.) on a SINGLE entry must never
132 # abort the whole walk — that would leave EVERY .pkl on disk while
133 # the caller logs "migration failed". Flag the entry and keep going.
134 try:
135 is_link = sub.is_symlink()
136 except OSError:
137 scan_errors.append(str(sub))
138 logger.exception(
139 f"RAG plaintext migration could not stat directory {sub}; "
140 "a legacy .pkl there may NOT have been removed."
141 )
142 continue
143 if is_link:
144 scan_errors.append(str(sub))
145 logger.error(
146 "RAG plaintext migration did not descend into symlinked "
147 f"directory {sub}; a legacy .pkl inside may NOT have been "
148 "removed — replace the symlink with a real directory and "
149 "re-run, or delete the .pkl manually."
150 )
151 for name in files:
152 p = Path(dirpath) / name
153 try:
154 is_pkl = _is_live_pkl(p) or _is_quarantined_pkl(p)
155 except OSError:
156 scan_errors.append(str(p))
157 logger.exception(
158 f"RAG plaintext migration could not stat file {p}; "
159 "a legacy .pkl there may NOT have been removed."
160 )
161 continue
162 if is_pkl:
163 # A symlinked .pkl is a trap: _is_*_pkl / is_file() follow the
164 # link, so it looks migratable, but _delete_file's unlink()
165 # removes only the symlink — the real plaintext file it points
166 # to (potentially OUTSIDE the cache root) survives untouched
167 # while the run logs clean success. Don't follow it: flag it as
168 # unresolved (incomplete migration) so an operator removes the
169 # real file. We deliberately do NOT auto-delete the target — it
170 # can be an arbitrary path we must not unlink blindly. A stat
171 # error here is treated as "can't tell" -> unresolved.
172 try:
173 is_link = p.is_symlink()
174 except OSError:
175 is_link = True
176 if is_link:
177 scan_errors.append(str(p))
178 logger.error(
179 f"RAG plaintext migration found a SYMLINKED (or "
180 f"un-statable) docstore {p}; deleting the link would "
181 "leave the real plaintext file intact. Remove the real "
182 "target manually, then re-run."
183 )
184 continue
185 found.append(p)
186 return sorted(found), scan_errors
189def _delete_file(path: Path) -> None:
190 """Delete ``path`` cross-platform (Windows + POSIX).
192 Windows raises ``PermissionError`` when a file carries the read-only
193 attribute; clear it and retry once. On POSIX ``unlink`` succeeds even while a
194 file is open. A still-failing delete propagates to the caller, which records
195 it — the re-scan then reports it.
196 """
197 try:
198 path.unlink()
199 return
200 except FileNotFoundError:
201 return
202 except PermissionError:
203 pass
204 try:
205 os.chmod(path, stat.S_IWRITE | stat.S_IREAD)
206 except OSError:
207 pass
208 path.unlink()
211def _write_sidecar_atomic(sidecar: Path, id_map: Dict[str, str]) -> None:
212 """Atomically write the text-free position->uuid map as JSON."""
213 sidecar.parent.mkdir(parents=True, exist_ok=True)
214 try:
215 sidecar.parent.chmod(0o700)
216 except OSError:
217 logger.warning(f"Could not chmod {sidecar.parent} to 0o700")
218 # Sweep stale sidecar temps from a previously hard-killed migration. Phase-1
219 # normally runs single-threaded at startup, but under a multi-process server
220 # (gunicorn -w N) two workers can start it concurrently — only sweep temps
221 # OLDER than a generous threshold so a concurrent worker's just-created,
222 # still-in-flight temp for the same sidecar is never mistaken for an orphan
223 # and unlinked (which would crash its write). Mirrors FaissVectorStore.persist.
224 now = time.time()
225 for stale in sidecar.parent.glob(f"{sidecar.name}.*.tmp"):
226 try:
227 if now - stale.stat().st_mtime < _STALE_TMP_AGE_SECONDS:
228 continue
229 stale.unlink()
230 except OSError:
231 logger.warning(f"Could not remove stale sidecar temp {stale}")
232 # Safe: this writes only the text-free ``.idmap.json`` sidecar — a plain
233 # position->uuid JSON map (see module docstring / `_ID_RE`), never chunk
234 # text/metadata. Not sensitive data at rest.
235 fd, tmp_name = tempfile.mkstemp(
236 dir=str(sidecar.parent), prefix=f"{sidecar.name}.", suffix=".tmp"
237 )
238 tmp = Path(tmp_name)
239 try:
240 # Safe: same text-free id-map as above, written via the fd from
241 # mkstemp — no chunk text/metadata ever passes through this path.
242 with os.fdopen(fd, "w", encoding="utf-8") as f:
243 json.dump(id_map, f)
244 f.flush()
245 os.fsync(
246 f.fileno()
247 ) # durable before the rename (crash-consistency)
248 os.replace(tmp, sidecar)
249 # fsync the directory so the rename itself survives a crash (mirrors
250 # FaissVectorStore.persist); unsupported on some platforms.
251 try:
252 dir_fd = os.open(str(sidecar.parent), os.O_RDONLY)
253 try:
254 os.fsync(dir_fd)
255 finally:
256 os.close(dir_fd)
257 except OSError:
258 pass
259 finally:
260 if tmp.exists(): 260 ↛ 261line 260 didn't jump to line 261 because the condition on line 260 was never true
261 try:
262 tmp.unlink()
263 except OSError:
264 logger.warning(f"Could not remove temp sidecar {tmp}")
267def _extract_map(pkl_path: Path) -> bool:
268 """Extract the text-free position->uuid map to a sidecar next to ``pkl_path``.
270 Returns True if the sidecar was written (so the login re-key can rebuild
271 without re-embedding). Returns False if the ``.pkl`` couldn't be read — the
272 caller still deletes the ``.pkl`` (plaintext removal is not optional) and
273 that collection will fall back to a full reindex.
274 """
275 try:
276 raw = load_index_to_docstore_id(str(pkl_path))
277 except Exception as exc:
278 logger.warning(
279 f"Could not read legacy docstore {pkl_path} "
280 f"({type(exc).__name__}); that collection will reindex"
281 )
282 return False
283 # Validate shape + content. A malformed/tampered map (non-dict, non-int
284 # positions, or non-uuid values that could smuggle text into the "text-free"
285 # sidecar) falls to the reindex fallback — it must never crash the whole
286 # migration (a DoS leaving other users' plaintext on disk) nor leak content.
287 if not isinstance(raw, dict):
288 logger.warning(
289 f"Legacy docstore {pkl_path} has an unexpected id-map shape; "
290 "that collection will reindex"
291 )
292 return False
293 id_map: Dict[str, str] = {}
294 for pos, key in raw.items():
295 # Require key to be a real str and match it DIRECTLY (never str(key)).
296 # A tampered .pkl can put an arbitrary unpickled object here; calling
297 # str() on a maliciously deep/shared nested structure is a pickle
298 # "billion laughs" amplification that can hang/OOM the whole startup
299 # migration. A legitimate index_to_docstore_id value is always a uuid
300 # string, so a non-str (or non-uuid) value is simply rejected to the
301 # reindex fallback — it can never reach str().
302 # pos must be a plausible faiss array position. Bounding it also stops
303 # str(pos) below from raising ValueError on a maliciously huge int
304 # (Python's int->str digit limit), which would otherwise escape this
305 # loop uncaught and abort the whole startup migration — leaving other
306 # users' plaintext .pkl files on disk. bool is an int subclass; the
307 # range check harmlessly admits True/False (1/0), rejected downstream.
308 if (
309 not isinstance(pos, int)
310 or isinstance(pos, bool) # bool is an int subclass; True/False
311 # would serialize as "True"/"False" and never match a str(pos)
312 # lookup at rekey time, silently orphaning that chunk's vector.
313 or not (0 <= pos < 2**31)
314 or not isinstance(key, str)
315 or not _ID_RE.fullmatch(key)
316 ):
317 logger.warning(
318 f"Legacy docstore {pkl_path} has an unexpected id-map entry; "
319 "that collection will reindex"
320 )
321 return False
322 id_map[str(pos)] = key
323 try:
324 _write_sidecar_atomic(_sidecar_for(pkl_path), id_map)
325 except OSError:
326 logger.exception(
327 f"Could not write id-map sidecar for {pkl_path}; that "
328 "collection will reindex"
329 )
330 return False
331 return True
334def migrate_legacy_docstores() -> Dict[str, int]:
335 """Extract each legacy ``.pkl``'s id-map to a text-free sidecar, then delete
336 the ``.pkl``, removing all plaintext docstores from disk.
338 Returns ``{"found", "extracted", "deleted", "reindex_fallback",
339 "remaining"}``. Re-scans afterward and logs a prominent error recommending
340 manual deletion if any docstore remains — plaintext must never pass
341 silently. Idempotent (a clean tree is a no-op).
342 """
343 root = _rag_cache_root()
344 # SECURITY: refuse to operate on a SYMLINKED cache root. os.walk (and
345 # Path.chmod below) follow a symlinked TOP-LEVEL path unconditionally —
346 # followlinks=False only guards descent into symlinked CHILDREN found
347 # mid-walk (which _find_legacy_docstores already flags). A symlinked
348 # rag_indices/ would otherwise make this migration chmod 0o700 and
349 # pattern-delete .pkl files OUTSIDE the intended tree (an admin who
350 # symlinked the cache onto a bigger volume, or a hostile account that can
351 # write the cache parent), reporting clean success. Treat it exactly like a
352 # symlinked subdir: log loudly and refuse rather than following it. Uses
353 # is_symlink() (an lstat, does NOT follow) so it is safe to call first.
354 if root.is_symlink():
355 logger.error(
356 f"RAG cache root {root} is a symlink; refusing to migrate through "
357 "it (following it would chmod/delete files outside the intended "
358 "cache tree). Replace it with a real directory and restart."
359 )
360 return {
361 "found": 0,
362 "extracted": 0,
363 "deleted": 0,
364 "reindex_fallback": 0,
365 "remaining": 0,
366 "scan_errors": 1,
367 }
368 # Harden the RAG cache ROOT at the earliest point it is touched so another
369 # local OS account can't traverse into any user's vector caches. Both the
370 # leaf-only chmod in _write_sidecar_atomic and _get_index_path otherwise
371 # leave the root at the process umask (typically world-traversable) between
372 # startup and the first index op.
373 if root.exists(): 373 ↛ 378line 373 didn't jump to line 378 because the condition on line 373 was always true
374 try:
375 root.chmod(0o700)
376 except OSError:
377 logger.warning(f"Could not chmod RAG cache root {root} to 0o700")
378 found, scan_errors = _find_legacy_docstores(root)
379 if not found and not scan_errors:
380 logger.debug(f"RAG cache: no legacy docstore files under {root}")
381 return {
382 "found": 0,
383 "extracted": 0,
384 "deleted": 0,
385 "reindex_fallback": 0,
386 "remaining": 0,
387 "scan_errors": 0,
388 }
390 extracted = 0
391 deleted = 0
392 reindex_fallback = 0
393 for pkl in found:
394 # Live .pkl: preserve the re-key map before deleting. Quarantined
395 # (.pkl.corrupt-*): dead — just delete, no sidecar.
396 if _is_live_pkl(pkl):
397 if _extract_map(pkl):
398 extracted += 1
399 else:
400 reindex_fallback += 1
401 try:
402 _delete_file(pkl)
403 deleted += 1
404 except OSError:
405 logger.exception(f"RAG cache: could not delete {pkl}")
407 # Re-scan is the source of truth (partial failures, POSIX unlink-of-open).
408 # A scan error is NOT clean success: an unreadable directory hides its
409 # contents from BOTH the delete pass above and this re-scan, so a plaintext
410 # .pkl inside it could survive while `remaining` reads empty. Treat any
411 # unscannable dir as an unresolved leftover.
412 remaining, rescan_errors = _find_legacy_docstores(root)
413 if remaining or rescan_errors:
414 parts = []
415 if remaining: 415 ↛ 416line 415 didn't jump to line 416 because the condition on line 415 was never true
416 parts.append(
417 f"{len(remaining)} legacy docstore file(s) remain on disk: "
418 + ", ".join(str(p) for p in remaining)
419 )
420 if rescan_errors: 420 ↛ 426line 420 didn't jump to line 426 because the condition on line 420 was always true
421 parts.append(
422 f"{len(rescan_errors)} director(ies) could not be scanned "
423 "(a plaintext .pkl inside may NOT have been removed): "
424 + ", ".join(rescan_errors)
425 )
426 logger.error(
427 "RAG cache: plaintext migration did NOT complete cleanly — "
428 + "; ".join(parts)
429 + ". Fix permissions and re-run, or delete the file(s) manually."
430 )
431 else:
432 logger.info(
433 f"RAG cache: extracted {extracted} id-map(s), removed {deleted} "
434 f"legacy docstore file(s) from {root}"
435 + (
436 f" ({reindex_fallback} unreadable -> will reindex)"
437 if reindex_fallback
438 else ""
439 )
440 )
441 return {
442 "found": len(found),
443 "extracted": extracted,
444 "deleted": deleted,
445 "reindex_fallback": reindex_fallback,
446 "remaining": len(remaining),
447 "scan_errors": len(rescan_errors),
448 }
451def rekey_index_file(
452 faiss_path: Path,
453 sidecar_path: Path,
454 uuid_to_id: Mapping[str, int],
455 *,
456 dimension: int,
457 index_type: str,
458 metric: str,
459 normalize: bool,
460) -> Dict[str, int]:
461 """Convert one old-format index to the new ``IndexIDMap2`` keyed by int id,
462 WITHOUT re-embedding — the login-time (phase-2) half of the migration.
464 Reads the raw old ``.faiss`` (reconstructs each vector by position), maps
465 position -> uuid (from the ``.idmap.json`` sidecar) -> int ``DocumentChunk.id``
466 (via ``uuid_to_id``), builds a fresh ``IndexIDMap2`` keyed by those ints,
467 writes it atomically over ``faiss_path``, and **reads it back and verifies**.
468 Positions whose uuid has no DB row are orphans and are dropped.
470 Does NOT delete the sidecar (nor acquire the write lock, nor re-record file
471 integrity) — the caller owns all three. The caller deletes the sidecar ONLY
472 after recording the new file's integrity, so a crash/failure between persist
473 and integrity-record leaves the sidecar as a retry signal; the ``IndexIDMap2``
474 guard above then makes that retry a clean no-op. Raises on a failed read-back
475 verification (the caller quarantines + falls back to a normal reindex).
476 """
477 # Import here to avoid a module-load cycle (implementations import base only).
478 from .implementations.faiss_store import FaissVectorStore
480 faiss_path = Path(faiss_path)
481 sidecar_path = Path(sidecar_path)
483 old = read_index(str(faiss_path))
485 # Crash-consistency guard. rekey persists the new IndexIDMap2 (atomic
486 # replace) and only unlinks the sidecar AFTER a successful read-back verify.
487 # If the process dies in that window, faiss_path is ALREADY the new format
488 # but the sidecar survives, so the next login re-enters here on an
489 # already-migrated file. Running the reconstruct path below on an
490 # IndexIDMap2 is not merely wrong — reconstruct_n(0, ntotal) would treat
491 # 0..ntotal-1 as *ids* to look up; a missing id raises a C++ faiss
492 # exception that SWIG does NOT surface as a catchable Python exception on
493 # this call path, aborting the whole process (and, if the int ids happen to
494 # fall in 0..ntotal-1, silently reconstructing the WRONG vectors instead).
495 # Either way the re-key is already done — report success WITHOUT deleting
496 # the sidecar. The CALLER owns sidecar deletion (only after it has recorded
497 # file integrity), so an interrupted prior run that persisted the new index
498 # but hadn't yet recorded integrity is recoverable: this re-entry succeeds,
499 # the caller (re)records integrity, THEN drops the sidecar. Deleting it here
500 # would lose that retry signal.
501 if isinstance(old, IndexIDMap2):
502 # Same dimension guard the fresh re-key path enforces below: an
503 # already-migrated index whose width no longer matches the configured
504 # embedding model (model changed post-migration) must be rebuilt, not
505 # silently accepted. Raise so the caller quarantines + reindexes now,
506 # rather than deferring to an opaque dimension failure at first search.
507 if old.d != dimension: 507 ↛ 508line 507 didn't jump to line 508 because the condition on line 507 was never true
508 raise ValueError(
509 f"rekey {faiss_path}: already-migrated index dim {old.d} != "
510 f"expected {dimension} (embedding model changed) — reindex"
511 )
512 # Cross-check the on-disk ids against what the CURRENT sidecar + DB
513 # resolution says they should be — the fresh path below does this via a
514 # read-back verify; this short-circuit must too. Without it, a stale or
515 # FOREIGN IndexIDMap2 sitting at this path (e.g. two collections/users
516 # colliding on the pre-per-user-scoping shared cache path, or a retry
517 # after the uuid->id resolution changed) would be silently adopted, and
518 # its baked-in ids would rehydrate the WRONG document's text on search.
519 try:
520 pos_to_uuid = json.loads(sidecar_path.read_text(encoding="utf-8"))
521 except FileNotFoundError:
522 # The sidecar is the caller's COMPLETION SIGNAL: it is deleted only
523 # AFTER a successful re-key records the new index's integrity (see
524 # rekey_user_indexes Phase B). A concurrent worker — another process
525 # (gunicorn -w N / the scheduler) with no cross-process mutex, or a
526 # thread whose sidecar check passed before that worker won the lock —
527 # can finalize and delete it between our caller's sidecar check and
528 # this read. A missing sidecar next to an already-IndexIDMap2 file of
529 # the right dimension therefore means "already finalized by another
530 # worker" — the SAME signature rekey_user_indexes treats as done at
531 # its top-of-loop no-sidecar branch (skipped_no_sidecar). Report
532 # success rather than letting the FileNotFoundError bubble up and
533 # QUARANTINE a healthy, just-rekeyed index (forcing a needless full
534 # re-embed). No id cross-check: there is no sidecar left to check
535 # against, exactly as the caller's no-sidecar branch also skips it.
536 logger.info(
537 f"Re-key skipped for {faiss_path.name}: already migrated to "
538 "IndexIDMap2; its sidecar was removed by a concurrent finalizer"
539 )
540 return {
541 "reconstructed": int(old.ntotal),
542 "kept": int(old.ntotal),
543 "orphaned": 0,
544 }
545 expected_ids = {
546 int(uuid_to_id[u]) for u in pos_to_uuid.values() if u in uuid_to_id
547 }
548 on_disk_id_list = (
549 vector_to_array(old.id_map).astype("int64").tolist()
550 if old.ntotal
551 else []
552 )
553 # Multiplicity, not just membership: a duplicate id (e.g. [5, 5, 6, 7])
554 # has the SAME set as {5, 6, 7}, so the set comparison below cannot see
555 # it (sets dedupe) — exactly the corruption the fresh-build path rejects
556 # at ``len(set(kept_ids)) != len(kept_ids)``. A duplicate baked into an
557 # already-migrated IndexIDMap2 double-counts/mis-scores that chunk in
558 # search, so reject it here too rather than silently adopt it.
559 if len(set(on_disk_id_list)) != len(on_disk_id_list):
560 raise ValueError(
561 f"rekey {faiss_path}: already-migrated index has duplicate "
562 "DocumentChunk.id(s) — refusing to adopt an id-collided index; "
563 "quarantine + reindex"
564 )
565 if set(on_disk_id_list) != expected_ids:
566 raise ValueError(
567 f"rekey {faiss_path}: already-migrated index ids do not match "
568 "the sidecar/DB resolution — refusing to adopt a stale/foreign "
569 "index; quarantine + reindex"
570 )
571 logger.info(
572 f"Re-key skipped for {faiss_path.name}: already migrated to "
573 "IndexIDMap2 (interrupted prior run — caller will finalize)"
574 )
575 return {
576 "reconstructed": int(old.ntotal),
577 "kept": int(old.ntotal),
578 "orphaned": 0,
579 }
581 # The crash hazard is the WHOLE IndexIDMap family, not just IndexIDMap2:
582 # reconstruct_n on a plain IndexIDMap (our own format's superclass) also
583 # raises an untranslated C++ faiss exception that aborts the process
584 # (SIGABRT), bypassing every try/except. A plain IndexIDMap is never a
585 # format we write (we only ever persist IndexIDMap2) nor the pre-migration
586 # raw base index we mean to reconstruct here — so it is foreign/corrupt.
587 # Raise a CATCHABLE ValueError so the caller quarantines + reindexes,
588 # instead of falling through to the process-killing reconstruct_n below.
589 if isinstance(old, IndexIDMap):
590 raise ValueError(
591 f"rekey {faiss_path}: on-disk index is {type(old).__name__} "
592 "(IndexIDMap family but not IndexIDMap2) — refusing to reconstruct "
593 "a foreign/corrupt index; quarantine + reindex instead"
594 )
596 # Positive whitelist. The IndexIDMap guard above is only one of MANY faiss
597 # classes that abort the process on reconstruct_n: any index type that isn't
598 # a plain reconstruct-safe base index (e.g. a corrupt/tampered/foreign file
599 # that happens to decode as IndexIVFPQFastScan) segfaults at reconstruct_n
600 # below — an untranslated C++ abort that bypasses the caller's try/except
601 # and, on the login-time rekey path, crash-loops the worker. Only the base
602 # classes _build_base_index actually produces are reconstruct-safe here;
603 # reject anything else with a CATCHABLE ValueError so the caller quarantines
604 # + reindexes.
605 if not isinstance(old, (IndexFlatL2, IndexFlatIP, IndexHNSWFlat)):
606 raise ValueError(
607 f"rekey {faiss_path}: on-disk index is {type(old).__name__}, not a "
608 "reconstruct-safe base index (IndexFlatL2/IndexFlatIP/IndexHNSWFlat)"
609 " — refusing a foreign/corrupt index; quarantine + reindex instead"
610 )
612 ntotal = int(old.ntotal)
613 if old.d != dimension: 613 ↛ 614line 613 didn't jump to line 614 because the condition on line 613 was never true
614 raise ValueError(
615 f"rekey {faiss_path}: index dim {old.d} != expected {dimension}"
616 )
618 try:
619 pos_to_uuid = json.loads(sidecar_path.read_text(encoding="utf-8"))
620 except FileNotFoundError:
621 # A concurrent worker (multi-process: gunicorn -w N / the scheduler,
622 # with no cross-process lock) may have finished re-keying this SAME
623 # index between our read_index() above and this read — persisting a new
624 # IndexIDMap2 and deleting the sidecar as its completion signal
625 # (legacy_rekey Phase B, which persists the new index BEFORE unlinking
626 # the sidecar). Re-read the on-disk file: if it is now an IndexIDMap2 of
627 # the right dimension, the migration is already done, so report success
628 # rather than letting FileNotFoundError bubble up — the caller treats
629 # that as un-rekeyable and QUARANTINES a perfectly good, just-migrated
630 # index (forcing a needless full re-embed). Only if the file is still
631 # NOT migrated (a genuinely missing sidecar next to an old-format file)
632 # do we re-raise so the caller quarantines. This mirrors the
633 # already-migrated short-circuit's own FileNotFoundError handling above.
634 now_on_disk = None
635 try:
636 now_on_disk = read_index(str(faiss_path))
637 except Exception:
638 # Probe only: a failed/partial read just means "not a valid
639 # migrated index yet" — now_on_disk stays None and we re-raise the
640 # original FileNotFoundError below.
641 logger.debug(
642 f"Concurrent-migration re-read probe failed for {faiss_path.name}"
643 )
644 if isinstance(now_on_disk, IndexIDMap2) and now_on_disk.d == dimension: 644 ↛ 654line 644 didn't jump to line 654 because the condition on line 644 was always true
645 logger.info(
646 f"Re-key skipped for {faiss_path.name}: a concurrent worker "
647 "migrated it (sidecar already removed)"
648 )
649 return {
650 "reconstructed": int(now_on_disk.ntotal),
651 "kept": int(now_on_disk.ntotal),
652 "orphaned": 0,
653 }
654 raise
656 # Reconstruct ALL vectors by position in one bounded call. reconstruct_n only
657 # touches positions 0..ntotal-1, avoiding the out-of-range garbage that a
658 # hand-rolled reconstruct(i) loop risks (faiss #4413).
659 if ntotal: 659 ↛ 662line 659 didn't jump to line 662 because the condition on line 659 was always true
660 vectors = np.asarray(old.reconstruct_n(0, ntotal), dtype="float32")
661 else:
662 vectors = np.empty((0, dimension), dtype="float32")
664 kept_ids: List[int] = []
665 kept_rows: List[int] = []
666 orphaned = 0
667 for pos in range(ntotal):
668 uid = pos_to_uuid.get(str(pos))
669 int_id = uuid_to_id.get(uid) if uid is not None else None
670 if int_id is None:
671 orphaned += 1
672 continue
673 kept_ids.append(int(int_id))
674 kept_rows.append(pos)
676 # Two legacy positions must never map to the SAME DocumentChunk.id: that
677 # would bake a duplicate id into the IndexIDMap2 (which tolerates dupes),
678 # producing phantom, mis-scored search hits — and the set()-based read-back
679 # verify below cannot see it (sets dedupe). A duplicate means the sidecar /
680 # uuid->id resolution is inconsistent, so refuse and let the caller
681 # quarantine + reindex cleanly from the DB rather than persist corruption.
682 if len(set(kept_ids)) != len(kept_ids):
683 raise ValueError(
684 f"rekey {faiss_path}: {len(kept_ids) - len(set(kept_ids))} duplicate "
685 "DocumentChunk.id(s) from the legacy id-map — refusing to persist an "
686 "id-collided index; quarantine + reindex instead"
687 )
689 kept_vecs = (
690 vectors[kept_rows]
691 if kept_rows
692 else np.empty((0, dimension), dtype="float32")
693 )
695 store = FaissVectorStore.create(
696 dimension=dimension,
697 index_type=index_type,
698 metric=metric,
699 normalize=normalize,
700 )
701 if kept_ids: 701 ↛ 703line 701 didn't jump to line 703 because the condition on line 701 was always true
702 store.add(kept_ids, kept_vecs)
703 store.persist(faiss_path) # atomic temp + os.replace
705 # Read-back verify BEFORE removing the sidecar (the only remaining re-key
706 # input): the new file must load as IndexIDMap2 with exactly the kept ids.
707 verify = FaissVectorStore.load(
708 faiss_path,
709 dimension=dimension,
710 index_type=index_type,
711 metric=metric,
712 normalize=normalize,
713 )
714 if set(verify.live_ids()) != set(kept_ids): 714 ↛ 715line 714 didn't jump to line 715 because the condition on line 714 was never true
715 raise RuntimeError(
716 f"rekey {faiss_path}: read-back verify failed "
717 f"({verify.count()} ids on disk != {len(kept_ids)} expected)"
718 )
720 # NOTE: the sidecar is deliberately NOT deleted here. The caller deletes it
721 # ONLY after it has recorded the new file's integrity, so a crash/failure
722 # between persist and integrity-record leaves the sidecar as a retry signal
723 # (the IndexIDMap2 guard at the top makes the retry a clean no-op re-key).
724 logger.info(
725 f"Re-keyed {faiss_path.name}: {len(kept_ids)} vectors "
726 f"({orphaned} orphaned, {ntotal} reconstructed)"
727 )
728 return {
729 "reconstructed": ntotal,
730 "kept": len(kept_ids),
731 "orphaned": orphaned,
732 }