From 360b35d6e43cb50b78bd178b9cec78e402098811 Mon Sep 17 00:00:00 2001 From: Arnav Agrawal Date: Tue, 8 Sep 2026 16:49:41 -0700 Subject: [PATCH] Fix document update scheduling with ingestion revisions --- core/database/postgres_database.py | 30 +- core/routes/ingest.py | 59 ++- core/services/ingestion_service.py | 70 +++- .../test_document_update_revisions.py | 357 ++++++++++++++++++ .../test_ingestion_service_metadata_update.py | 8 +- core/workers/ingestion_worker.py | 97 ++++- docs/document-updates.md | 57 +++ docs/iqor-on-prem.md | 6 +- 8 files changed, 652 insertions(+), 32 deletions(-) create mode 100644 core/tests/integration/test_document_update_revisions.py create mode 100644 docs/document-updates.md diff --git a/core/database/postgres_database.py b/core/database/postgres_database.py index 9de86d31..8197ac75 100644 --- a/core/database/postgres_database.py +++ b/core/database/postgres_database.py @@ -1,5 +1,6 @@ import json import logging +from contextlib import asynccontextmanager from datetime import UTC, datetime from typing import Any, Dict, List, Optional @@ -7,6 +8,7 @@ from sqlalchemy.exc import ProgrammingError from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import NullPool from core.config import get_settings from core.utils.folder_utils import normalize_folder_path @@ -172,8 +174,30 @@ def __init__( connect_args={"server_settings": {"statement_timeout": "30000"}}, # 30 second timeout ) self.async_session = sessionmaker(self.engine, class_=AsyncSession, expire_on_commit=False) + # Ingestion holds a lock across parsing and embedding. Keep those connections + # out of the query pool so busy workers cannot exhaust it and deadlock writes. + self._ingestion_lock_engine = create_async_engine(uri, poolclass=NullPool) self._initialized = False + @asynccontextmanager + async def document_ingestion_lock(self, document_id: str, *, wait: bool = False): + """Serialize content replacement and all worker writes for one document. + + Transaction-scoped advisory locks are released on cancellation/disconnect, + with no expiring lease that could admit another writer during a long ingest. + Callers must read the document/revision *after* acquiring the lock. + """ + async with self._ingestion_lock_engine.begin() as connection: + key = {"key": f"document-ingestion:{document_id}"} + if wait: + await connection.execute(text("SELECT pg_advisory_xact_lock(hashtextextended(:key, 0))"), key) + acquired = True + else: + acquired = await connection.scalar( + text("SELECT pg_try_advisory_xact_lock(hashtextextended(:key, 0))"), key + ) + yield acquired + async def initialize(self): """Initialize database tables and indexes.""" if self._initialized: @@ -296,7 +320,9 @@ async def store_document( logger.error(f"Error storing document metadata: {str(e)}") return False - async def get_document(self, document_id: str, auth: AuthContext) -> Optional[Document]: + async def get_document( + self, document_id: str, auth: AuthContext, *, raise_on_error: bool = False + ) -> Optional[Document]: """Retrieve document metadata by ID if user has access.""" try: async with self.async_session() as session: @@ -320,6 +346,8 @@ async def get_document(self, document_id: str, auth: AuthContext) -> Optional[Do except Exception as e: logger.error(f"Error retrieving document metadata: {str(e)}") + if raise_on_error: + raise return None async def get_document_by_filename( diff --git a/core/routes/ingest.py b/core/routes/ingest.py index 55086891..4097c039 100644 --- a/core/routes/ingest.py +++ b/core/routes/ingest.py @@ -3,6 +3,7 @@ from typing import Any, Dict, List, Optional, Set import arq +from arq.jobs import Job, JobStatus from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile from core.auth_utils import verify_token @@ -287,11 +288,39 @@ async def requeue_ingest_jobs( results: List[RequeueIngestionResult] = [] async def _process_document(doc: Document, override_flag: Optional[bool]) -> None: + if doc.external_id in processed_ids: + return + async with ingestion_service.db.document_ingestion_lock(doc.external_id) as acquired: + if not acquired: + results.append( + RequeueIngestionResult( + external_id=doc.external_id, + status="already_queued", + message="Document ingestion is in progress", + ) + ) + processed_ids.add(doc.external_id) + return + current = await ingestion_service.db.get_document(doc.external_id, auth, raise_on_error=True) + if current is None: + results.append( + RequeueIngestionResult( + external_id=doc.external_id, + status="error", + message="Document no longer exists", + ) + ) + processed_ids.add(doc.external_id) + return + await _process_locked_document(current, override_flag) + + async def _process_locked_document(doc: Document, override_flag: Optional[bool]) -> None: ext_id = doc.external_id if ext_id in processed_ids: return processed_ids.add(ext_id) + revision_persisted = False try: auth_for_doc = AuthContext( @@ -332,11 +361,28 @@ async def _process_document(doc: Document, override_flag: Optional[bool]) -> Non if isinstance(system_metadata, str): system_metadata = json.loads(system_metadata) sanitized_system_metadata = IngestionService._reset_processing_metadata(system_metadata) - await ingestion_service.db.update_document( + revision = int(system_metadata.get("ingestion_revision", 0)) + current_job = Job(f"ingest:{ext_id}:{revision}", redis, _queue_name=redis.default_queue_name) + if await current_job.status() in {JobStatus.queued, JobStatus.deferred, JobStatus.in_progress}: + results.append( + RequeueIngestionResult( + external_id=ext_id, + status="already_queued", + message="An ingestion job is already pending", + ) + ) + return + # A retained result may represent a failed attempt. A manual requeue + # gets a new revision, fencing off any delayed attempt of the old job. + sanitized_system_metadata["ingestion_revision"] = revision + 1 + success = await ingestion_service.db.update_document( document_id=ext_id, updates={"system_metadata": sanitized_system_metadata}, auth=auth_for_doc, ) + if not success: + raise RuntimeError("Failed to persist requeue revision") + revision_persisted = True job_payload = IngestionService._build_ingestion_job_payload( document_id=ext_id, file_key=key, @@ -349,17 +395,12 @@ async def _process_document(doc: Document, override_flag: Optional[bool]) -> Non folder_path=doc.folder_path, folder_leaf=doc.folder_name, end_user_id=doc.end_user_id, + ingestion_revision=revision + 1, ) job = await redis.enqueue_job("process_ingestion_job", **job_payload) if job is None: - results.append( - RequeueIngestionResult( - external_id=ext_id, - status="already_queued", - message="An ingestion job is already pending for this document", - ) - ) + raise RuntimeError("Requeue job ID is already present in Redis; no new job was queued") else: results.append( RequeueIngestionResult( @@ -372,6 +413,8 @@ async def _process_document(doc: Document, override_flag: Optional[bool]) -> Non raise except Exception as exc: # noqa: BLE001 logger.error("Failed to requeue ingestion for document %s: %s", ext_id, exc, exc_info=True) + if revision_persisted: + await ingestion_service._mark_document_failed(doc, auth, f"Failed to requeue ingestion: {exc}") results.append( RequeueIngestionResult( external_id=ext_id, diff --git a/core/services/ingestion_service.py b/core/services/ingestion_service.py index 8f3120df..5e5394a7 100644 --- a/core/services/ingestion_service.py +++ b/core/services/ingestion_service.py @@ -428,9 +428,11 @@ def _build_ingestion_job_payload( folder_path: Optional[str] = None, folder_leaf: Optional[str] = None, end_user_id: Optional[str] = None, + ingestion_revision: int = 0, ) -> Dict[str, Any]: return { - "_job_id": f"ingest:{document_id}", + "_job_id": f"ingest:{document_id}:{ingestion_revision}", + "ingestion_revision": ingestion_revision, "_expires": timedelta(days=7), "document_id": document_id, "file_key": file_key, @@ -562,6 +564,38 @@ async def ingest_file_content( end_user_id: Optional[str] = None, use_colpali: Optional[bool] = False, external_id: Optional[str] = None, + ) -> Document: + document_id = external_id or str(uuid.uuid4()) + async with self.db.document_ingestion_lock(document_id) as acquired: + if not acquired: + raise HTTPException(status_code=409, detail="Document ingestion is already in progress; retry later") + return await self._ingest_file_content_locked( + file_content_bytes, + filename, + content_type, + metadata, + auth, + redis, + metadata_types, + folder_name, + end_user_id, + use_colpali, + document_id, + ) + + async def _ingest_file_content_locked( + self, + file_content_bytes: bytes, + filename: str, + content_type: Optional[str], + metadata: Optional[Dict[str, Any]], + auth: AuthContext, + redis: arq.ArqRedis, + metadata_types: Optional[Dict[str, str]], + folder_name: Optional[Union[str, List[str]]], + end_user_id: Optional[str], + use_colpali: Optional[bool], + external_id: str, ) -> Document: """ Ingests file content from bytes. Saves to storage, creates document record, @@ -594,6 +628,7 @@ async def ingest_file_content( folder_path=folder_path, ) doc.system_metadata = self._reset_processing_metadata(doc.system_metadata) + doc.system_metadata["ingestion_revision"] = 0 await self._verify_ingest_and_storage_limits(auth, len(file_content_bytes), doc.external_id) @@ -685,7 +720,7 @@ async def ingest_file_content( ) job = await redis.enqueue_job("process_ingestion_job", **job_payload) if job is None: - logger.info("Connector file ingestion job already queued (doc_id=%s)", doc.external_id) + raise RuntimeError("Ingestion job ID is already present in Redis; no new job was queued") else: logger.info( "Connector file ingestion job queued with ID: %s for document: %s", job.job_id, doc.external_id @@ -712,6 +747,33 @@ async def queue_document_update( metadata: Optional[Dict[str, Any]] = None, metadata_types: Optional[Dict[str, str]] = None, use_colpali: Optional[bool] = None, + ) -> Optional[Document]: + async with self.db.document_ingestion_lock(document_id) as acquired: + if not acquired: + raise HTTPException(status_code=409, detail="Document ingestion is already in progress; retry later") + return await self._queue_document_update_locked( + document_id, + auth, + redis, + content, + file, + filename, + metadata, + metadata_types, + use_colpali, + ) + + async def _queue_document_update_locked( + self, + document_id: str, + auth: AuthContext, + redis: arq.ArqRedis, + content: Optional[str], + file: Optional[UploadFile], + filename: Optional[str], + metadata: Optional[Dict[str, Any]], + metadata_types: Optional[Dict[str, str]], + use_colpali: Optional[bool], ) -> Optional[Document]: """ Update a document by replacing its content and re-queueing ingestion. @@ -781,6 +843,7 @@ async def queue_document_update( raise HTTPException(status_code=500, detail=f"Failed to upload updated file to storage: {str(e)}") doc.system_metadata = self._reset_processing_metadata(doc.system_metadata) + doc.system_metadata["ingestion_revision"] = int(doc.system_metadata.get("ingestion_revision", 0)) + 1 updates = { "metadata": doc.metadata, @@ -832,10 +895,11 @@ async def queue_document_update( folder_path=doc.folder_path, folder_leaf=doc.folder_name, end_user_id=doc.end_user_id, + ingestion_revision=doc.system_metadata["ingestion_revision"], ) job = await redis.enqueue_job("process_ingestion_job", **job_payload) if job is None: - logger.info("Update ingestion job already queued (doc_id=%s)", doc.external_id) + raise RuntimeError("Update job ID is already present in Redis; no new job was queued") else: logger.info("Update ingestion job queued (job_id=%s, doc=%s)", job.job_id, doc.external_id) except Exception as e: diff --git a/core/tests/integration/test_document_update_revisions.py b/core/tests/integration/test_document_update_revisions.py new file mode 100644 index 00000000..f1f7d607 --- /dev/null +++ b/core/tests/integration/test_document_update_revisions.py @@ -0,0 +1,357 @@ +"""Real Redis/ARQ and PostgreSQL/pgvector regression tests for content updates. + +Set CORE_UPDATE_TEST_POSTGRES_URI and CORE_UPDATE_TEST_REDIS_URL to disposable +services. Documents and queue keys are isolated per test; Redis is never flushed. +The parser and worker are real. Embeddings are deterministic and local here; +provider-backed API/download/retrieval verification is a separate runtime proof. +""" + +import asyncio +import os +import uuid +from io import BytesIO +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from arq import create_pool +from arq.connections import RedisSettings +from arq.jobs import Job, JobStatus +from arq.worker import Retry, Worker +from fastapi import HTTPException, UploadFile +from sqlalchemy import text +from sqlalchemy.exc import OperationalError + +pytestmark = pytest.mark.integration +POSTGRES_URI = os.environ.get("CORE_UPDATE_TEST_POSTGRES_URI") +REDIS_URL = os.environ.get("CORE_UPDATE_TEST_REDIS_URL") + + +@pytest.fixture +async def runtime(tmp_path, monkeypatch): + if not POSTGRES_URI or not REDIS_URL: + pytest.skip("Set CORE_UPDATE_TEST_POSTGRES_URI and CORE_UPDATE_TEST_REDIS_URL") + + from core.config import get_settings + from core.database.postgres_database import PostgresDatabase + from core.models.auth import AuthContext + from core.parser.morphik_parser import MorphikParser + from core.services.ingestion_service import IngestionService + from core.storage.local_storage import LocalStorage + from core.vector_store.pgvector_store import PGVectorStore + from core.workers.ingestion_worker import process_ingestion_job + + settings = get_settings() + monkeypatch.setattr(settings, "ENABLE_COLPALI", False) + monkeypatch.setattr(settings, "MODE", "self_hosted") + queue_name = f"update-test:{uuid.uuid4()}" + redis = await create_pool(RedisSettings.from_dsn(REDIS_URL), default_queue_name=queue_name) + db = PostgresDatabase(POSTGRES_URI) + store = PGVectorStore(POSTGRES_URI) + assert await db.initialize() + assert await store.initialize() + storage = LocalStorage(str(tmp_path)) + embedding = [1.0] + [0.0] * (settings.VECTOR_DIMENSIONS - 1) + model = SimpleNamespace(embed_for_ingestion=AsyncMock(side_effect=lambda chunks: [embedding for _ in chunks])) + parser = MorphikParser(chunk_size=80, chunk_overlap=0) + service = IngestionService(db, store, model, storage, parser) + auth = AuthContext(user_id="update-regression") + ctx = dict(database=db, vector_store=store, embedding_model=model, storage=storage, parser=parser) + ids = [] + + async def ingest(content="ORIGINAL-CONTENT " * 40): + doc = await service.ingest_file_content( + content.encode(), + "policy.txt", + "text/plain", + {"proof": queue_name}, + auth, + redis, + ) + ids.append(doc.external_id) + return doc + + def payload(doc): + return dict( + document_id=doc.external_id, + file_key=doc.storage_info["key"], + bucket=doc.storage_info["bucket"], + original_filename=doc.filename, + content_type=doc.content_type, + auth_dict={"user_id": auth.user_id}, + use_colpali=False, + ingestion_revision=doc.system_metadata.get("ingestion_revision", 0), + ) + + async def drain(): + worker = Worker( + [process_ingestion_job], + redis_pool=redis, + queue_name=queue_name, + ctx=ctx, + burst=True, + poll_delay=0.01, + handle_signals=False, + keep_result=3600, + ) + await asyncio.wait_for(worker.async_run(), timeout=30) + assert worker.jobs_failed == 0 + + async def snapshot(doc): + current = await db.get_document(doc.external_id, auth) + async with db.engine.connect() as conn: + rows = ( + await conn.execute( + text( + "SELECT chunk_number, content FROM vector_embeddings WHERE document_id = :id ORDER BY chunk_number" + ), + {"id": doc.external_id}, + ) + ).all() + return current.model_dump(mode="json"), [tuple(row) for row in rows] + + r = SimpleNamespace( + db=db, + store=store, + redis=redis, + service=service, + auth=auth, + ctx=ctx, + ingest=ingest, + payload=payload, + drain=drain, + snapshot=snapshot, + model=model, + process=process_ingestion_job, + embedding=embedding, + queue_name=queue_name, + ) + try: + yield r + finally: + for document_id in ids: + await store.delete_chunks_by_document_id(document_id) + await db.delete_document(document_id, auth) + for pattern in (f"arq:*:ingest:{document_id}*",): + keys = [key async for key in redis.scan_iter(match=pattern)] + if keys: + await redis.delete(*keys) + await redis.delete(queue_name, queue_name + ":health-check") + await redis.aclose() + await store.engine.dispose() + await db.engine.dispose() + await db._ingestion_lock_engine.dispose() + + +async def test_update_with_retained_result_replaces_all_chunks_and_preserves_identity(runtime): + r = runtime + doc = await r.ingest() + # Reproduce the exact legacy ID collision with a completed real ARQ job. + job_id = f"ingest:{doc.external_id}:0" + await r.redis.zrem(r.queue_name, job_id) + await r.redis.delete("arq:job:" + job_id) + legacy_payload = r.payload(doc) + legacy_payload.pop("ingestion_revision") + legacy = await r.redis.enqueue_job("process_ingestion_job", _job_id=f"ingest:{doc.external_id}", **legacy_payload) + await r.drain() + assert await legacy.status() == JobStatus.complete + before, old_rows = await r.snapshot(doc) + assert len(old_rows) > 1 + assert await r.redis.ttl("arq:result:" + legacy.job_id) > 0 + + corrected = "CORRECTED-CONTENT retention is twenty-one days." + for revision in (1, 2): + updated = await r.service.queue_document_update(doc.external_id, r.auth, r.redis, content=corrected) + assert updated.external_id == doc.external_id + assert updated.metadata == before["metadata"] + assert updated.system_metadata["ingestion_revision"] == revision + new_job = Job(f"ingest:{doc.external_id}:{revision}", r.redis, _queue_name=r.queue_name) + assert await new_job.status() == JobStatus.queued + # Same revision retries remain deduplicated without losing result retention. + assert await r.redis.enqueue_job("process_ingestion_job", _job_id=new_job.job_id, **r.payload(updated)) is None + await r.drain() + current, rows = await r.snapshot(doc) + assert current["system_metadata"]["status"] == "completed" + assert current["system_metadata"]["indexed_revision"] == revision + assert rows == [(0, corrected)] + assert ( + await r.service.storage.download_file(**{k: current["storage_info"][k] for k in ("bucket", "key")}) + == corrected.encode() + ) + retrieved = await r.store.query_similar(r.embedding, 100, doc_ids=[doc.external_id]) + assert [chunk.content for chunk in retrieved] == [corrected] + assert await r.redis.ttl("arq:result:" + legacy.job_id) > 0 + assert await new_job.status() == JobStatus.complete + + +async def test_superseded_and_duplicate_workers_make_no_writes(runtime): + r = runtime + initial = await r.ingest() + old = await r.service.queue_document_update(initial.external_id, r.auth, r.redis, content="OLDER-UPDATE") + latest = await r.service.queue_document_update(initial.external_id, r.auth, r.redis, content="LATEST-UPDATE") + before = await r.snapshot(latest) + for doc in (initial, old): + result = await r.process(r.ctx, **r.payload(doc)) + assert result["status"] == "superseded" + assert await r.snapshot(latest) == before + assert r.model.embed_for_ingestion.await_count == 0 + await r.drain() + finished = await r.snapshot(latest) + assert finished[1] == [(0, "LATEST-UPDATE")] + calls = r.model.embed_for_ingestion.await_count + assert (await r.process(r.ctx, **r.payload(latest)))["status"] == "completed" + assert await r.snapshot(latest) == finished + assert r.model.embed_for_ingestion.await_count == calls + + +@pytest.mark.parametrize("fail_worker", [False, True]) +async def test_active_worker_blocks_updates_until_all_writes_finish(runtime, fail_worker): + r = runtime + doc = await r.ingest() + started, release = asyncio.Event(), asyncio.Event() + + async def delayed_embedding(chunks): + started.set() + await release.wait() + if fail_worker: + raise ValueError("synthetic parse/embedding failure") + return [r.embedding for _ in chunks] + + r.model.embed_for_ingestion.side_effect = delayed_embedding + task = asyncio.create_task(r.process(r.ctx, **r.payload(doc))) + try: + await asyncio.wait_for(started.wait(), timeout=10) + before = await r.snapshot(doc) + with pytest.raises(HTTPException) as error: + await r.service.queue_document_update(doc.external_id, r.auth, r.redis, content="REJECTED") + assert error.value.status_code == 409 + assert await r.snapshot(doc) == before + finally: + release.set() + result = await task + assert result["status"] == ("failed" if fail_worker else "completed") + r.model.embed_for_ingestion.side_effect = lambda chunks: [r.embedding for _ in chunks] + updated = await r.service.queue_document_update(doc.external_id, r.auth, r.redis, content="AFTER-WORKER") + await r.drain() + assert (await r.snapshot(updated))[1] == [(0, "AFTER-WORKER")] + + +async def test_partial_store_retry_cleans_untracked_chunks(runtime, monkeypatch): + r = runtime + doc = await r.ingest("RETRY-CONTENT") + store_embeddings = r.store.store_embeddings + + async def partially_store(*args, **kwargs): + await store_embeddings(*args, **kwargs) + raise ConnectionError("synthetic connection failure after commit") + + monkeypatch.setattr(r.store, "store_embeddings", partially_store) + with pytest.raises(Retry): + await r.process({**r.ctx, "job_try": 1}, **r.payload(doc)) + before, rows = await r.snapshot(doc) + assert rows and not before["chunk_ids"] + monkeypatch.setattr(r.store, "store_embeddings", store_embeddings) + await r.process({**r.ctx, "job_try": 2}, **r.payload(doc)) + current, rows = await r.snapshot(doc) + assert current["system_metadata"]["status"] == "completed" + assert rows == [(0, "RETRY-CONTENT")] + + +@pytest.mark.parametrize("outcome", [None, ConnectionError("synthetic Redis outage")]) +async def test_enqueue_failure_is_reported_and_retry_gets_a_new_revision(runtime, monkeypatch, outcome): + r = runtime + doc = await r.ingest() + await r.drain() + enqueue = r.redis.enqueue_job + monkeypatch.setattr(r.redis, "enqueue_job", AsyncMock(return_value=None, side_effect=outcome)) + with pytest.raises(HTTPException) as error: + await r.service.queue_document_update(doc.external_id, r.auth, r.redis, content="CORRECTION") + assert error.value.status_code == 500 + failed = await r.db.get_document(doc.external_id, r.auth) + assert failed.system_metadata["status"] == "failed" + monkeypatch.setattr(r.redis, "enqueue_job", enqueue) + updated = await r.service.queue_document_update(doc.external_id, r.auth, r.redis, content="CORRECTION") + assert updated.system_metadata["ingestion_revision"] == 2 + await r.drain() + assert (await r.snapshot(updated))[1] == [(0, "CORRECTION")] + + +async def test_file_update_and_cleanup_failure_cannot_publish_stale_chunks(runtime, monkeypatch): + r = runtime + doc = await r.ingest() + await r.drain() + updated = await r.service.queue_document_update( + doc.external_id, + r.auth, + r.redis, + file=UploadFile(BytesIO(b"FILE-CORRECTION"), filename="replacement.txt"), + ) + delete = r.store.delete_chunks_by_document_id + monkeypatch.setattr(r.store, "delete_chunks_by_document_id", AsyncMock(return_value=False)) + result = await r.process(r.ctx, **r.payload(updated)) + assert result["status"] == "failed" + assert (await r.db.get_document(doc.external_id, r.auth)).system_metadata["status"] == "failed" + monkeypatch.setattr(r.store, "delete_chunks_by_document_id", delete) + await r.process({**r.ctx, "job_try": 2}, **r.payload(updated)) + assert (await r.snapshot(updated))[1] == [(0, "FILE-CORRECTION")] + + +async def test_manual_requeue_recovers_retained_result_and_deduplicates_pending_job(runtime, monkeypatch): + from core.models.request import RequeueIngestionRequest + from core.routes import ingest as routes + + r = runtime + monkeypatch.setattr(routes, "ingestion_service", r.service) + doc = await r.ingest("REQUEUE-CONTENT") + await r.drain() + await r.db.update_document(doc.external_id, {"system_metadata": {"status": "failed"}}, r.auth) + request = RequeueIngestionRequest(jobs=[{"external_id": doc.external_id}]) + response = await routes.requeue_ingest_jobs(request, r.auth, r.redis) + assert response.results[0].status == "requeued" + pending = await r.db.get_document(doc.external_id, r.auth) + assert pending.system_metadata["ingestion_revision"] == 1 + before = await r.snapshot(pending) + repeated = await routes.requeue_ingest_jobs(request, r.auth, r.redis) + assert repeated.results[0].status == "already_queued" + assert await r.snapshot(pending) == before + await r.drain() + assert (await r.snapshot(doc))[1] == [(0, "REQUEUE-CONTENT")] + + +async def test_cancelled_worker_releases_lock_and_cannot_resume_over_new_revision(runtime): + r = runtime + doc = await r.ingest("CANCELLED-CONTENT") + started = asyncio.Event() + + async def paused_embedding(chunks): + started.set() + await asyncio.Event().wait() + + r.model.embed_for_ingestion.side_effect = paused_embedding + task = asyncio.create_task(r.process(r.ctx, **r.payload(doc))) + await asyncio.wait_for(started.wait(), timeout=10) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + r.model.embed_for_ingestion.side_effect = lambda chunks: [r.embedding for _ in chunks] + updated = await r.service.queue_document_update(doc.external_id, r.auth, r.redis, content="AFTER-CANCEL") + await r.drain() + before = await r.snapshot(updated) + assert (await r.process({**r.ctx, "job_try": 2}, **r.payload(doc)))["status"] == "superseded" + assert await r.snapshot(updated) == before + assert before[1] == [(0, "AFTER-CANCEL")] + + +async def test_revision_read_error_retries_without_writes(runtime, monkeypatch): + r = runtime + doc = await r.ingest("READ-RETRY") + before = await r.snapshot(doc) + session = r.db.async_session + monkeypatch.setattr(r.db, "async_session", Mock(side_effect=OperationalError("SELECT", {}, Exception("offline")))) + with pytest.raises(Retry): + await r.process(r.ctx, **r.payload(doc)) + assert r.model.embed_for_ingestion.await_count == 0 + monkeypatch.setattr(r.db, "async_session", session) + assert await r.snapshot(doc) == before + await r.drain() + assert (await r.snapshot(doc))[1] == [(0, "READ-RETRY")] diff --git a/core/tests/unit/test_ingestion_service_metadata_update.py b/core/tests/unit/test_ingestion_service_metadata_update.py index cc49e6ff..00b44e00 100644 --- a/core/tests/unit/test_ingestion_service_metadata_update.py +++ b/core/tests/unit/test_ingestion_service_metadata_update.py @@ -2,6 +2,7 @@ import os import sys +from contextlib import asynccontextmanager from pathlib import Path from types import ModuleType, SimpleNamespace @@ -31,6 +32,10 @@ def __init__(self, doc: Document): self.doc = doc self.update_calls = [] + @asynccontextmanager + async def document_ingestion_lock(self, document_id, *, wait=False): + yield True + async def get_document(self, document_id: str, auth: AuthContext): if document_id == self.doc.external_id: return self.doc @@ -294,4 +299,5 @@ async def no_stored_size(*args, **kwargs): assert queued["function_name"] == "process_ingestion_job" assert queued["payload"]["document_id"] == "doc-1" assert queued["payload"]["file_key"] == "ingest_uploads/replacement/report.txt" - assert queued["payload"]["_job_id"] == "ingest:doc-1" + assert queued["payload"]["_job_id"] == "ingest:doc-1:1" + assert queued["payload"]["ingestion_revision"] == 1 diff --git a/core/workers/ingestion_worker.py b/core/workers/ingestion_worker.py index 46c907ba..4c6a2348 100644 --- a/core/workers/ingestion_worker.py +++ b/core/workers/ingestion_worker.py @@ -15,6 +15,7 @@ from arq.worker import Retry from opentelemetry.trace import Status, StatusCode, get_current_span from sqlalchemy import text +from sqlalchemy.exc import SQLAlchemyError from core.config import get_settings from core.database.postgres_database import PostgresDatabase @@ -45,7 +46,8 @@ # library can never break error classification. Used to recognise S3 / turbopuffer # backpressure (throttling, 5xx, connection blips) as transient-and-retryable. try: - from botocore.exceptions import BotoCoreError, ClientError as BotoClientError + from botocore.exceptions import BotoCoreError + from botocore.exceptions import ClientError as BotoClientError except Exception: # noqa: BLE001 BotoCoreError = BotoClientError = None @@ -458,6 +460,65 @@ async def process_ingestion_job( folder_path: Optional[str] = None, folder_leaf: Optional[str] = None, end_user_id: Optional[str] = None, + ingestion_revision: int = 0, +) -> Dict[str, Any]: + """Run only the current revision, holding the document lock through every write. + + Legacy queue messages have revision zero. Checking their storage location too + prevents an old message from reading a file replaced before this rollout. + """ + database = ctx["database"] + auth = AuthContext( + user_id=auth_dict.get("user_id") or auth_dict.get("entity_id", ""), + app_id=auth_dict.get("app_id"), + ) + try: + async with database.document_ingestion_lock(document_id, wait=True): + doc = await database.get_document(document_id, auth, raise_on_error=True) + if ( + doc is None + or int(doc.system_metadata.get("ingestion_revision", 0)) != ingestion_revision + or doc.storage_info.get("key") != file_key + or doc.storage_info.get("bucket") != bucket + ): + logger.info("Skipping superseded ingestion job for %s revision %s", document_id, ingestion_revision) + return {"document_id": document_id, "status": "superseded"} + if doc.system_metadata.get("status") == "completed": + return {"document_id": document_id, "status": "completed"} + return await _process_ingestion_job_locked( + ctx, + document_id, + file_key, + bucket, + original_filename, + content_type, + auth_dict, + use_colpali, + folder_name, + folder_path, + folder_leaf, + end_user_id, + ingestion_revision, + ) + except (SQLAlchemyError, OSError, asyncio.TimeoutError) as exc: + # A failed lock/read is not evidence that this revision was superseded. + raise Retry(defer=30) from exc + + +async def _process_ingestion_job_locked( + ctx: Dict[str, Any], + document_id: str, + file_key: str, + bucket: str, + original_filename: str, + content_type: str, + auth_dict: Dict[str, Any], + use_colpali: bool, + folder_name: Optional[str] = None, + folder_path: Optional[str] = None, + folder_leaf: Optional[str] = None, + end_user_id: Optional[str] = None, + ingestion_revision: int = 0, ) -> Dict[str, Any]: """ Background worker task that processes file ingestion jobs. @@ -1088,16 +1149,13 @@ def _meta_resolver(): # noqa: D401 # duplicate; deletion is by document_id, so it works even though # chunk_ids is only persisted on success. is_retry_attempt = int(ctx.get("job_try") or 1) > 1 - if doc.chunk_ids or is_retry_attempt: + if doc.chunk_ids or is_retry_attempt or ingestion_revision > 0: logger.info( "Cleanup before storing for %s (%s): deleting existing chunks (%d tracked)", document_id, "arq retry" if is_retry_attempt and not doc.chunk_ids else "re-ingestion", len(doc.chunk_ids), ) - deletion_tasks = [] - if hasattr(vector_store, "delete_chunks_by_document_id"): - deletion_tasks.append(vector_store.delete_chunks_by_document_id(document_id, auth.app_id)) # Always try to clean colpali store — the doc may have been ingested # with colpali previously even if this re-ingestion doesn't use it cleanup_colpali_store = colpali_vector_store @@ -1105,23 +1163,25 @@ def _meta_resolver(): # noqa: D401 try: cleanup_colpali_store = await _get_worker_colpali_store(database) except Exception as e: - logger.warning(f"Could not init colpali store for cleanup: {e}") + raise RuntimeError("Could not initialize ColPali store to remove old chunks") from e + deletion_tasks = [] + if hasattr(vector_store, "delete_chunks_by_document_id"): + deletion_tasks.append(vector_store.delete_chunks_by_document_id(document_id, auth.app_id)) if cleanup_colpali_store and hasattr(cleanup_colpali_store, "delete_chunks_by_document_id"): deletion_tasks.append(cleanup_colpali_store.delete_chunks_by_document_id(document_id, auth.app_id)) chunk_v2_store = ctx.get("chunk_v2_store") if chunk_v2_store and auth.app_id and hasattr(chunk_v2_store, "delete_chunks_by_document_id"): deletion_tasks.append(chunk_v2_store.delete_chunks_by_document_id(document_id, auth)) if deletion_tasks: - try: - results = await asyncio.wait_for( - asyncio.gather(*deletion_tasks, return_exceptions=True), - timeout=30, - ) - for i, result in enumerate(results): - if isinstance(result, Exception): - logger.error(f"Error deleting old chunks (task {i}): {result}") - except asyncio.TimeoutError: - logger.error(f"Timeout deleting old chunks for {document_id}, proceeding anyway") + results = await asyncio.wait_for( + asyncio.gather(*deletion_tasks, return_exceptions=True), + timeout=30, + ) + for result in results: + if isinstance(result, BaseException): + raise result + if result is not True: + raise RuntimeError("Failed to delete old chunks; refusing to complete ingestion") # 12. Handle ColPali embeddings chunk_objects_multivector = [] @@ -1375,6 +1435,7 @@ def _meta_resolver(): # noqa: D401 # Final update to mark as completed completion_update = { + "indexed_revision": ingestion_revision, "page_count": final_page_count, "status": "completed", "use_colpali": using_colpali, @@ -1386,9 +1447,11 @@ def _meta_resolver(): # noqa: D401 if no_content_extracted: completion_update["content_extraction_status"] = "no_content_extracted" completion_update["content_extraction_warning"] = content_extraction_warning - await ingestion_service.db.update_document( + completed = await ingestion_service.db.update_document( document_id=document_id, updates={"system_metadata": completion_update}, auth=auth ) + if not completed: + raise RuntimeError("Failed to persist completed ingestion status") # 13. Log successful completion logger.info(f"Successfully completed ingestion for {original_filename}, document ID: {doc.external_id}") diff --git a/docs/document-updates.md b/docs/document-updates.md new file mode 100644 index 00000000..6fa82985 --- /dev/null +++ b/docs/document-updates.md @@ -0,0 +1,57 @@ +# Document update scheduling + +Content updates preserve the document ID, user metadata, and folder association. Each accepted update increments +`system_metadata.ingestion_revision` and queues `ingest::`. Redis can retain the result of +an earlier ingestion without blocking the update. Completion records `system_metadata.indexed_revision`. + +The API and worker take the same PostgreSQL advisory lock for a document. The lock covers initial upload, update +scheduling, manual requeue, and the worker's entire processing attempt, including progress and failure writes. +Workers check the stored revision and source location after taking the lock. Superseded jobs return without +changing files, chunks, or status; duplicate delivery after completion also makes no changes. + +An update attempted during active processing returns HTTP 409 and leaves the document unchanged. The caller can +retry after processing finishes. Updates accepted while earlier jobs are still queued supersede those jobs; +the latest accepted revision is indexed. Repeating an HTTP content update creates another revision, even when +the bytes are identical. There is no client `If-Match` precondition or HTTP idempotency key in this change. + +`completed` is set only after chunk replacement and its document update succeed. Cleanup includes untracked +chunks from interrupted attempts and rejects failed deletions. While processing or failed, the document remains +excluded from normal retrieval, as before. An enqueue exception or unexpected `None` result returns an error +and marks the persisted revision failed instead of reporting successful scheduling. + +`POST /ingest/requeue` reads the current document under the same lock. A queued or active job returns +`already_queued` without resetting status. A finished or missing job gets a new revision, so a retained failed +result cannot prevent recovery. Requeue also recovers a revision left processing if the API exits between the +PostgreSQL write and Redis enqueue; these operations are not a distributed transaction. + +## Deployment + +Deploy the API and ingestion workers together. Drain or stop workers running the old code before accepting +updates through the new API. Older binaries do not take the lock or accept revision arguments. The new worker +accepts legacy queued messages without a revision, treating them as revision zero and checking their source +location. Existing documents need no migration; an absent revision is zero. + +The lock uses one additional PostgreSQL connection per active ingestion or scheduling operation. These +connections use a separate unpooled engine so long-running jobs cannot exhaust the ordinary query pool. + +## Verification + +`core/tests/integration/test_document_update_revisions.py` runs the real ARQ worker, Redis, PostgreSQL/pgvector, +LocalStorage, and text parser with deterministic local embeddings. Point the following variables at disposable +services and run it with the project's installed dependencies and test configuration: + +```bash +export CORE_UPDATE_TEST_POSTGRES_URI='postgresql+asyncpg://morphik:morphik@127.0.0.1:55438/morphik' +export CORE_UPDATE_TEST_REDIS_URL='redis://127.0.0.1:56388/0' +python -m pytest core/tests/integration/test_document_update_revisions.py +``` + +The tests retain an actual completed legacy ARQ result, replace many chunks with one, repeat updates, replay +superseded and completed jobs, reject updates during active workers, retry partial writes, cancel workers, +exercise enqueue/cleanup/read failures, and verify manual requeue. Each test removes only its own data and keys. + +On September 8, 2026, the supplied full runtime proof also passed against the working tree based on `7f72d712`. +It used the real authenticated Core API and worker, real `text-embedding-3-small` embeddings, and isolated +PostgreSQL/pgvector and Redis containers. Corrected text was downloaded and retrieved under the same document +ID and metadata, with no original text in retrieval. The correction survived container recreation and API/worker +restart. This verifies a synthetic standard-text case; it does not cover production images, PDF/OCR, or ColPali. diff --git a/docs/iqor-on-prem.md b/docs/iqor-on-prem.md index 8309d1f0..c7f8dab1 100644 --- a/docs/iqor-on-prem.md +++ b/docs/iqor-on-prem.md @@ -2,6 +2,8 @@ Audit baseline: `origin/main` at `8c51b8d` on 2026-09-02. +The September 8 content-update follow-up is documented in [Document update scheduling](document-updates.md). + This document separates Morphik Core behavior from iQor's MCP wrapper and UI. It also records which findings have an executable test. A code path alone is not counted as a passing deployment check. @@ -14,8 +16,8 @@ executable test. A code path alone is not counted as a passing deployment check. | Start rewrites Compose state and fixed container names collide | Implemented; static verification passed | The API port now uses `MORPHIK_API_PORT`; production services use Compose project-scoped names; static lifecycle tests and `docker compose config` pass. | Morphik Core | | Documents survive PostgreSQL container recreation | Verified in Docker | `scripts/test_postgres_persistence.sh` inserts a document row, recreates PostgreSQL, and checks the original ID and metadata. | Morphik Core / iQor infrastructure | | Text update preserves document identity and existing metadata | Verified in a unit test | `test_queued_text_update_preserves_identity_metadata_and_queues_reindex` passes. | Morphik Core | -| Text update queues changed content for re-indexing and exposes processing status | Partially verified | The unit test proves the replacement object and `process_ingestion_job` payload use the same document ID and that the returned status is `processing`. Existing SDK status tests pass. No end-to-end test in this audit proves the changed text is retrievable after the worker finishes. | Morphik Core | -| Text update prevents lost updates | Fails | There is no content revision precondition. Every update uses ARQ job ID `ingest:{document_id}`. A second update can receive a successful API response while `enqueue_job` returns `None`, and the first queued job may refer to an object the second update deleted. | Morphik Core, then iQor caller adoption | +| Text update queues changed content for re-indexing and exposes processing status | Verified with synthetic standard text | The September 8 runtime proof uses the real API, worker, embeddings, and pgvector. Corrected downloads and retrieval survive container recreation. See [update verification](document-updates.md#verification). | Morphik Core | +| Content updates and ingestion workers cannot overwrite a newer accepted revision | Implemented and integration tested | Persisted revisions give updates distinct job IDs. API and worker share a document lock; active processing returns 409, queued superseded jobs skip all writes. Client editing preconditions remain outside this change. See [update scheduling](document-updates.md). | Morphik Core, then iQor caller adoption | | `min_score` affects retrieval | Fixed and unit verified in this branch | `test_min_score_zero_keeps_zero_and_positive_scores` and `test_min_score_filters_on_the_final_score` pass. | Morphik Core | | Work item 47490 returns five distinct QA backlog items | Not verified | The repository has no iQor corpus, query text, auth token, or captured response. `scripts/verify_iqor_retrieval.sh` captures and validates the response on iQor's deployment. | iQor MCP wrapper / iQor acceptance test | | Default Docker config keeps document and query data on premises | Fails by default | `morphik.docker.toml` selects OpenAI for completion and standard embeddings. Telemetry is enabled unless `TELEMETRY=false`. See the data boundary below. | Joint configuration decision |