diff --git a/backend/app/alembic/versions/087_add_name_to_collection_jobs.py b/backend/app/alembic/versions/087_add_name_to_collection_jobs.py new file mode 100644 index 000000000..94e69ebce --- /dev/null +++ b/backend/app/alembic/versions/087_add_name_to_collection_jobs.py @@ -0,0 +1,42 @@ +"""Add name to collection_jobs with in-flight unique index + +Revision ID: 087 +Revises: 086 +Create Date: 2026-10-28 00:00:00.000000 +""" + +import sqlalchemy as sa +from alembic import op + +revision = "087" +down_revision = "086" +branch_labels = None +depends_on = None + +INDEX_NAME = "uq_collection_jobs_project_id_name_in_flight" + + +def upgrade(): + op.add_column( + "collection_jobs", + sa.Column( + "name", + sa.String(), + nullable=True, + comment="Requested collection name; keeps in-flight CREATE jobs unique per project", + ), + ) + op.create_index( + INDEX_NAME, + "collection_jobs", + ["project_id", "name"], + unique=True, + postgresql_where=sa.text( + "action_type = 'CREATE' AND status IN ('PENDING', 'PROCESSING')" + ), + ) + + +def downgrade(): + op.drop_index(INDEX_NAME, table_name="collection_jobs") + op.drop_column("collection_jobs", "name") diff --git a/backend/app/api/docs/collections/create.md b/backend/app/api/docs/collections/create.md index 6a50b252c..9eacee353 100644 --- a/backend/app/api/docs/collections/create.md +++ b/backend/app/api/docs/collections/create.md @@ -14,3 +14,8 @@ been created, information about the collection will be returned to the user via the callback URL. If a callback URL is not provided, clients can check the `collection job info` endpoint with the `job_id`, to retrieve information about the created collection. + +Requests are retry-safe by `name`: if a create job with the same `name` is still +`PENDING` or `PROCESSING` in the project, the response returns that existing +job's `job_id` instead of creating a second job. Once the job finishes, an +existing collection with that name still returns 409. diff --git a/backend/app/api/routes/collections.py b/backend/app/api/routes/collections.py index 828db79e5..807a42e07 100644 --- a/backend/app/api/routes/collections.py +++ b/backend/app/api/routes/collections.py @@ -20,6 +20,7 @@ CollectionJobStatus, CollectionActionType, CollectionJobCreate, + CollectionJobUpdate, CollectionJobPublic, CollectionJobImmediatePublic, CollectionWithDocsPublic, @@ -113,23 +114,36 @@ def create_collection( unique_documents = list(dict.fromkeys(request.documents)) collection_job_crud = CollectionJobCrud(session, current_user.project_.id) - collection_job = collection_job_crud.create( + collection_job, created = collection_job_crud.create_or_get_in_flight( CollectionJobCreate( action_type=CollectionActionType.CREATE, project_id=current_user.project_.id, status=CollectionJobStatus.PENDING, + name=request.name, docs_num=len(unique_documents), documents=[str(doc_id) for doc_id in unique_documents], ) ) - create_service.start_job( - db=session, - request=request, - collection_job_id=collection_job.id, - project_id=current_user.project_.id, - organization_id=current_user.organization_.id, - ) + if created: + try: + create_service.start_job( + db=session, + request=request, + collection_job_id=collection_job.id, + project_id=current_user.project_.id, + organization_id=current_user.organization_.id, + ) + except Exception as err: + # A PENDING job that never reached the queue would hold its name forever. + collection_job_crud.update( + collection_job.id, + CollectionJobUpdate( + status=CollectionJobStatus.FAILED, + error_message=f"Failed to enqueue collection job: {err}", + ), + ) + raise return APIResponse[CollectionJobImmediatePublic].success_response( CollectionJobImmediatePublic.model_validate(collection_job), diff --git a/backend/app/crud/collection/collection_job.py b/backend/app/crud/collection/collection_job.py index 1cefd39a5..95255ba0a 100644 --- a/backend/app/crud/collection/collection_job.py +++ b/backend/app/crud/collection/collection_job.py @@ -3,10 +3,13 @@ from typing import List from fastapi import HTTPException -from sqlmodel import Session, select +from sqlalchemy.exc import IntegrityError +from sqlmodel import Session, col, select from app.models.collection_job import ( + CollectionActionType, CollectionJob, + CollectionJobStatus, CollectionJobUpdate, CollectionJobCreate, ) @@ -94,3 +97,45 @@ def create(self, collection_job: CollectionJobCreate) -> CollectionJob: raise return job + + def read_in_flight_create(self, name: str) -> CollectionJob | None: + """Return the PENDING/PROCESSING CREATE job holding this name, if any.""" + statement = select(CollectionJob).where( + CollectionJob.project_id == self.project_id, + CollectionJob.name == name, + CollectionJob.action_type == CollectionActionType.CREATE, + col(CollectionJob.status).in_( + [CollectionJobStatus.PENDING, CollectionJobStatus.PROCESSING] + ), + ) + return self.session.exec(statement).first() + + def create_or_get_in_flight( + self, collection_job: CollectionJobCreate + ) -> tuple[CollectionJob, bool]: + """Create a CREATE job, or return the in-flight job with the same name. + + The bool is True only when a new row was inserted, so callers enqueue work once. + """ + if collection_job.name: + existing = self.read_in_flight_create(collection_job.name) + if existing: + return existing, False + + try: + job = CollectionJob(**collection_job.model_dump()) + self.session.add(job) + self.session.commit() + except IntegrityError: + self.session.rollback() + existing = ( + self.read_in_flight_create(collection_job.name) + if collection_job.name + else None + ) + if existing is None: + raise + return existing, False + + self.session.refresh(job) + return job, True diff --git a/backend/app/models/collection_job.py b/backend/app/models/collection_job.py index ea4423d08..dbabd06bd 100644 --- a/backend/app/models/collection_job.py +++ b/backend/app/models/collection_job.py @@ -2,6 +2,7 @@ from enum import Enum from uuid import UUID, uuid4 +from sqlalchemy import Index, text from sqlmodel import JSON, Column, Field, SQLModel, Text from app.core.util import now @@ -24,6 +25,17 @@ class CollectionJob(SQLModel, table=True): """Database model for CollectionJob operations.""" __tablename__ = "collection_jobs" # pyright: ignore[reportAssignmentType] + __table_args__ = ( + Index( + "uq_collection_jobs_project_id_name_in_flight", + "project_id", + "name", + unique=True, + postgresql_where=text( + "action_type = 'CREATE' AND status IN ('PENDING', 'PROCESSING')" + ), + ), + ) id: UUID = Field( default_factory=uuid4, @@ -43,6 +55,14 @@ class CollectionJob(SQLModel, table=True): description="Type of operation", sa_column_kwargs={"comment": "Type of operation (CREATE, DELETE)"}, ) + name: str | None = Field( + default=None, + nullable=True, + description="Requested collection name", + sa_column_kwargs={ + "comment": "Requested collection name; keeps in-flight CREATE jobs unique per project" + }, + ) task_id: str | None = Field( default=None, nullable=True, @@ -153,6 +173,7 @@ def job_updated_at(self) -> datetime: # Request models class CollectionJobCreate(SQLModel): collection_id: UUID | None = None + name: str | None = None status: CollectionJobStatus action_type: CollectionActionType docs_num: int | None = None diff --git a/backend/app/tests/api/routes/collections/test_create_collections.py b/backend/app/tests/api/routes/collections/test_create_collections.py index 5934394b3..9960e8217 100644 --- a/backend/app/tests/api/routes/collections/test_create_collections.py +++ b/backend/app/tests/api/routes/collections/test_create_collections.py @@ -3,11 +3,12 @@ from typing import Any from fastapi.testclient import TestClient -from sqlmodel import Session +from sqlmodel import Session, select from app.core.config import settings from app.tests.utils.auth import TestAuthContext -from app.models import CollectionJobStatus, Document +import pytest +from app.models import Collection, CollectionJob, CollectionJobStatus, Document from app.models.collection import CreationRequest @@ -69,3 +70,83 @@ def test_collection_creation_calls_start_job_and_returns_job( assert kwargs["request"].model_dump(mode="json") == creation_data.model_dump( mode="json" ) + + +@patch("app.api.routes.collections.create_service.start_job") +def test_duplicate_in_flight_name_returns_existing_job( + mock_start_job: Any, + client: TestClient, + user_api_key_header: dict[str, str], + user_api_key: TestAuthContext, + db: Session, +) -> None: + doc = _create_test_document(db, user_api_key.project_id) + payload = {"name": "kb-retry", "documents": [str(doc.id)]} + url = f"{settings.API_V1_STR}/collections" + + first = client.post(url, json=payload, headers=user_api_key_header) + second = client.post(url, json=payload, headers=user_api_key_header) + + assert first.status_code == 200 + assert second.status_code == 200 + assert second.json()["data"]["job_id"] == first.json()["data"]["job_id"] + mock_start_job.assert_called_once() + + +@patch("app.api.routes.collections.create_service.start_job") +def test_existing_collection_name_still_returns_conflict( + mock_start_job: Any, + client: TestClient, + user_api_key_header: dict[str, str], + user_api_key: TestAuthContext, + db: Session, +) -> None: + doc = _create_test_document(db, user_api_key.project_id) + db.add( + Collection( + provider="openai", + knowledge_base_id="vs_existing", + knowledge_base_provider="OpenAI", + name="kb-taken", + project_id=user_api_key.project_id, + ) + ) + db.commit() + + resp = client.post( + f"{settings.API_V1_STR}/collections", + json={"name": "kb-taken", "documents": [str(doc.id)]}, + headers=user_api_key_header, + ) + + assert resp.status_code == 409 + mock_start_job.assert_not_called() + + +@patch("app.api.routes.collections.create_service.start_job") +def test_enqueue_failure_marks_job_failed_and_frees_name( + mock_start_job: Any, + client: TestClient, + user_api_key_header: dict[str, str], + user_api_key: TestAuthContext, + db: Session, +) -> None: + doc = _create_test_document(db, user_api_key.project_id) + payload = {"name": "kb-enqueue", "documents": [str(doc.id)]} + url = f"{settings.API_V1_STR}/collections" + + mock_start_job.side_effect = RuntimeError("broker down") + with pytest.raises(RuntimeError): + client.post(url, json=payload, headers=user_api_key_header) + + failed = db.exec( + select(CollectionJob).where(CollectionJob.name == "kb-enqueue") + ).one() + assert failed.status == CollectionJobStatus.FAILED + assert "broker down" in (failed.error_message or "") + + mock_start_job.side_effect = None + retry = client.post(url, json=payload, headers=user_api_key_header) + + assert retry.status_code == 200 + assert retry.json()["data"]["job_id"] != str(failed.id) diff --git a/backend/app/tests/crud/collections/collection_job/test_collection_jobs.py b/backend/app/tests/crud/collections/collection_job/test_collection_jobs.py index 24d7076a4..4da03beff 100644 --- a/backend/app/tests/crud/collections/collection_job/test_collection_jobs.py +++ b/backend/app/tests/crud/collections/collection_job/test_collection_jobs.py @@ -1,11 +1,19 @@ from typing import Any +from unittest.mock import patch from uuid import uuid4 import pytest from sqlmodel import Session from sqlalchemy.exc import IntegrityError -from app.models import CollectionJob, CollectionJobStatus, CollectionActionType, Project +from app.models import ( + CollectionJob, + CollectionJobCreate, + CollectionJobStatus, + CollectionJobUpdate, + CollectionActionType, + Project, +) from app.crud import CollectionJobCrud from app.core.util import now from app.tests.utils.utils import get_project @@ -126,3 +134,129 @@ def test_create_collection_job_with_invalid_data( with pytest.raises(IntegrityError): collection_job_crud.create(collection_job) + + +def _create_request(project_id: int, name: str | None) -> CollectionJobCreate: + return CollectionJobCreate( + action_type=CollectionActionType.CREATE, + project_id=project_id, + status=CollectionJobStatus.PENDING, + name=name, + ) + + +def test_create_or_get_in_flight_creates_new_job( + db: Session, sample_project: Project +) -> None: + crud = CollectionJobCrud(db, sample_project.id) + + job, created = crud.create_or_get_in_flight( + _create_request(sample_project.id, "kb-new") + ) + + assert created is True + assert job.name == "kb-new" + assert job.status == CollectionJobStatus.PENDING + + +def test_create_or_get_in_flight_returns_existing_job( + db: Session, sample_project: Project +) -> None: + crud = CollectionJobCrud(db, sample_project.id) + first, _ = crud.create_or_get_in_flight( + _create_request(sample_project.id, "kb-dup") + ) + + second, created = crud.create_or_get_in_flight( + _create_request(sample_project.id, "kb-dup") + ) + + assert created is False + assert second.id == first.id + + +def test_create_or_get_in_flight_returns_existing_on_insert_race( + db: Session, sample_project: Project +) -> None: + """Both requests passed the lookup; the unique index rejects the second insert.""" + crud = CollectionJobCrud(db, sample_project.id) + crud.create_or_get_in_flight(_create_request(sample_project.id, "kb-race")) + winner = CollectionJob( + project_id=sample_project.id, + action_type=CollectionActionType.CREATE, + status=CollectionJobStatus.PROCESSING, + name="kb-race", + ) + + with patch.object( + CollectionJobCrud, + "read_in_flight_create", + side_effect=[None, winner], + ): + second, created = crud.create_or_get_in_flight( + _create_request(sample_project.id, "kb-race") + ) + + assert created is False + assert second is winner + + +def test_create_or_get_in_flight_reraises_when_no_in_flight_job( + db: Session, sample_project: Project +) -> None: + crud = CollectionJobCrud(db, sample_project.id) + crud.create_or_get_in_flight(_create_request(sample_project.id, "kb-gone")) + + with patch.object( + CollectionJobCrud, "read_in_flight_create", side_effect=[None, None] + ): + with pytest.raises(IntegrityError): + crud.create_or_get_in_flight(_create_request(sample_project.id, "kb-gone")) + + +def test_create_or_get_in_flight_allows_name_after_job_finished( + db: Session, sample_project: Project +) -> None: + crud = CollectionJobCrud(db, sample_project.id) + first, _ = crud.create_or_get_in_flight( + _create_request(sample_project.id, "kb-reuse") + ) + crud.update(first.id, CollectionJobUpdate(status=CollectionJobStatus.FAILED)) + + second, created = crud.create_or_get_in_flight( + _create_request(sample_project.id, "kb-reuse") + ) + + assert created is True + assert second.id != first.id + + +def test_create_or_get_in_flight_without_name_always_creates( + db: Session, sample_project: Project +) -> None: + crud = CollectionJobCrud(db, sample_project.id) + + first, first_created = crud.create_or_get_in_flight( + _create_request(sample_project.id, None) + ) + second, second_created = crud.create_or_get_in_flight( + _create_request(sample_project.id, None) + ) + + assert first_created and second_created + assert first.id != second.id + + +def test_read_in_flight_create_ignores_other_action_and_status( + db: Session, sample_project: Project +) -> None: + crud = CollectionJobCrud(db, sample_project.id) + done = CollectionJob( + project_id=sample_project.id, + action_type=CollectionActionType.CREATE, + status=CollectionJobStatus.SUCCESSFUL, + name="kb-read", + ) + crud.create(done) + + assert crud.read_in_flight_create("kb-read") is None diff --git a/docs/wiki/modules/knowledge-base.md b/docs/wiki/modules/knowledge-base.md index b585165a4..fb51dd367 100644 --- a/docs/wiki/modules/knowledge-base.md +++ b/docs/wiki/modules/knowledge-base.md @@ -55,3 +55,4 @@ All paths relative to `backend/app/`. - `documents_uploaded` is deduped on append — a task re-run (e.g. redelivered after worker loss, since `acks_late` is on) re-adds its own IDs, and `DocumentCrud.read_each` raises when duplicates collapse in its `IN` clause. - Do not "fix" batch timing with a self-requeueing continuation task — a task that returns normally is not redelivered, and a lost continuation strands the job in `PROCESSING`, which nothing monitors or recovers. Deep dive §11.5. - Uploads happen per batch in `execute_batch_job`, not in setup; setup only plans batches and creates the vector store. +- Collection create is retry-safe by name while a job is in flight: `collection_jobs.name` + partial unique index `uq_collection_jobs_project_id_name_in_flight` (`project_id, name` where CREATE and PENDING/PROCESSING). `CollectionJobCrud.create_or_get_in_flight` returns the existing job (and `created=False`, so `POST /collections` skips `start_job`) instead of a duplicate job and vector store. The `collection` row only exists after the worker finishes, so `ensure_unique_name` alone cannot catch retries. A completed collection still gets 409. Unnamed requests have no key and are not covered. A job stuck in `PROCESSING` keeps its name reserved.