diff --git a/BROKER_PREREQUISITES_ENHANCEMENT.md b/BROKER_PREREQUISITES_ENHANCEMENT.md new file mode 100644 index 0000000..1dc7afc --- /dev/null +++ b/BROKER_PREREQUISITES_ENHANCEMENT.md @@ -0,0 +1,115 @@ +# Broker Prerequisites Enhancement + +## Problem +The broker API was returning null values for `prerequisites`, `validation_hints`, and `file_metadata` fields because: +1. Missing accession data in submission payloads +2. No way to distinguish between "not required" vs "required but not yet submitted" + +## Solution +Enhanced the broker contract to include **existing** accessions (resolved from `accession_registry`), allowing clients to decide whether dependency state is complete enough to submit. + +### Key Architecture Understanding + +**Crucial clarification**: Accessions are stored in different places: +- **Existing accessions**: Stored in database row fields (`project_accession`, `sample_accession`, etc.) and `accession_registry` table +- **prepared_payload**: Contains data for ENA submission, NOT existing accessions +- **Required accessions**: Specified in payload using `expected_*_accession` fields for dependencies not yet submitted + +### Changes Made + +#### 1. BrokerPrerequisites Schema +```python +class BrokerPrerequisites(BaseModel): + # Existing accessions (from database row fields) + project_accession: Optional[str] = None + sample_accession: Optional[str] = None + experiment_accession: Optional[str] = None + run_accession: Optional[str] = None + study_accession: Optional[str] = None + analysis_accession: Optional[str] = None +``` + +#### 2. Enhanced Prerequisite Extraction Logic +- **Existing accessions**: Looked up from `accession_registry` via FK relationships on submission tables + +#### 3. Updated Test Cases +- Tests now verify both existing and required accessions +- Demonstrates mixed states where some dependencies exist and others don't + +### Usage Examples + +#### Sample Submission with Existing Project Accession +**Database state**: `accession_registry` contains a project accession for the sample's `project_id` +```json +{ + "alias": "sample-1", + "requires_project_accession": true +} +``` +Result: +```json +{ + "prerequisites": { + "project_accession": "PRJ123456" // ✅ Resolved from accession_registry + } +} +``` + +#### Experiment with Missing Study Accession +**Database state**: `accession_registry` does not contain a project accession for the experiment's `project_id` +```json +{ + "alias": "experiment-1", + "requires_study_accession": true, + "expected_study_accession": "PRJ123456" // Required but not submitted +} +``` +Result: +```json +{ + "prerequisites": { + "sample_accession": "SAMEA123456", // ✅ Resolved from accession_registry + "study_accession": null // ❌ Missing from registry + } +} +``` + +#### Run with Missing Experiment Accession +**Database state**: `accession_registry` does not contain an experiment accession for the run's `experiment_id` +```json +{ + "alias": "run-1", + "file_name": "reads.fastq.gz", + "file_format": "fastq", + "expected_experiment_accession": "ERX123456" // Required but not submitted +} +``` +Result: +```json +{ + "prerequisites": { + "experiment_accession": null // ❌ Missing from registry + }, + "files": [ + {"filename": "reads.fastq.gz", "filetype": "fastq"} + ] +} +``` + +### Data Flow Summary + +1. **prepared_payload**: Contains ENA submission data (metadata, files, etc.) +2. **Database row fields**: Store actual accessions once submitted (`*_accession` fields) +3. **Broker response**: Shows resolved accessions (or `null`), and clients decide completeness + +### Client Benefits +1. **Clear dependency visibility**: See what exists vs what is missing +2. **Single source of truth**: Accessions resolved from `accession_registry` + +### Payload Fields for Required Accessions +- `expected_*_accession`: Placeholder for the accession that will be assigned +- `requires_*_accession`: Boolean flag indicating the dependency is required +- `requires_project_accession`: For samples +- `requires_study_accession`: For experiments (maps to `required_project_accession` in broker prerequisites) + +This enhancement maintains backward compatibility while providing much richer dependency information to broker clients, with the correct understanding of where different types of data are stored. diff --git a/README.md b/README.md index 66bc37a..11ffaf3 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ A dedicated broker workflow enables integration with external submission pipelin - Submission workflow endpoints (sample-submissions, experiment-submissions, read-submissions) - Broker endpoints to support external submission pipelines: - Claim drafts and obtain a lease: `/api/v1/broker/organisms/{organism_key}/claim` + - ENA broker contract endpoints: `/api/v1/broker/claims/ready`, `/api/v1/broker/claims/entity`, `/api/v1/broker/validation`, `/api/v1/broker/reports/{attempt_id}` - Renew lease, finalise, and report results: `/api/v1/broker/attempts/{attempt_id}/...` - Attempt listing and summaries for dashboard views - Bulk import endpoints for organisms, samples, and experiments @@ -236,6 +237,7 @@ pyproject.toml, uv.lock, docker-compose.yml, Dockerfile, schema.sql, scripts/, d -d '{"samples": [], "experiments": [], "reads": [], "projects": []}' ``` + For the flat ENA broker contract used by Canopy, see [docs/ena_broker_contract.md](docs/ena_broker_contract.md). For a deeper overview of attempt leasing and statuses, see the `broker` endpoints in `app/api/v1/endpoints/broker.py` and the interactive docs. ## Bulk Import API diff --git a/alembic/versions/0006_consolidated_submission_schema_refactor.py b/alembic/versions/0006_consolidated_submission_schema_refactor.py new file mode 100644 index 0000000..2648e36 --- /dev/null +++ b/alembic/versions/0006_consolidated_submission_schema_refactor.py @@ -0,0 +1,136 @@ +"""Consolidated submission schema refactor. + +This migration consolidates the following changes: +- Add project_id to sample_submission (NOT NULL, FK to project) +- Add project_id to experiment (NOT NULL, FK to project) +- Drop project_id from read_submission (derivable via experiment) +- Drop sample_id and project_id from experiment_submission (derivable via experiment) + +Note: No backfill logic included - database will be repopulated after merge. + +Revision ID: 0006_consolidated_refactor +Revises: 0005_org_sci_name_nullable +Create Date: 2026-04-22 00:00:00.000000 +""" + +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "0006_consolidated_refactor" +down_revision = "0005_org_sci_name_nullable" +branch_labels = None +depends_on = None + + +def upgrade(): + # 1) Add sample_submission.project_id (NOT NULL, FK to project) + op.add_column( + "sample_submission", + sa.Column("project_id", postgresql.UUID(as_uuid=True), nullable=False), + ) + op.create_foreign_key( + "fk_sample_submission_project_id", + "sample_submission", + "project", + ["project_id"], + ["id"], + ) + op.create_index( + "idx_sample_submission_project_id", + "sample_submission", + ["project_id"], + unique=False, + ) + + # 2) Add experiment.project_id (NOT NULL, FK to project with CASCADE) + op.add_column( + "experiment", + sa.Column("project_id", postgresql.UUID(as_uuid=True), nullable=False), + ) + op.create_foreign_key( + "fk_experiment_project_id", + "experiment", + "project", + ["project_id"], + ["id"], + ondelete="CASCADE", + ) + op.create_index( + "idx_experiment_project_id", + "experiment", + ["project_id"], + unique=False, + ) + + # 3) Drop read_submission.project_id (derivable via experiment) + op.execute( + "ALTER TABLE read_submission DROP CONSTRAINT IF EXISTS read_submission_project_id_fkey" + ) + op.execute( + "ALTER TABLE read_submission DROP CONSTRAINT IF EXISTS read_submission_project_id_fkey1" + ) + op.execute("ALTER TABLE read_submission DROP COLUMN IF EXISTS project_id") + + # 4) Drop experiment_submission.sample_id and project_id (derivable via experiment) + op.execute( + "ALTER TABLE experiment_submission DROP CONSTRAINT IF EXISTS experiment_submission_sample_id_fkey" + ) + op.execute( + "ALTER TABLE experiment_submission DROP CONSTRAINT IF EXISTS experiment_submission_project_id_fkey" + ) + op.execute("ALTER TABLE experiment_submission DROP COLUMN IF EXISTS sample_id") + op.execute("ALTER TABLE experiment_submission DROP COLUMN IF EXISTS project_id") + + +def downgrade(): + # Reverse order of upgrade + + # 4) Re-add experiment_submission.sample_id and project_id + op.add_column( + "experiment_submission", + sa.Column("sample_id", postgresql.UUID(as_uuid=True), nullable=True), + ) + op.add_column( + "experiment_submission", + sa.Column("project_id", postgresql.UUID(as_uuid=True), nullable=True), + ) + op.create_foreign_key( + "experiment_submission_sample_id_fkey", + "experiment_submission", + "sample", + ["sample_id"], + ["id"], + ) + op.create_foreign_key( + "experiment_submission_project_id_fkey", + "experiment_submission", + "project", + ["project_id"], + ["id"], + ) + + # 3) Re-add read_submission.project_id + op.add_column( + "read_submission", + sa.Column("project_id", postgresql.UUID(as_uuid=True), nullable=True), + ) + op.create_foreign_key( + "read_submission_project_id_fkey", + "read_submission", + "project", + ["project_id"], + ["id"], + ) + + # 2) Drop experiment.project_id + op.drop_index("idx_experiment_project_id", table_name="experiment") + op.drop_constraint("fk_experiment_project_id", "experiment", type_="foreignkey") + op.drop_column("experiment", "project_id") + + # 1) Drop sample_submission.project_id + op.drop_index("idx_sample_submission_project_id", table_name="sample_submission") + op.drop_constraint("fk_sample_submission_project_id", "sample_submission", type_="foreignkey") + op.drop_column("sample_submission", "project_id") diff --git a/app/api/v1/endpoints/broker.py b/app/api/v1/endpoints/broker.py index 1c0af98..2d476d6 100644 --- a/app/api/v1/endpoints/broker.py +++ b/app/api/v1/endpoints/broker.py @@ -1,3 +1,4 @@ +import logging from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Optional from uuid import UUID, uuid4 @@ -6,6 +7,7 @@ from pydantic import BaseModel, Field from sqlalchemy import func, or_ from sqlalchemy.dialects.postgresql import insert +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from app.core.dependencies import get_current_active_user, get_db, has_role @@ -18,8 +20,29 @@ from app.models.read import Read, ReadSubmission from app.models.sample import Sample, SampleSubmission from app.models.user import User +from app.schemas.broker_contract import ( + BrokerBatchClaimRequest, + BrokerBatchClaimResponse, + BrokerClaimEntity, + BrokerEntityType, + BrokerFileMetadata, + BrokerPrerequisites, + BrokerReadyClaimRequest, + BrokerReadyClaimResponse, + BrokerReportRequest, + BrokerReportResponse, + BrokerTargetedClaimRequest, + BrokerTargetedClaimResponse, + BrokerValidationIssue, + BrokerValidationRequest, + BrokerValidationResponse, +) router = APIRouter() +CLAIMABLE_SUBMISSION_STATES = ("draft", "ready") + + +logger = logging.getLogger(__name__) # ---------- Pydantic models for request/response ---------- @@ -118,6 +141,952 @@ class ReportResult(BaseModel): updated_counts: Dict[str, int] +def _normalise_tax_id_value(value: str | int) -> int: + try: + return int(str(value)) + except ValueError as exc: + raise HTTPException(status_code=422, detail="tax_id must be an integer value") from exc + + +def _coerce_entity_type(entity_type: BrokerEntityType | str) -> BrokerEntityType: + if isinstance(entity_type, BrokerEntityType): + return entity_type + return BrokerEntityType(entity_type) + + +def _prerequisites_to_dict(prerequisites: Optional[BrokerPrerequisites]) -> Dict[str, str]: + if prerequisites is None: + return {} + return {key: value for key, value in prerequisites.model_dump().items() if value is not None} + + +def _first_present_value(*values: Any) -> Optional[str]: + for value in values: + if value is not None: + return str(value) + return None + + +def _lookup_scientific_name(db: Session, *, tax_id: str | int) -> Optional[str]: + try: + tax_id_int = int(tax_id) + except Exception: + return None + return db.query(Organism.scientific_name).filter(Organism.tax_id == tax_id_int).scalar() + + +def _create_new_draft_submission_after_rejection( + db: Session, + *, + entity_type: BrokerEntityType, + row: Any, +) -> None: + """Create a new draft submission row after a rejection so the entity can be re-claimed. + + Copies the original row data (where applicable) but clears response_payload and resets status. + """ + + model_cls = None + entity_fk_field = None + if entity_type == BrokerEntityType.SAMPLE: + model_cls = SampleSubmission + entity_fk_field = "sample_id" + elif entity_type == BrokerEntityType.EXPERIMENT: + model_cls = ExperimentSubmission + entity_fk_field = "experiment_id" + elif entity_type == BrokerEntityType.RUN: + model_cls = ReadSubmission + entity_fk_field = "read_id" + elif entity_type == BrokerEntityType.PROJECT: + model_cls = ProjectSubmission + entity_fk_field = "project_id" + + if model_cls is None or entity_fk_field is None: + return + + entity_id = getattr(row, entity_fk_field, None) + prepared_payload = getattr(row, "prepared_payload", None) or None + authority = getattr(row, "authority", "ENA") + entity_type_const = getattr(row, "entity_type_const", entity_type.value) + + if entity_id is None or prepared_payload is None: + return + + kwargs: Dict[str, Any] = { + entity_fk_field: entity_id, + "authority": authority, + "entity_type_const": entity_type_const, + "prepared_payload": prepared_payload, + "status": "draft", + "response_payload": None, + "accession": getattr(row, "accession", None), + "submitted_at": None, + "attempt_id": None, + "finalised_attempt_id": None, + "lock_acquired_at": None, + "lock_expires_at": None, + } + + # Preserve optional extra fields where present + model_column_names = set(getattr(model_cls, "__table__").columns.keys()) + if hasattr(row, "project_id") and "project_id" in model_column_names: + kwargs["project_id"] = getattr(row, "project_id", None) + if hasattr(row, "biosample_accession") and "biosample_accession" in model_column_names: + kwargs["biosample_accession"] = getattr(row, "biosample_accession", None) + + # Drop keys that don't exist on the model (helps keep this generic) + kwargs = {k: v for k, v in kwargs.items() if k in model_column_names} + + db.add(model_cls(**kwargs)) + + +def _derive_claim_title( + db: Session, *, entity_type: BrokerEntityType, entity_id: UUID, tax_id: str | int +) -> Optional[str]: + try: + tax_id_int = int(tax_id) + except Exception: + return None + + scientific_name = ( + db.query(Organism.scientific_name).filter(Organism.tax_id == tax_id_int).scalar() + ) + if not scientific_name: + return None + + if entity_type == BrokerEntityType.EXPERIMENT: + bpa_package_id = ( + db.query(Experiment.bpa_package_id).filter(Experiment.id == entity_id).scalar() + ) + if not bpa_package_id: + return None + return f"Bioplatforms Australia dataset {bpa_package_id} for {scientific_name}" + + if entity_type == BrokerEntityType.SAMPLE: + sample_row = ( + db.query(Sample.kind, Sample.specimen_id, Sample.bpa_sample_id) + .filter(Sample.id == entity_id) + .first() + ) + if not sample_row: + return None + kind, specimen_id, bpa_sample_id = sample_row + if kind == "specimen": + if not specimen_id: + return None + return f"Specimen {specimen_id} for {scientific_name}" + if kind == "derived": + if not bpa_sample_id: + return None + return f"Sample {bpa_sample_id} for {scientific_name}" + + return None + + +def _as_payload(payload: Any) -> Dict[str, Any]: + if isinstance(payload, dict): + return payload + return {} + + +def _get_accession_for_entity( + db: Session, entity_type: str, entity_id: UUID, authority: str +) -> Optional[str]: + """Lookup accession from accession_registry for a given entity.""" + if not entity_id: + return None + return ( + db.query(AccessionRegistry.accession) + .filter( + AccessionRegistry.entity_type == entity_type, + AccessionRegistry.entity_id == entity_id, + AccessionRegistry.authority == authority, + ) + .scalar() + ) + + +def _extract_broker_prerequisites( + db: Session, entity_type: BrokerEntityType | str, prepared_payload: Dict[str, Any], row: Any +) -> BrokerPrerequisites: + entity_type = _coerce_entity_type(entity_type) + authority = getattr(row, "authority", "ENA") + + # Lookup existing accessions from accession_registry via FK relationships + project_id = getattr(row, "project_id", None) + sample_id = getattr(row, "sample_id", None) + experiment_id = getattr(row, "experiment_id", None) + + # Runs and reads may not carry project_id directly; derive via experiment. + if project_id is None and experiment_id is not None: + project_id = db.query(Experiment.project_id).filter(Experiment.id == experiment_id).scalar() + + # Experiment submissions may not carry sample_id directly; derive via experiment. + if sample_id is None and experiment_id is not None: + sample_id = db.query(Experiment.sample_id).filter(Experiment.id == experiment_id).scalar() + + project_accession = _get_accession_for_entity(db, "project", project_id, authority) + sample_accession = _get_accession_for_entity(db, "sample", sample_id, authority) + experiment_accession = _get_accession_for_entity(db, "experiment", experiment_id, authority) + run_accession = getattr(row, "accession", None) if entity_type == BrokerEntityType.RUN else None + + if entity_type == BrokerEntityType.SAMPLE: + project_accession = None + + study_accession = project_accession + analysis_accession = prepared_payload.get("analysis_accession") # This might only be in payload + + return BrokerPrerequisites( + # Existing accessions (from database) + project_accession=project_accession, + sample_accession=sample_accession, + experiment_accession=experiment_accession, + run_accession=run_accession, + study_accession=study_accession, + analysis_accession=analysis_accession, + ) + + +def _extract_run_files(prepared_payload: Dict[str, Any]) -> Optional[List[BrokerFileMetadata]]: + if isinstance(prepared_payload.get("files"), list): + files: List[BrokerFileMetadata] = [] + for item in prepared_payload["files"]: + if isinstance(item, dict) and item.get("filename") and item.get("filetype"): + files.append( + BrokerFileMetadata(filename=item["filename"], filetype=item["filetype"]) + ) + if files: + return files + + filename = prepared_payload.get("filename") or prepared_payload.get("file_name") + filetype = prepared_payload.get("filetype") or prepared_payload.get("file_format") + if filename and filetype: + return [BrokerFileMetadata(filename=filename, filetype=filetype)] + return None + + +def _build_contract_entity( + db: Session, entity_type: BrokerEntityType | str, entity_id: UUID, tax_id: str | int, row: Any +) -> BrokerClaimEntity: + entity_type = _coerce_entity_type(entity_type) + prepared_payload = _as_payload(getattr(row, "prepared_payload", None)) + prerequisites = None + if entity_type != BrokerEntityType.PROJECT: + prerequisites = _extract_broker_prerequisites(db, entity_type, prepared_payload, row) + files = _extract_run_files(prepared_payload) if entity_type == BrokerEntityType.RUN else None + scientific_name = _lookup_scientific_name(db, tax_id=tax_id) + + if entity_type in (BrokerEntityType.EXPERIMENT, BrokerEntityType.SAMPLE): + title = _derive_claim_title(db, entity_type=entity_type, entity_id=entity_id, tax_id=tax_id) + if title is not None: + prepared_payload["title"] = title + + return BrokerClaimEntity( + type=entity_type, + id=entity_id, + tax_id=str(tax_id), + scientific_name=scientific_name, + payload=prepared_payload or None, + prerequisites=prerequisites if _prerequisites_to_dict(prerequisites) else None, + files=files, + ) + + +def _claim_submission_row( + db: Session, + *, + attempt: SubmissionAttempt, + row: Any, + entity_type: str, + now: datetime, +) -> None: + row.status = "submitting" + row.attempt_id = attempt.id + row.lock_acquired_at = now + row.lock_expires_at = attempt.lock_expires_at + db.add( + SubmissionEvent( + attempt_id=attempt.id, + entity_type=entity_type, + submission_id=row.id, + action="claimed", + ) + ) + + +def _create_contract_attempt( + db: Session, *, organism_key: Optional[str], lease_minutes: int = 30 +) -> SubmissionAttempt: + now = datetime.now(timezone.utc) + attempt = SubmissionAttempt( + organism_key=organism_key, + status="processing", + lock_acquired_at=now, + lock_expires_at=now + timedelta(minutes=lease_minutes), + ) + db.add(attempt) + db.flush() + return attempt + + +def _get_organism_by_tax_id(db: Session, tax_id: int) -> Optional[Organism]: + return db.query(Organism).filter(Organism.tax_id == tax_id).first() + + +def _query_ready_project_submissions(db: Session, tax_id: int) -> List[ProjectSubmission]: + return ( + db.query(ProjectSubmission) + .join(Project, ProjectSubmission.project_id == Project.id) + .join(Organism, Project.organism_key == Organism.grouping_key) + .filter( + Organism.tax_id == tax_id, + ProjectSubmission.status.in_(CLAIMABLE_SUBMISSION_STATES), + ) + .order_by(ProjectSubmission.created_at.desc()) + .with_for_update(skip_locked=True) + .all() + ) + + +def _query_ready_sample_submissions(db: Session, tax_id: int) -> List[SampleSubmission]: + return ( + db.query(SampleSubmission) + .join(Sample, SampleSubmission.sample_id == Sample.id) + .join(Organism, Sample.organism_key == Organism.grouping_key) + .filter( + Organism.tax_id == tax_id, + SampleSubmission.status.in_(CLAIMABLE_SUBMISSION_STATES), + ) + .order_by(SampleSubmission.created_at.desc()) + .with_for_update(skip_locked=True) + .all() + ) + + +def _query_ready_experiment_submissions(db: Session, tax_id: int) -> List[ExperimentSubmission]: + return ( + db.query(ExperimentSubmission) + .join(Experiment, ExperimentSubmission.experiment_id == Experiment.id) + .join(Sample, Experiment.sample_id == Sample.id) + .join(Organism, Sample.organism_key == Organism.grouping_key) + .filter( + Organism.tax_id == tax_id, + ExperimentSubmission.status.in_(CLAIMABLE_SUBMISSION_STATES), + ) + .order_by(ExperimentSubmission.created_at.desc()) + .with_for_update(skip_locked=True) + .all() + ) + + +def _query_ready_run_submissions(db: Session, tax_id: int) -> List[ReadSubmission]: + return ( + db.query(ReadSubmission) + .join(Read, ReadSubmission.read_id == Read.id) + .join(Experiment, Read.experiment_id == Experiment.id) + .join(Sample, Experiment.sample_id == Sample.id) + .join(Organism, Sample.organism_key == Organism.grouping_key) + .filter( + Organism.tax_id == tax_id, + ReadSubmission.status.in_(CLAIMABLE_SUBMISSION_STATES), + ) + .order_by(ReadSubmission.created_at.desc()) + .with_for_update(skip_locked=True) + .all() + ) + + +def _latest_per_entity(rows: List[Any], entity_attr: str) -> List[Any]: + latest_rows: List[Any] = [] + seen_entity_ids: set[UUID] = set() + for row in rows: + entity_id = getattr(row, entity_attr, None) + if entity_id is None or entity_id in seen_entity_ids: + continue + seen_entity_ids.add(entity_id) + latest_rows.append(row) + return latest_rows + + +def _find_latest_claimable_entity_submission( + db: Session, entity_type: BrokerEntityType | str, entity_id: UUID +) -> Optional[Any]: + entity_type = _coerce_entity_type(entity_type) + if entity_type == BrokerEntityType.PROJECT: + return ( + db.query(ProjectSubmission) + .filter( + ProjectSubmission.project_id == entity_id, + ProjectSubmission.status.in_(CLAIMABLE_SUBMISSION_STATES), + ) + .order_by(ProjectSubmission.created_at.desc()) + .with_for_update(skip_locked=True) + .first() + ) + if entity_type == BrokerEntityType.SAMPLE: + return ( + db.query(SampleSubmission) + .filter( + SampleSubmission.sample_id == entity_id, + SampleSubmission.status.in_(CLAIMABLE_SUBMISSION_STATES), + ) + .order_by(SampleSubmission.created_at.desc()) + .with_for_update(skip_locked=True) + .first() + ) + if entity_type == BrokerEntityType.EXPERIMENT: + return ( + db.query(ExperimentSubmission) + .filter( + ExperimentSubmission.experiment_id == entity_id, + ExperimentSubmission.status.in_(CLAIMABLE_SUBMISSION_STATES), + ) + .order_by(ExperimentSubmission.created_at.desc()) + .with_for_update(skip_locked=True) + .first() + ) + return ( + db.query(ReadSubmission) + .filter( + ReadSubmission.read_id == entity_id, + ReadSubmission.status.in_(CLAIMABLE_SUBMISSION_STATES), + ) + .order_by(ReadSubmission.created_at.desc()) + .with_for_update(skip_locked=True) + .first() + ) + + +def _find_latest_submission_for_validation( + db: Session, entity_type: BrokerEntityType | str, entity_id: UUID +) -> Optional[Any]: + entity_type = _coerce_entity_type(entity_type) + if entity_type == BrokerEntityType.PROJECT: + return ( + db.query(ProjectSubmission) + .filter(ProjectSubmission.project_id == entity_id) + .order_by(ProjectSubmission.created_at.desc()) + .first() + ) + if entity_type == BrokerEntityType.SAMPLE: + return ( + db.query(SampleSubmission) + .filter(SampleSubmission.sample_id == entity_id) + .order_by(SampleSubmission.created_at.desc()) + .first() + ) + if entity_type == BrokerEntityType.EXPERIMENT: + return ( + db.query(ExperimentSubmission) + .filter(ExperimentSubmission.experiment_id == entity_id) + .order_by(ExperimentSubmission.created_at.desc()) + .first() + ) + return ( + db.query(ReadSubmission) + .filter(ReadSubmission.read_id == entity_id) + .order_by(ReadSubmission.created_at.desc()) + .first() + ) + + +def _lookup_taxonomy_for_entity( + db: Session, entity_type: BrokerEntityType | str, entity_id: UUID +) -> tuple[Optional[str], Optional[int]]: + entity_type = _coerce_entity_type(entity_type) + if entity_type == BrokerEntityType.PROJECT: + row = ( + db.query(Project.organism_key, Organism.tax_id) + .join(Organism, Project.organism_key == Organism.grouping_key) + .filter(Project.id == entity_id) + .first() + ) + elif entity_type == BrokerEntityType.SAMPLE: + row = ( + db.query(Sample.organism_key, Organism.tax_id) + .join(Organism, Sample.organism_key == Organism.grouping_key) + .filter(Sample.id == entity_id) + .first() + ) + elif entity_type == BrokerEntityType.EXPERIMENT: + row = ( + db.query(Sample.organism_key, Organism.tax_id) + .join(Experiment, Experiment.sample_id == Sample.id) + .join(Organism, Sample.organism_key == Organism.grouping_key) + .filter(Experiment.id == entity_id) + .first() + ) + else: + row = ( + db.query(Sample.organism_key, Organism.tax_id) + .join(Experiment, Experiment.sample_id == Sample.id) + .join(Read, Read.experiment_id == Experiment.id) + .join(Organism, Sample.organism_key == Organism.grouping_key) + .filter(Read.id == entity_id) + .first() + ) + + if not row: + return None, None + return row[0], row[1] + + +def _merge_prerequisites( + stored: BrokerPrerequisites, overrides: BrokerPrerequisites +) -> Dict[str, str]: + merged = _prerequisites_to_dict(stored) + for key, value in overrides.model_dump().items(): + if value is not None: + merged[key] = value + return merged + + +def _validate_contract_entity( + claimed_entity: BrokerClaimEntity, overrides: BrokerPrerequisites +) -> BrokerValidationResponse: + entity_type = _coerce_entity_type(claimed_entity.type) + stored_prerequisites = claimed_entity.prerequisites or BrokerPrerequisites() + resolved_prerequisites = _merge_prerequisites(stored_prerequisites, overrides) + issues: List[BrokerValidationIssue] = [] + + if entity_type == BrokerEntityType.SAMPLE: + if not claimed_entity.payload: + issues.append( + BrokerValidationIssue(field="payload", message="sample payload is required") + ) + elif entity_type == BrokerEntityType.EXPERIMENT: + if not claimed_entity.payload: + issues.append( + BrokerValidationIssue(field="payload", message="experiment payload is required") + ) + if not resolved_prerequisites.get("sample_accession"): + issues.append( + BrokerValidationIssue( + field="sample_accession", + message="sample accession is required", + ) + ) + if not resolved_prerequisites.get("project_accession"): + issues.append( + BrokerValidationIssue( + field="project_accession", + message="project accession is required", + ) + ) + elif entity_type == BrokerEntityType.RUN: + if not claimed_entity.payload: + issues.append(BrokerValidationIssue(field="payload", message="run payload is required")) + if not resolved_prerequisites.get("experiment_accession"): + issues.append( + BrokerValidationIssue( + field="experiment_accession", + message="experiment accession is required", + ) + ) + if not (claimed_entity.files or claimed_entity.file_metadata): + issues.append( + BrokerValidationIssue(field="files", message="run file metadata is required") + ) + + return BrokerValidationResponse( + entity_type=claimed_entity.type, + entity_id=claimed_entity.id, + valid=not issues, + issues=issues, + resolved_prerequisites=resolved_prerequisites, + ) + + +def _map_report_state_to_submission_status(state: str) -> str: + lowered = state.lower() + if lowered in {"completed", "accepted", "success"}: + return "accepted" + if lowered in {"failed", "rejected", "error"}: + return "rejected" + if lowered in {"submitting", "processing"}: + return "submitting" + raise HTTPException(status_code=422, detail=f"Unsupported report state '{state}'") + + +def _find_submission_for_attempt( + db: Session, entity_type: BrokerEntityType | str, entity_id: UUID, attempt_id: UUID +) -> Optional[Any]: + entity_type = _coerce_entity_type(entity_type) + + if entity_type == BrokerEntityType.PROJECT: + return ( + db.query(ProjectSubmission) + .filter( + ProjectSubmission.project_id == entity_id, + ProjectSubmission.attempt_id == attempt_id, + ) + .order_by(ProjectSubmission.created_at.desc()) + .first() + ) + if entity_type == BrokerEntityType.SAMPLE: + return ( + db.query(SampleSubmission) + .filter( + SampleSubmission.sample_id == entity_id, SampleSubmission.attempt_id == attempt_id + ) + .order_by(SampleSubmission.created_at.desc()) + .first() + ) + if entity_type == BrokerEntityType.EXPERIMENT: + result = ( + db.query(ExperimentSubmission) + .filter( + ExperimentSubmission.experiment_id == entity_id, + ExperimentSubmission.attempt_id == attempt_id, + ) + .order_by(ExperimentSubmission.created_at.desc()) + .first() + ) + return result + if entity_type == BrokerEntityType.RUN: + result = ( + db.query(ReadSubmission) + .filter(ReadSubmission.read_id == entity_id, ReadSubmission.attempt_id == attempt_id) + .order_by(ReadSubmission.created_at.desc()) + .first() + ) + return result + + +def _register_submission_accession( + db: Session, + entity_type: BrokerEntityType | str, + row: Any, + accession: Optional[str], + secondary_accession: Optional[str] = None, +) -> None: + entity_type = _coerce_entity_type(entity_type) + if not accession: + return + + if entity_type == BrokerEntityType.PROJECT: + registry_entity_type = "project" + registry_entity_id = row.project_id + elif entity_type == BrokerEntityType.SAMPLE: + registry_entity_type = "sample" + registry_entity_id = row.sample_id + elif entity_type == BrokerEntityType.EXPERIMENT: + registry_entity_type = "experiment" + registry_entity_id = row.experiment_id + else: + registry_entity_type = "read" + registry_entity_id = row.read_id + + if registry_entity_id is None: + return + + stmt = insert(AccessionRegistry).values( + authority=getattr(row, "authority", None) or "ENA", + accession=accession, + secondary_accession=secondary_accession, + entity_type=registry_entity_type, + entity_id=registry_entity_id, + accepted_at=datetime.now(timezone.utc), + ) + stmt = stmt.on_conflict_do_nothing(index_elements=[AccessionRegistry.accession]) + db.execute(stmt) + + +@router.post("/claims/ready", response_model=BrokerReadyClaimResponse) +@policy("broker:claim") +def claim_ready_entities( + *, + payload: BrokerReadyClaimRequest, + current_user: User = Depends(get_current_active_user), + db: Session = Depends(get_db), +) -> BrokerReadyClaimResponse: + expire_stale_leases(db) + + tax_id = _normalise_tax_id_value(payload.tax_id) + organism = _get_organism_by_tax_id(db, tax_id) + if not organism: + raise HTTPException(status_code=404, detail=f"Organism with tax_id {tax_id} not found") + organism_key = organism.grouping_key + + project_rows = _latest_per_entity(_query_ready_project_submissions(db, tax_id), "project_id") + sample_rows = _latest_per_entity(_query_ready_sample_submissions(db, tax_id), "sample_id") + experiment_rows = _latest_per_entity( + _query_ready_experiment_submissions(db, tax_id), "experiment_id" + ) + run_rows = _latest_per_entity(_query_ready_run_submissions(db, tax_id), "read_id") + + if not any([project_rows, sample_rows, experiment_rows, run_rows]): + # Return empty response when organism exists but has no claimable entities + return BrokerReadyClaimResponse( + attempt_id=None, + tax_id=str(tax_id), + scope="full", + entities=[], + ) + + attempt = _create_contract_attempt(db, organism_key=organism_key) + now = datetime.now(timezone.utc) + + entities: List[BrokerClaimEntity] = [] + for row in project_rows: + _claim_submission_row(db, attempt=attempt, row=row, entity_type="project", now=now) + entities.append( + _build_contract_entity(db, BrokerEntityType.PROJECT, row.project_id, tax_id, row) + ) + for row in sample_rows: + _claim_submission_row(db, attempt=attempt, row=row, entity_type="sample", now=now) + entities.append( + _build_contract_entity(db, BrokerEntityType.SAMPLE, row.sample_id, tax_id, row) + ) + for row in experiment_rows: + _claim_submission_row(db, attempt=attempt, row=row, entity_type="experiment", now=now) + entities.append( + _build_contract_entity(db, BrokerEntityType.EXPERIMENT, row.experiment_id, tax_id, row) + ) + for row in run_rows: + _claim_submission_row(db, attempt=attempt, row=row, entity_type="read", now=now) + entities.append(_build_contract_entity(db, BrokerEntityType.RUN, row.read_id, tax_id, row)) + + db.commit() + + return BrokerReadyClaimResponse( + attempt_id=attempt.id, + tax_id=str(tax_id), + scope="full", + entities=entities, + ) + + +@router.post("/claims/entity", response_model=BrokerTargetedClaimResponse) +@policy("broker:claim") +def claim_specific_entity( + *, + payload: BrokerTargetedClaimRequest, + current_user: User = Depends(get_current_active_user), + db: Session = Depends(get_db), +) -> BrokerTargetedClaimResponse: + expire_stale_leases(db) + + row = _find_latest_claimable_entity_submission(db, payload.entity_type, payload.entity_id) + if not row: + raise HTTPException(status_code=404, detail="No claimable submission found for entity") + + organism_key, tax_id = _lookup_taxonomy_for_entity(db, payload.entity_type, payload.entity_id) + if tax_id is None: + raise HTTPException(status_code=404, detail="Entity not found") + + attempt = _create_contract_attempt(db, organism_key=organism_key) + _claim_submission_row( + db, + attempt=attempt, + row=row, + entity_type="read" + if payload.entity_type == BrokerEntityType.RUN + else payload.entity_type.value, + now=datetime.now(timezone.utc), + ) + db.commit() + + return BrokerTargetedClaimResponse( + attempt_id=attempt.id, + tax_id=str(tax_id), + entities=[_build_contract_entity(db, payload.entity_type, payload.entity_id, tax_id, row)], + ) + + +@router.post("/claims/batch", response_model=BrokerBatchClaimResponse) +@policy("broker:claim") +def claim_batch_entities( + *, + payload: BrokerBatchClaimRequest, + current_user: User = Depends(get_current_active_user), + db: Session = Depends(get_db), +) -> BrokerBatchClaimResponse: + """ + Claim multiple specific entities by their IDs. + Supports claiming one or more entities across different types in a single request. + """ + expire_stale_leases(db) + + # Collect all entity IDs with their types + entity_requests: List[tuple[BrokerEntityType, UUID]] = [] + + for project_id in payload.project_ids: + entity_requests.append((BrokerEntityType.PROJECT, project_id)) + for sample_id in payload.sample_ids: + entity_requests.append((BrokerEntityType.SAMPLE, sample_id)) + for experiment_id in payload.experiment_ids: + entity_requests.append((BrokerEntityType.EXPERIMENT, experiment_id)) + for run_id in payload.run_ids: + entity_requests.append((BrokerEntityType.RUN, run_id)) + + if not entity_requests: + raise HTTPException(status_code=400, detail="At least one entity ID must be provided") + + # Find all claimable submissions + rows_with_metadata: List[tuple[Any, BrokerEntityType, UUID, str, int]] = [] + + for entity_type, entity_id in entity_requests: + row = _find_latest_claimable_entity_submission(db, entity_type, entity_id) + if not row: + raise HTTPException( + status_code=404, + detail=f"No claimable submission found for {entity_type.value} {entity_id}", + ) + + organism_key, tax_id = _lookup_taxonomy_for_entity(db, entity_type, entity_id) + if tax_id is None: + raise HTTPException( + status_code=404, detail=f"Entity not found: {entity_type.value} {entity_id}" + ) + + rows_with_metadata.append((row, entity_type, entity_id, organism_key, tax_id)) + + # Check if all entities belong to the same organism + unique_tax_ids = set(tax_id for _, _, _, _, tax_id in rows_with_metadata) + unique_organism_keys = set(org_key for _, _, _, org_key, _ in rows_with_metadata) + + # Use the first organism_key for the attempt (or None if multi-organism) + organism_key = rows_with_metadata[0][3] if len(unique_organism_keys) == 1 else None + tax_id_str = str(rows_with_metadata[0][4]) if len(unique_tax_ids) == 1 else None + + # Create a single attempt for all entities + attempt = _create_contract_attempt(db, organism_key=organism_key) + now = datetime.now(timezone.utc) + + # Claim all submissions and build response entities + entities: List[BrokerClaimEntity] = [] + for row, entity_type, entity_id, _, tax_id in rows_with_metadata: + _claim_submission_row( + db, + attempt=attempt, + row=row, + entity_type="read" if entity_type == BrokerEntityType.RUN else entity_type.value, + now=now, + ) + entities.append(_build_contract_entity(db, entity_type, entity_id, tax_id, row)) + + db.commit() + + return BrokerBatchClaimResponse( + attempt_id=attempt.id, + tax_id=tax_id_str, + entities=entities, + ) + + +@router.post("/validation", response_model=BrokerValidationResponse) +@policy("broker:claim") +def validate_entity_submission( + *, + payload: BrokerValidationRequest, + current_user: User = Depends(get_current_active_user), + db: Session = Depends(get_db), +) -> BrokerValidationResponse: + row = _find_latest_submission_for_validation(db, payload.entity_type, payload.entity_id) + if not row: + raise HTTPException(status_code=404, detail="Submission entity not found") + + _, tax_id = _lookup_taxonomy_for_entity(db, payload.entity_type, payload.entity_id) + claimed_entity = _build_contract_entity( + db, payload.entity_type, payload.entity_id, tax_id or "", row + ) + return _validate_contract_entity(claimed_entity, payload.overrides) + + +@router.post("/reports/{attempt_id}", response_model=BrokerReportResponse) +@policy("broker:claim") +def report_submission_outcomes( + *, + attempt_id: UUID, + payload: BrokerReportRequest, + current_user: User = Depends(get_current_active_user), + db: Session = Depends(get_db), +) -> BrokerReportResponse: + if not payload.results: + raise HTTPException(status_code=400, detail="results must contain at least one record") + + for result in payload.results: + row = _find_submission_for_attempt(db, result.entity_type, result.entity_id, attempt_id) + if not row: + raise HTTPException(status_code=404, detail="Submission attempt item not found") + if row.status != "submitting": + raise HTTPException(status_code=409, detail="Submission is not in submitting state") + + row.status = _map_report_state_to_submission_status(result.status) + + # Build response payload from provided fields or use full response_payload if provided + if result.response_payload: + row.response_payload = result.response_payload + else: + row.response_payload = { + "receipt_path": result.receipt_path, + "message": result.message, + "errors": result.errors, + } + + # Handle accessions + if result.accession: + row.accession = result.accession + # Register both primary and secondary accessions in accession_registry + _register_submission_accession( + db, result.entity_type, row, result.accession, result.secondary_accession + ) + + # Store secondary accession (public accession) in biosample_accession field for samples + if result.secondary_accession and hasattr(row, "biosample_accession"): + row.biosample_accession = result.secondary_accession + + if row.status != "submitting": + row.attempt_id = None + row.lock_acquired_at = None + row.lock_expires_at = None + row.finalised_attempt_id = attempt_id + db.add( + SubmissionEvent( + attempt_id=attempt_id, + entity_type="read" + if _coerce_entity_type(result.entity_type) == BrokerEntityType.RUN + else _coerce_entity_type(result.entity_type).value, + submission_id=row.id, + action="accepted" if row.status == "accepted" else "rejected", + accession=result.accession, + details=row.response_payload, + ) + ) + + if row.status == "rejected": + _create_new_draft_submission_after_rejection( + db, + entity_type=_coerce_entity_type(result.entity_type), + row=row, + ) + + try: + db.commit() + except IntegrityError as exc: + db.rollback() + # This is an expected class of failure when accession_registry contains conflicting rows + # (e.g. the accession already exists for a different entity so the registry insert is skipped, + # then a FK on the submission table fails). Avoid emitting a full stack trace in logs. + logger.warning( + "Report outcomes failed due to integrity error (attempt_id=%s): %s", + attempt_id, + str(getattr(exc, "orig", exc)), + ) + raise HTTPException( + status_code=409, + detail=( + "Report failed due to database integrity constraints. " + "A common cause is an accession already existing in accession_registry for a different entity " + "(so the registry insert is skipped/blocked), which then violates the submission table FK. " + "Check accession_registry for conflicting rows for the reported accession(s)." + ), + ) from exc + + return BrokerReportResponse(attempt_id=attempt_id, accepted=True, message="reported") + + # ---------- Endpoints ---------- # NOTE: More specific routes must come before parameterized routes in FastAPI # Otherwise /claim would match /organisms/{organism_key}/claim with organism_key="claim" diff --git a/app/api/v1/endpoints/samples.py b/app/api/v1/endpoints/samples.py index 7352901..51a993e 100644 --- a/app/api/v1/endpoints/samples.py +++ b/app/api/v1/endpoints/samples.py @@ -14,6 +14,7 @@ from app.core.policy import policy from app.models.experiment import Experiment from app.models.organism import Organism +from app.models.project import Project from app.models.sample import Sample, SampleSubmission from app.models.user import User from app.schemas.bulk_import import BulkImportResponse, BulkSampleImport @@ -31,6 +32,20 @@ router = APIRouter() +def _get_genomic_data_project_id(db: Session, organism_key: str) -> UUID: + """Get the genomic_data project ID for an organism.""" + project = ( + db.query(Project) + .filter(Project.organism_key == organism_key, Project.project_type == "genomic_data") + .first() + ) + if not project: + raise HTTPException( + status_code=404, detail=f"No genomic_data project found for organism {organism_key}" + ) + return project.id + + @router.get("/", response_model=List[SampleSchema]) def read_samples( db: Session = Depends(get_db), @@ -108,13 +123,16 @@ def create_sample( sample_id = uuid.uuid4() # Compute required NOT NULL fields and fallbacks - lifestage = sample_in.lifestage or "unknown" - sex = sample_in.sex or "unknown" - organism_part = sample_in.organism_part or "unknown" - region_and_locality = getattr(sample_in, "region_and_locality", None) or "unknown" - country_or_sea = getattr(sample_in, "country_or_sea", None) or "unknown" - habitat = sample_in.habitat or "unknown" - collection_date_val = getattr(sample_in, "collection_date", None) or None + lifestage = sample_in.get("lifestage") or "unknown" + sex = sample_in.get("sex") or "unknown" + organism_part = sample_in.get("organism_part") or "unknown" + region_and_locality = sample_in.get("region_and_locality") or "unknown" + country_or_sea = sample_in.get("country_or_sea") or "unknown" + habitat = sample_in.get("habitat") or "unknown" + collection_date_val = sample_in.get("collection_date") or "unknown" + collected_by = sample_in.get("collected_by") or "unknown" + collecting_institution = sample_in.get("collecting_institution") or "unknown" + # Accept raw string and allow missing collection_date # Build kwargs dynamically so we don't pass None for DB server_default columns @@ -182,15 +200,15 @@ def create_sample( state_or_region=sample_in.state_or_region, country_or_sea=country_or_sea, indigenous_location=sample_in.indigenous_location, - latitude=to_float(sample_in.decimal_latitude), - longitude=to_float(sample_in.decimal_longitude), + latitude=to_float(sample_in.latitude), + longitude=to_float(sample_in.longitude), elevation=to_float(sample_in.elevation), depth=to_float(sample_in.depth), habitat=habitat, - collection_method=sample_in.description_of_collection_method, + collected_by=collected_by, + collecting_institution=collecting_institution, + collection_method=sample_in.collection_method, collection_date=collection_date_val, - collected_by=sample_in.collected_by, - collecting_institute=sample_in.collector_institute, collection_permit=sample_in.collection_permit, data_context=sample_in.data_context, bioplatforms_project_id=sample_in.bioplatforms_project_id, @@ -230,12 +248,16 @@ def create_sample( if atol_key in sample_data: prepared_payload[ena_key] = sample_data[atol_key] + # Get project_id for this organism + project_id = _get_genomic_data_project_id(db, sample.organism_key) + sample_submission = SampleSubmission( sample_id=sample_id, authority=sample_in.authority, entity_type_const="sample", prepared_payload=prepared_payload, status=SubmissionStatus.DRAFT, + project_id=project_id, ) db.add(sample_submission) db.commit() @@ -271,16 +293,21 @@ def _create_sample_with_submission( lifestage = sample_data.get("lifestage") or "unknown" sex = sample_data.get("sex") or "unknown" organism_part = sample_data.get("organism_part") or "unknown" - region_and_locality = ( - sample_data.get("region_and_locality") - or sample_data.get("collection_location") - or "unknown" - ) + + region_and_locality = getattr(sample_data, "region_and_locality", None) or "unknown" + country_or_sea = getattr(sample_data, "country_or_sea", None) or "unknown" + habitat = sample_data.get("habitat") or "unknown" + collection_date_val = getattr(sample_data, "collection_date", None) or None + + region_and_locality = sample_data.get("region_and_locality") or "unknown" + country_or_sea = sample_data.get("country_or_sea") or "unknown" habitat = sample_data.get("habitat") or "unknown" collection_date_val = sample_data.get("date_of_collection") or sample_data.get( "collection_date" ) + collected_by = sample_data.get("collected_by") or "unknown" + collecting_institution = sample_data.get("collecting_institution") or "unknown" sample_kwargs = dict( id=sample_id, @@ -301,6 +328,8 @@ def _create_sample_with_submission( habitat=habitat, collection_method=sample_data.get("description_of_collection_method") or sample_data.get("collection_method"), + collected_by=collected_by, + collecting_institution=collecting_institution, collection_date=collection_date_val, collection_permit=sample_data.get("collection_permit"), data_context=sample_data.get("data_context"), @@ -314,8 +343,8 @@ def _create_sample_with_submission( preservation_temperature=sample_data.get("preservation_temperature"), project_name=sample_data.get("project_name"), biosample_accession=sample_data.get("biosample_accession"), - latitude=to_float(sample_data.get("decimal_latitude")), - longitude=to_float(sample_data.get("decimal_longitude")), + latitude=to_float(sample_data.get("latitude")), + longitude=to_float(sample_data.get("longitude")), elevation=to_float(sample_data.get("elevation")), depth=to_float(sample_data.get("depth")), # Parent-child relationship fields @@ -342,6 +371,9 @@ def _create_sample_with_submission( if atol_key in sample_data: prepared_payload[ena_key] = sample_data[atol_key] + # Get project_id for this organism + project_id = _get_genomic_data_project_id(db, organism_key) + # Create sample_submission record sample_submission = SampleSubmission( id=uuid.uuid4(), @@ -349,6 +381,7 @@ def _create_sample_with_submission( authority="ENA", entity_type_const="sample", prepared_payload=prepared_payload, + project_id=project_id, ) return sample, sample_submission @@ -423,6 +456,7 @@ def bulk_import_specimen_samples( ) skipped_count += 1 continue + sample_data["organism_part"] = "WHOLE ORGANISM" # Create specimen sample (bpa_sample_id is optional for specimens) sample, sample_submission = _create_sample_with_submission( @@ -656,12 +690,14 @@ def update_sample( new_sample_submission = None latest_sample_submission = {} if not sample_submission: + project_id = _get_genomic_data_project_id(db, sample.organism_key) new_sample_submission = SampleSubmission( sample_id=sample_id, - authority=sample_submission.authority, + authority="ENA", entity_type_const="sample", prepared_payload=prepared_payload, status="draft", + project_id=project_id, ) db.add(new_sample_submission) else: @@ -678,6 +714,7 @@ def update_sample( ): # leave the old record for logs and create a new record # retain accessions if they exist (accessions may not exist if status is 'rejected' and the sample has not successfully been submitted in the past) + project_id = _get_genomic_data_project_id(db, sample.organism_key) new_sample_submission = SampleSubmission( sample_id=sample_id, authority=sample_submission.authority, @@ -687,6 +724,7 @@ def update_sample( accession=sample_submission.accession, biosample_accession=sample_submission.biosample_accession, status="draft", + project_id=project_id, ) db.add(new_sample_submission) @@ -695,6 +733,7 @@ def update_sample( # retain accessions setattr(latest_sample_submission, "status", "replaced") db.add(latest_sample_submission) + project_id = _get_genomic_data_project_id(db, sample.organism_key) new_sample_submission = SampleSubmission( sample_id=sample_id, authority=sample_submission.authority, @@ -704,6 +743,7 @@ def update_sample( accession=sample_submission.accession, biosample_accession=sample_submission.biosample_accession, status="draft", + project_id=project_id, ) db.add(new_sample_submission) elif ( @@ -821,16 +861,10 @@ def bulk_import_samples( lifestage = sample_data.get("lifestage") or "unknown" sex = sample_data.get("sex") or "unknown" organism_part = sample_data.get("organism_part") or "unknown" - region_and_locality = ( - sample_data.get("region_and_locality") - or sample_data.get("collection_location") - or "unknown" - ) + region_and_locality = sample_data.get("region_and_locality") or "unknown" country_or_sea = sample_data.get("country_or_sea") or "unknown" habitat = sample_data.get("habitat") or "unknown" - collection_date_val = sample_data.get("date_of_collection") or sample_data.get( - "collection_date" - ) + collection_date_val = sample_data.get("collection_date") or "unknown" # Accept raw string and allow missing collection_date # Determine sample kind - default to specimen for bulk imports @@ -878,8 +912,8 @@ def bulk_import_samples( collection_permit=sample_data.get("collection_permit"), data_context=sample_data.get("data_context"), bioplatforms_project_id=sample_data.get("bioplatforms_project_id"), - latitude=to_float(sample_data.get("decimal_latitude")), - longitude=to_float(sample_data.get("decimal_longitude")), + latitude=to_float(sample_data.get("latitude")), + longitude=to_float(sample_data.get("longitude")), elevation=to_float(sample_data.get("elevation")), depth=to_float(sample_data.get("depth")), # Parent-child relationship fields @@ -904,6 +938,9 @@ def bulk_import_samples( if atol_key in sample_data: prepared_payload[ena_key] = sample_data[atol_key] + # Get project_id for this organism + project_id = _get_genomic_data_project_id(db, organism_key) + # Create sample_submission record sample_submission = SampleSubmission( id=uuid.uuid4(), @@ -911,6 +948,7 @@ def bulk_import_samples( authority="ENA", entity_type_const="sample", prepared_payload=prepared_payload, + project_id=project_id, ) db.add(sample_submission) diff --git a/app/models/experiment.py b/app/models/experiment.py index 876ce7c..deca0ec 100644 --- a/app/models/experiment.py +++ b/app/models/experiment.py @@ -21,6 +21,9 @@ class Experiment(Base): sample_id = Column( UUID(as_uuid=True), ForeignKey("sample.id", ondelete="CASCADE"), nullable=False ) + project_id = Column( + UUID(as_uuid=True), ForeignKey("project.id", ondelete="CASCADE"), nullable=False + ) bpa_package_id = Column(Text, unique=True, nullable=False) design_description = Column(Text, nullable=True) bpa_library_id = Column(Text, nullable=True) @@ -59,6 +62,9 @@ class Experiment(Base): sample = relationship( "Sample", backref=backref("exp_sample_records", cascade="all, delete-orphan") ) + project = relationship( + "Project", backref=backref("exp_project_records", cascade="all, delete-orphan") + ) class ExperimentSubmission(Base): @@ -91,16 +97,6 @@ class ExperimentSubmission(Base): default="draft", ) - sample_id = Column( - UUID(as_uuid=True), ForeignKey("sample.id", ondelete="SET NULL"), nullable=False - ) - project_id = Column( - UUID(as_uuid=True), ForeignKey("project.id", ondelete="SET NULL"), nullable=True - ) - - project_accession = Column(Text, nullable=True) - sample_accession = Column(Text, nullable=True) - prepared_payload = Column(JSONB, nullable=True) response_payload = Column(JSONB, nullable=True) accession = Column(Text, nullable=True) @@ -128,12 +124,6 @@ class ExperimentSubmission(Base): experiment = relationship( "Experiment", backref=backref("exp_submission_records", cascade="all, delete-orphan") ) - sample = relationship( - "Sample", backref=backref("exp_sample_submission_records", cascade="all, delete-orphan") - ) - project = relationship( - "Project", backref=backref("exp_project_submission_records", cascade="all, delete-orphan") - ) # Table constraints __table_args__ = ( @@ -150,18 +140,6 @@ class ExperimentSubmission(Base): deferrable=True, initially="DEFERRED", ), - # Foreign key constraint for project accession - ForeignKeyConstraint( - ["project_accession", "authority"], - ["accession_registry.accession", "accession_registry.authority"], - name="fk_proj_acc", - ), - # Foreign key constraint for sample accession - ForeignKeyConstraint( - ["sample_accession", "authority"], - ["accession_registry.accession", "accession_registry.authority"], - name="fk_samp_acc", - ), # This is a simplified version of the SQL constraint: # UNIQUE (experiment_id, authority) WHERE (status = 'accepted' AND accession IS NOT NULL) # SQLAlchemy doesn't directly support WHERE clauses in constraints, so this would need custom SQL diff --git a/app/models/read.py b/app/models/read.py index 76f1b91..0a6bb90 100644 --- a/app/models/read.py +++ b/app/models/read.py @@ -103,11 +103,6 @@ class ReadSubmission(Base): experiment_id = Column( UUID(as_uuid=True), ForeignKey("experiment.id", ondelete="CASCADE"), nullable=True ) - project_id = Column( - UUID(as_uuid=True), ForeignKey("project.id", ondelete="SET NULL"), nullable=True - ) - - experiment_accession = Column(Text, nullable=True) accession = Column(Text, nullable=True) @@ -129,9 +124,6 @@ class ReadSubmission(Base): experiment = relationship( "Experiment", backref=backref("read_exp_submission_records", cascade="all, delete-orphan") ) - project = relationship( - "Project", backref=backref("read_proj_submission_records", cascade="all, delete-orphan") - ) # Broker lease/claim fields attempt_id = Column(UUID(as_uuid=True), nullable=True) @@ -154,12 +146,6 @@ class ReadSubmission(Base): deferrable=True, initially="DEFERRED", ), - # Foreign key constraint for experiment accession - ForeignKeyConstraint( - ["experiment_accession", "authority"], - ["accession_registry.accession", "accession_registry.authority"], - name="fk_exp_acc", - ), # This is a simplified version of the SQL constraint: # UNIQUE (read_id, authority) WHERE (status = 'accepted' AND accession IS NOT NULL) # SQLAlchemy doesn't directly support WHERE clauses in constraints, so this would need custom SQL diff --git a/app/models/sample.py b/app/models/sample.py index 3eaf47c..79fb3ae 100644 --- a/app/models/sample.py +++ b/app/models/sample.py @@ -51,8 +51,9 @@ class Sample(Base): habitat = Column(Text, nullable=False) collection_method = Column(Text, nullable=True) collection_date = Column(Text, nullable=True) - collected_by = Column(Text, nullable=False, server_default=text("'not provided'")) - collecting_institution = Column(Text, nullable=False, server_default=text("'not provided'")) + # TODO unify approach to default values - either set in DB or in code, but not both - currently a mix of both approaches + collected_by = Column(Text, nullable=False) + collecting_institution = Column(Text, nullable=False) collection_permit = Column(Text, nullable=True) data_context = Column(Text, nullable=True) bioplatforms_project_id = Column(Text, nullable=True) @@ -128,6 +129,10 @@ class SampleSubmission(Base): accession = Column(Text, nullable=True) biosample_accession = Column(Text, nullable=True) entity_type_const = Column(Text, nullable=False, default="sample", server_default="sample") + + # Prerequisite FK (for normalized accession lookups) + project_id = Column(UUID(as_uuid=True), ForeignKey("project.id"), nullable=False) + submitted_at = Column(DateTime(timezone=True), nullable=True) created_at = Column(DateTime(timezone=True), nullable=False, server_default=func.now()) updated_at = Column( diff --git a/app/schemas/broker_contract.py b/app/schemas/broker_contract.py new file mode 100644 index 0000000..42931d7 --- /dev/null +++ b/app/schemas/broker_contract.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +from enum import Enum +from typing import Any, Dict, List, Optional +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field + + +class BrokerEntityType(str, Enum): + PROJECT = "project" + SAMPLE = "sample" + EXPERIMENT = "experiment" + RUN = "run" + + +class BrokerPrerequisites(BaseModel): + # Existing accessions (already submitted and available) + project_accession: Optional[str] = None + sample_accession: Optional[str] = None + experiment_accession: Optional[str] = None + run_accession: Optional[str] = None + study_accession: Optional[str] = None + analysis_accession: Optional[str] = None + + +class BrokerFileMetadata(BaseModel): + filename: str + filetype: str + + +class BrokerClaimEntity(BaseModel): + model_config = ConfigDict(use_enum_values=True) + + type: BrokerEntityType + id: UUID + tax_id: str + scientific_name: Optional[str] = None + payload: Optional[Dict[str, Any]] = None + prerequisites: Optional[BrokerPrerequisites] = None + files: Optional[List[BrokerFileMetadata]] = None + file_metadata: Optional[List[BrokerFileMetadata]] = None + + +class BrokerReadyClaimRequest(BaseModel): + tax_id: str + + +class BrokerReadyClaimResponse(BaseModel): + attempt_id: Optional[UUID] + tax_id: str + scope: str = "full" + entities: List[BrokerClaimEntity] = Field(default_factory=list) + + +class BrokerTargetedClaimRequest(BaseModel): + entity_type: BrokerEntityType + entity_id: UUID + + +class BrokerTargetedClaimResponse(BaseModel): + attempt_id: UUID + tax_id: str + entities: List[BrokerClaimEntity] = Field(default_factory=list) + + +class BrokerBatchClaimRequest(BaseModel): + project_ids: List[UUID] = Field(default_factory=list) + sample_ids: List[UUID] = Field(default_factory=list) + experiment_ids: List[UUID] = Field(default_factory=list) + run_ids: List[UUID] = Field(default_factory=list) + + +class BrokerBatchClaimResponse(BaseModel): + attempt_id: UUID + tax_id: Optional[str] = None # None if multi-organism batch + entities: List[BrokerClaimEntity] = Field(default_factory=list) + + +class BrokerValidationIssue(BaseModel): + field: str + message: str + + +class BrokerValidationRequest(BaseModel): + entity_type: BrokerEntityType + entity_id: UUID + overrides: BrokerPrerequisites = Field(default_factory=BrokerPrerequisites) + + +class BrokerValidationResponse(BaseModel): + model_config = ConfigDict(use_enum_values=True) + + entity_type: BrokerEntityType + entity_id: UUID + valid: bool + issues: List[BrokerValidationIssue] = Field(default_factory=list) + resolved_prerequisites: Dict[str, str] = Field(default_factory=dict) + + +class BrokerReportRecord(BaseModel): + model_config = ConfigDict(use_enum_values=True) + + entity_type: BrokerEntityType + entity_id: UUID + status: ( + str # "completed"/"accepted"/"success" -> accepted, "failed"/"rejected"/"error" -> rejected + ) + accession: Optional[str] = None # Primary/ENA accession (internal ENA ID) + secondary_accession: Optional[str] = None # Public accession (SAMEA*, PRJEB*, etc.) + receipt_path: Optional[str] = None + message: Optional[str] = None + errors: List[Any] = Field(default_factory=list) + response_payload: Optional[Dict[str, Any]] = None # Full ENA response + + +class BrokerReportRequest(BaseModel): + tax_id: Optional[str | int] = None + results: List[BrokerReportRecord] = Field(default_factory=list) + + +class BrokerReportResponse(BaseModel): + attempt_id: UUID + accepted: bool + message: str diff --git a/app/schemas/experiment.py b/app/schemas/experiment.py index fbaaa7e..1fb4e7d 100644 --- a/app/schemas/experiment.py +++ b/app/schemas/experiment.py @@ -107,8 +107,6 @@ class Experiment(ExperimentInDBBase): class ExperimentSubmissionBase(BaseModel): """Base ExperimentSubmission schema with common attributes.""" - sample_id: UUID - project_id: Optional[UUID] = None authority: str = Field(default="ENA", description="Authority for the submission") status: SubmissionStatus = Field( default=SubmissionStatus.DRAFT, description="Status of the submission" @@ -135,8 +133,6 @@ class ExperimentSubmissionCreate(ExperimentSubmissionBase): class ExperimentSubmissionUpdate(BaseModel): """Schema for updating an existing experiment submission.""" - sample_id: Optional[UUID] = None - project_id: Optional[UUID] = None authority: Optional[str] = None project_accession: Optional[str] = None sample_accession: Optional[str] = None diff --git a/app/services/experiment_service.py b/app/services/experiment_service.py index d003000..99a6d55 100644 --- a/app/services/experiment_service.py +++ b/app/services/experiment_service.py @@ -73,10 +73,27 @@ def create_experiment(self, db: Session, *, experiment_in: ExperimentCreate) -> """Create experiment and corresponding submission with prepared payload.""" experiment_id = uuid.uuid4() + sample = db.query(Sample).filter(Sample.id == experiment_in.sample_id).first() + if not sample: + raise RuntimeError(f"Sample not found: {experiment_in.sample_id}") + + project = ( + db.query(Project) + .filter( + Project.organism_key == sample.organism_key, + Project.project_type == "genomic_data", + ) + .first() + ) + if not project: + raise RuntimeError( + f"No genomic_data project found for organism_key '{sample.organism_key}'" + ) + # Auto-map fields from Pydantic schema to Experiment columns using shared mapper exp_data = experiment_in.model_dump(exclude_unset=True) transforms = {"insert_size": (lambda v: str(v) if v is not None else None)} - inject = {"id": experiment_id} + inject = {"id": experiment_id, "project_id": project.id} experiment_kwargs = map_to_model_columns( Experiment, exp_data, @@ -96,8 +113,6 @@ def create_experiment(self, db: Session, *, experiment_in: ExperimentCreate) -> experiment_submission = ExperimentSubmission( experiment_id=experiment_id, - sample_id=experiment_in.sample_id, - project_id=exp_data.get("project_id"), entity_type_const="experiment", prepared_payload=prepared_payload, status=SubmissionStatus.DRAFT, @@ -154,8 +169,6 @@ def update_experiment( if not latest_experiment_submission: new_experiment_submission = ExperimentSubmission( experiment_id=experiment_id, - sample_id=experiment.sample_id, - project_id=experiment.project_id, authority="ENA", entity_type_const="experiment", prepared_payload=prepared_payload, @@ -170,8 +183,6 @@ def update_experiment( elif latest_experiment_submission.status in ("rejected", "replaced"): new_experiment_submission = ExperimentSubmission( experiment_id=experiment_id, - sample_id=experiment.sample_id, - project_id=experiment.project_id, authority=latest_experiment_submission.authority, entity_type_const="experiment", prepared_payload=prepared_payload, @@ -187,8 +198,6 @@ def update_experiment( db.add(latest_experiment_submission) new_experiment_submission = ExperimentSubmission( experiment_id=experiment_id, - sample_id=experiment.sample_id, - project_id=experiment.project_id, authority=latest_experiment_submission.authority, entity_type_const="experiment", prepared_payload=prepared_payload, @@ -261,12 +270,6 @@ def bulk_import_experiments( # Still process reads for existing experiment experiment_id = existing_experiment.id - # Find project for read submissions - project_id = None - project = db.query(Project).first() - if project: - project_id = project.id - # Process reads even though experiment exists if isinstance(experiment_data.get("runs"), list): for run in experiment_data["runs"]: @@ -319,7 +322,6 @@ def bulk_import_experiments( id=uuid.uuid4(), read_id=read.id, experiment_id=experiment_id, - project_id=project_id, authority="ENA", entity_type_const="read", prepared_payload=run_prepared_payload, @@ -382,9 +384,32 @@ def bulk_import_experiments( # Create experiment experiment_id = uuid.uuid4() sample_id = sample.id + + organism_key = sample.organism_key + + project = ( + db.query(Project) + .filter( + Project.organism_key == organism_key, + Project.project_type == "genomic_data", + ) + .first() + ) + if not project: + raise RuntimeError( + f"No genomic_data project found for organism_key '{organism_key}'" + ) + + project_id = project.id + aliases = {"GAL": "gal", "extraction_protocol_DOI": "extraction_protocol_doi"} transforms = {"insert_size": (lambda v: str(v) if v is not None else None)} - inject = {"id": experiment_id, "sample_id": sample_id, "bpa_package_id": package_id} + inject = { + "id": experiment_id, + "sample_id": sample_id, + "project_id": project_id, + "bpa_package_id": package_id, + } experiment_kwargs = map_to_model_columns( Experiment, experiment_data, @@ -395,12 +420,6 @@ def bulk_import_experiments( experiment = Experiment(**experiment_kwargs) db.add(experiment) - # Find a project (fallback to any project for now) - project_id = None - project = db.query(Project).first() - if project: - project_id = project.id - # Build prepared payload for experiment submission prepared_payload: Dict[str, Any] = {} for ena_key, atol_key in experiment_mapping.items(): @@ -410,8 +429,6 @@ def bulk_import_experiments( experiment_submission = ExperimentSubmission( id=uuid.uuid4(), experiment_id=experiment_id, - sample_id=sample_id, - project_id=project_id, authority="ENA", entity_type_const="experiment", prepared_payload=prepared_payload, @@ -459,7 +476,6 @@ def bulk_import_experiments( id=uuid.uuid4(), read_id=read.id, experiment_id=experiment_id, - project_id=project_id, authority="ENA", entity_type_const="read", prepared_payload=run_prepared_payload, @@ -522,14 +538,18 @@ def get_by_experiment_id(self, db: Session, experiment_id: UUID) -> List[Experim def get_by_sample_id(self, db: Session, sample_id: UUID) -> List[ExperimentSubmission]: """Get submission experiments by sample ID.""" return ( - db.query(ExperimentSubmission).filter(ExperimentSubmission.sample_id == sample_id).all() + db.query(ExperimentSubmission) + .join(Experiment, Experiment.id == ExperimentSubmission.experiment_id) + .filter(Experiment.sample_id == sample_id) + .all() ) def get_by_project_id(self, db: Session, project_id: UUID) -> List[ExperimentSubmission]: """Get submission experiments by project ID.""" return ( db.query(ExperimentSubmission) - .filter(ExperimentSubmission.project_id == project_id) + .join(Experiment, Experiment.id == ExperimentSubmission.experiment_id) + .filter(Experiment.project_id == project_id) .all() ) diff --git a/app/services/read_service.py b/app/services/read_service.py index d3b183c..a63c8a7 100644 --- a/app/services/read_service.py +++ b/app/services/read_service.py @@ -3,6 +3,7 @@ from sqlalchemy.orm import Session +from app.models.experiment import Experiment from app.models.read import Read, ReadSubmission from app.schemas.read import ReadCreate, ReadUpdate from app.services.base_service import BaseService @@ -50,7 +51,12 @@ def get_by_experiment_id(self, db: Session, experiment_id: UUID) -> List[ReadSub def get_by_project_id(self, db: Session, project_id: UUID) -> List[ReadSubmission]: """Get submission reads by project ID.""" - return db.query(ReadSubmission).filter(ReadSubmission.project_id == project_id).all() + return ( + db.query(ReadSubmission) + .join(Experiment, Experiment.id == ReadSubmission.experiment_id) + .filter(Experiment.project_id == project_id) + .all() + ) def get_by_accession(self, db: Session, accession: str) -> Optional[ReadSubmission]: """Get submission read by accession.""" diff --git a/schema.sql b/schema.sql index 2351e58..69c6829 100644 --- a/schema.sql +++ b/schema.sql @@ -138,6 +138,7 @@ CREATE TABLE project ( description TEXT NOT NULL, centre_name TEXT, study_attributes JSONB, + -- TODO remove submitted_at, should only be for submission table submitted_at TIMESTAMPTZ, status submission_status NOT NULL DEFAULT 'draft', authority authority_type NOT NULL DEFAULT 'ENA', @@ -221,8 +222,8 @@ CREATE TABLE sample ( habitat TEXT NOT NULL, collection_method TEXT, collection_date TEXT, - collected_by TEXT NOT NULL DEFAULT 'not provided', - collecting_institution TEXT NOT NULL DEFAULT 'not provided', + collected_by TEXT NOT NULL, + collecting_institution TEXT NOT NULL, collection_permit TEXT, data_context TEXT, bioplatforms_project_id TEXT, @@ -269,6 +270,10 @@ CREATE TABLE sample_submission ( -- constant to help the composite FK entity_type_const entity_type NOT NULL DEFAULT 'sample' CHECK (entity_type_const = 'sample'), + -- prerequisite FK (for normalized accession lookups) + -- TODO remove: samples are not project-scoped; project context (if any) is derived via experiments + project_id UUID NOT NULL REFERENCES project(id), + -- attempt linkage attempt_id UUID, finalised_attempt_id UUID, @@ -294,6 +299,7 @@ CREATE INDEX IF NOT EXISTS idx_sample_submission_attempt ON sample_submission (a CREATE INDEX IF NOT EXISTS idx_sample_submission_finalised_attempt ON sample_submission (finalised_attempt_id); CREATE INDEX IF NOT EXISTS idx_sample_submission_status ON sample_submission (status); CREATE INDEX IF NOT EXISTS idx_sample_submission_lock_expires_at ON sample_submission (lock_expires_at); +CREATE INDEX IF NOT EXISTS idx_sample_submission_project_id ON sample_submission (project_id); -- Support parent/child lookups for derived samples CREATE INDEX IF NOT EXISTS idx_sample_derived_from_sample_id ON sample(derived_from_sample_id); @@ -323,6 +329,7 @@ CREATE INDEX IF NOT EXISTS idx_sample_organism_specimen_lookup CREATE TABLE experiment ( id UUID PRIMARY KEY DEFAULT uuid_generate_v4(), sample_id UUID NOT NULL REFERENCES sample(id) ON DELETE CASCADE, + project_id UUID NOT NULL REFERENCES project(id) ON DELETE CASCADE, bpa_package_id TEXT UNIQUE NOT NULL, design_description TEXT, bpa_library_id TEXT, @@ -362,13 +369,6 @@ CREATE TABLE experiment_submission ( authority authority_type NOT NULL DEFAULT 'ENA', status submission_status NOT NULL DEFAULT 'draft', - sample_id UUID REFERENCES sample(id) NOT NULL, -- nullable if sample doesn't exist yet? or are we happy to have the constraint that a sample & project needs to exist before we create an experiment, probs - project_id UUID REFERENCES project(id), - -- TO DO do we even need this ? I don't think so - - project_accession TEXT, - sample_accession TEXT, - prepared_payload JSONB, response_payload JSONB, @@ -392,15 +392,7 @@ CREATE TABLE experiment_submission ( CONSTRAINT fk_self_accession FOREIGN KEY (accession, authority, entity_type_const, experiment_id) REFERENCES accession_registry (accession, authority, entity_type, entity_id) - DEFERRABLE INITIALLY DEFERRED, - - -- Upstream accessions also validated (drafts allowed to be NULL): - CONSTRAINT fk_proj_acc - FOREIGN KEY (project_accession, authority) - REFERENCES accession_registry (accession, authority), - CONSTRAINT fk_samp_acc - FOREIGN KEY (sample_accession, authority) - REFERENCES accession_registry (accession, authority) + DEFERRABLE INITIALLY DEFERRED ); -- Broker lease/claim index @@ -456,9 +448,6 @@ CREATE TABLE read_submission ( response_payload JSONB, experiment_id UUID NOT NULL REFERENCES experiment(id) ON DELETE CASCADE, - project_id UUID REFERENCES project(id), - - experiment_accession TEXT, accession TEXT, @@ -480,12 +469,7 @@ CREATE TABLE read_submission ( CONSTRAINT fk_self_accession FOREIGN KEY (accession, authority, entity_type_const, read_id) REFERENCES accession_registry (accession, authority, entity_type, entity_id) - DEFERRABLE INITIALLY DEFERRED, - - -- Upstream accessions also validated (drafts allowed to be NULL): - CONSTRAINT fk_exp_acc - FOREIGN KEY (experiment_accession, authority) - REFERENCES accession_registry (accession, authority) + DEFERRABLE INITIALLY DEFERRED ); CREATE UNIQUE INDEX uq_read_one_accepted @@ -507,6 +491,7 @@ CREATE TABLE submission_attempt ( id UUID PRIMARY KEY DEFAULT uuid_generate_v4(), organism_key TEXT REFERENCES organism(grouping_key), campaign_label TEXT, + -- TODO make ENUM status TEXT NOT NULL DEFAULT 'processing', lock_acquired_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), lock_expires_at TIMESTAMPTZ, diff --git a/scripts/create_user.py b/scripts/create_user.py index 6ec472b..747a851 100755 --- a/scripts/create_user.py +++ b/scripts/create_user.py @@ -73,7 +73,7 @@ def main(): # Database connection parameters parser.add_argument("--host", default="localhost", help="Database host") parser.add_argument("--port", default="5433", help="Database port") - parser.add_argument("--dbname", default="atol_database", help="Database name") + parser.add_argument("--dbname", default="atol_db", help="Database name") parser.add_argument("--user", default="postgres", help="Database user") parser.add_argument("--password", default="postgres", help="Database password") parser.add_argument( diff --git a/tests/unit/endpoints/test_endpoints_broker_contract.py b/tests/unit/endpoints/test_endpoints_broker_contract.py new file mode 100644 index 0000000..9aeadf0 --- /dev/null +++ b/tests/unit/endpoints/test_endpoints_broker_contract.py @@ -0,0 +1,912 @@ +from datetime import datetime, timezone +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from fastapi import HTTPException +from sqlalchemy.exc import IntegrityError + +from app.api.v1.endpoints import broker +from app.models.experiment import Experiment +from app.models.organism import Organism +from app.models.sample import Sample, SampleSubmission +from app.schemas.broker_contract import ( + BrokerBatchClaimRequest, + BrokerEntityType, + BrokerReadyClaimRequest, + BrokerReportRecord, + BrokerReportRequest, + BrokerTargetedClaimRequest, + BrokerValidationRequest, +) + + +class FakeQuery: + """Mock query object for FakeSession.""" + + def __init__(self, result=None): + self._result = result + + def filter(self, *args, **kwargs): + return self + + def scalar(self): + return self._result + + def first(self): + return self._result + + +class FakeSession: + def __init__(self): + self.added = [] + self.committed = False + self.flushed = False + self.rolled_back = False + self.commit_exception = None + self.executed = [] + self._accession_lookups = {} # Map of (entity_type, entity_id) -> accession + self._query_results = {} # Map of tuple(query_args) -> scalar/first result + + def add(self, obj): + if getattr(obj, "id", None) is None: + try: + obj.id = uuid4() + except Exception: + pass + self.added.append(obj) + + def flush(self): + self.flushed = True + + def commit(self): + if self.commit_exception is not None: + raise self.commit_exception + self.committed = True + + def rollback(self): + self.rolled_back = True + + def execute(self, stmt): + self.executed.append(stmt) + + def query(self, *args): + """Mock query method that returns configured results by query args.""" + return FakeQuery(result=self._query_results.get(tuple(args), None)) + + +def _broker_user(): + return SimpleNamespace(is_superuser=False, roles=["broker"], is_active=True) + + +def test_claims_ready_returns_flat_entity_contract(monkeypatch): + db = FakeSession() + project_id = uuid4() + sample_id = uuid4() + experiment_id = uuid4() + run_id = uuid4() + + project_row = SimpleNamespace( + id=uuid4(), + project_id=project_id, + status="ready", + prepared_payload={"alias": "project-1"}, + accession=None, + authority="ENA", + ) + sample_row = SimpleNamespace( + id=uuid4(), + sample_id=sample_id, + project_id=project_id, # FK to project + status="ready", + prepared_payload={ + "alias": "sample-1", + "requires_project_accession": True, + "expected_project_accession": "PRJ000001", # Required but may not exist yet + }, + accession=None, + authority="ENA", + ) + experiment_row = SimpleNamespace( + id=uuid4(), + experiment_id=experiment_id, + sample_id=sample_id, # FK to sample + project_id=project_id, # FK to project + status="ready", + prepared_payload={ + "alias": "experiment-1", + "requires_study_accession": True, + "expected_study_accession": "PRJ000001", # Required but not yet submitted + }, + accession=None, + authority="ENA", + ) + run_row = SimpleNamespace( + id=uuid4(), + read_id=run_id, + experiment_id=experiment_id, # FK to experiment + project_id=project_id, # FK to project + status="ready", + prepared_payload={ + "alias": "run-1", + "file_name": "reads_1.fastq.gz", + "file_format": "fastq", + "expected_experiment_accession": "ERX000001", # Required but not submitted + }, + accession=None, + authority="ENA", + ) + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + + db._query_results[(Organism.scientific_name,)] = "Homo sapiens" + db._query_results[(Experiment.bpa_package_id,)] = "PKG1" + db._query_results[(Sample.kind, Sample.specimen_id, Sample.bpa_sample_id)] = ( + "specimen", + "SP1", + "BPA-S1", + ) + monkeypatch.setattr( + broker, + "_get_organism_by_tax_id", + lambda db_arg, tax_id: SimpleNamespace(grouping_key="org-1"), + ) + monkeypatch.setattr( + broker, "_query_ready_project_submissions", lambda db_arg, tax_id: [project_row] + ) + monkeypatch.setattr( + broker, "_query_ready_sample_submissions", lambda db_arg, tax_id: [sample_row] + ) + monkeypatch.setattr( + broker, "_query_ready_experiment_submissions", lambda db_arg, tax_id: [experiment_row] + ) + monkeypatch.setattr(broker, "_query_ready_run_submissions", lambda db_arg, tax_id: [run_row]) + + response = broker.claim_ready_entities( + payload=BrokerReadyClaimRequest(tax_id="9606"), + current_user=_broker_user(), + db=db, + ) + + assert response.tax_id == "9606" + assert response.scope == "full" + assert [entity.type for entity in response.entities] == [ + BrokerEntityType.PROJECT, + BrokerEntityType.SAMPLE, + BrokerEntityType.EXPERIMENT, + BrokerEntityType.RUN, + ] + + assert response.entities[0].scientific_name == "Homo sapiens" + assert response.entities[1].scientific_name == "Homo sapiens" + assert response.entities[2].scientific_name == "Homo sapiens" + assert response.entities[3].scientific_name == "Homo sapiens" + # With normalized lookups, accessions are None unless they exist in accession_registry + # Projects don't have prerequisite accessions + assert response.entities[0].prerequisites is None # No prerequisites for projects + + assert response.entities[1].prerequisites.project_accession is None + assert response.entities[1].prerequisites.study_accession is None + assert response.entities[1].payload["title"] == "Specimen SP1 for Homo sapiens" + + assert ( + response.entities[2].prerequisites.sample_accession is None + ) # Experiment - sample not in registry + assert response.entities[2].prerequisites.study_accession is None # Project not in registry + assert ( + response.entities[2].payload["title"] + == "Bioplatforms Australia dataset PKG1 for Homo sapiens" + ) + + assert ( + response.entities[3].prerequisites.experiment_accession is None + ) # Missing experiment accession + assert response.entities[3].files[0].filename == "reads_1.fastq.gz" + assert response.entities[3].file_metadata is None + assert sample_row.status == "submitting" + assert db.committed is True + + +def test_report_returns_409_on_integrity_error_during_commit(monkeypatch): + db = FakeSession() + attempt_id = uuid4() + entity_id = uuid4() + row = SimpleNamespace( + id=uuid4(), + sample_id=entity_id, + attempt_id=attempt_id, + status="submitting", + prepared_payload={"alias": "sample-1"}, + authority="ENA", + accession=None, + response_payload=None, + lock_acquired_at=None, + lock_expires_at=None, + finalised_attempt_id=None, + biosample_accession=None, + ) + + monkeypatch.setattr( + broker, + "_find_submission_for_attempt", + lambda db_arg, entity_type_arg, entity_id_arg, attempt_id_arg: row, + ) + monkeypatch.setattr(broker, "_register_submission_accession", lambda *args, **kwargs: None) + + db.commit_exception = IntegrityError( + statement="COMMIT", + params=None, + orig=Exception("fk_self_accession violation"), + ) + + with pytest.raises(HTTPException) as exc: + broker.report_submission_outcomes( + attempt_id=attempt_id, + payload=BrokerReportRequest( + results=[ + BrokerReportRecord( + entity_type=BrokerEntityType.SAMPLE, + entity_id=entity_id, + status="accepted", + accession="ERS123456", + secondary_accession="SAMEA123456", + errors=[], + ) + ], + ), + current_user=_broker_user(), + db=db, + ) + + assert exc.value.status_code == 409 + assert "Report failed due to database integrity constraints" in str(exc.value.detail) + assert db.rolled_back is True + + +def test_project_entities_never_return_prerequisites_even_if_registry_has_accession(monkeypatch): + db = FakeSession() + project_id = uuid4() + + project_row = SimpleNamespace( + id=uuid4(), + project_id=project_id, + status="ready", + prepared_payload={"alias": "project-1"}, + accession=None, + authority="ENA", + ) + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + monkeypatch.setattr( + broker, + "_find_latest_claimable_entity_submission", + lambda db_arg, entity_type_arg, entity_id_arg: project_row, + ) + monkeypatch.setattr( + broker, + "_lookup_taxonomy_for_entity", + lambda db_arg, entity_type_arg, entity_id_arg: ("org-1", 9606), + ) + + # If prerequisites were calculated for projects, we'd see them here. + # This simulates a (possibly incorrect) accession_registry hit. + def _fake_prereqs(db_arg, entity_type_arg, prepared_payload_arg, row_arg): + return broker.BrokerPrerequisites( + project_accession="PRJEB99999", + study_accession="PRJEB99999", + ) + + monkeypatch.setattr(broker, "_extract_broker_prerequisites", _fake_prereqs) + + response = broker.claim_batch_entities( + payload=BrokerBatchClaimRequest(project_ids=[project_id]), + current_user=_broker_user(), + db=db, + ) + + assert len(response.entities) == 1 + assert response.entities[0].type == BrokerEntityType.PROJECT + assert response.entities[0].prerequisites is None + + +@pytest.mark.parametrize( + ("entity_type", "row_attr"), + [ + (BrokerEntityType.PROJECT, "project_id"), + (BrokerEntityType.SAMPLE, "sample_id"), + (BrokerEntityType.EXPERIMENT, "experiment_id"), + (BrokerEntityType.RUN, "read_id"), + ], +) +def test_claims_entity_returns_requested_entity_only(monkeypatch, entity_type, row_attr): + db = FakeSession() + entity_id = uuid4() + row = SimpleNamespace( + id=uuid4(), + status="ready", + prepared_payload={"alias": f"{entity_type.value}-1"}, + accession=None, + project_accession=None, + sample_accession=None, + experiment_accession=None, + ) + setattr(row, row_attr, entity_id) + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + monkeypatch.setattr( + broker, + "_find_latest_claimable_entity_submission", + lambda db_arg, entity_type_arg, entity_id_arg: row, + ) + monkeypatch.setattr( + broker, + "_lookup_taxonomy_for_entity", + lambda db_arg, entity_type_arg, entity_id_arg: ("org-1", 9606), + ) + + response = broker.claim_specific_entity( + payload=BrokerTargetedClaimRequest(entity_type=entity_type, entity_id=entity_id), + current_user=_broker_user(), + db=db, + ) + + assert response.tax_id == "9606" + assert len(response.entities) == 1 + assert response.entities[0].type == entity_type + assert response.entities[0].id == entity_id + assert row.status == "submitting" + + +def test_validation_returns_success_with_override_merge(monkeypatch): + experiment_id = uuid4() + row = SimpleNamespace( + prepared_payload={"alias": "experiment-1"}, + sample_accession=None, + project_accession=None, + ) + + monkeypatch.setattr( + broker, + "_find_latest_submission_for_validation", + lambda db_arg, entity_type_arg, entity_id_arg: row, + ) + monkeypatch.setattr( + broker, + "_lookup_taxonomy_for_entity", + lambda db_arg, entity_type_arg, entity_id_arg: ("org-1", 9606), + ) + + response = broker.validate_entity_submission( + payload=BrokerValidationRequest( + entity_type=BrokerEntityType.EXPERIMENT, + entity_id=experiment_id, + overrides={"sample_accession": "SAMEA000001"}, + ), + current_user=_broker_user(), + db=FakeSession(), + ) + + assert response.valid is True + assert response.issues == [] + assert response.resolved_prerequisites == {"sample_accession": "SAMEA000001"} + + +def test_validation_returns_failure_when_required_prerequisite_missing(monkeypatch): + experiment_id = uuid4() + row = SimpleNamespace( + prepared_payload={"alias": "experiment-1"}, + sample_accession=None, + project_accession=None, + ) + + monkeypatch.setattr( + broker, + "_find_latest_submission_for_validation", + lambda db_arg, entity_type_arg, entity_id_arg: row, + ) + monkeypatch.setattr( + broker, + "_lookup_taxonomy_for_entity", + lambda db_arg, entity_type_arg, entity_id_arg: ("org-1", 9606), + ) + + response = broker.validate_entity_submission( + payload=BrokerValidationRequest( + entity_type=BrokerEntityType.EXPERIMENT, + entity_id=experiment_id, + overrides={}, + ), + current_user=_broker_user(), + db=FakeSession(), + ) + + assert response.valid is False + assert response.resolved_prerequisites == {} + assert response.issues[0].field == "sample_accession" + + +def test_reports_attempt_acceptance(monkeypatch): + db = FakeSession() + attempt_id = uuid4() + entity_id = uuid4() + row = SimpleNamespace( + id=uuid4(), + sample_id=entity_id, + status="submitting", + attempt_id=attempt_id, + authority="ENA", + accession=None, + lock_acquired_at=datetime.now(timezone.utc), + lock_expires_at=datetime.now(timezone.utc), + finalised_attempt_id=None, + ) + + monkeypatch.setattr( + broker, + "_find_submission_for_attempt", + lambda db_arg, entity_type_arg, entity_id_arg, attempt_id_arg: row, + ) + monkeypatch.setattr( + broker, + "_register_submission_accession", + lambda db_arg, entity_type_arg, row_arg, accession_arg, secondary_accession=None: None, + ) + + response = broker.report_submission_outcomes( + attempt_id=attempt_id, + payload=BrokerReportRequest( + tax_id=9606, + results=[ + BrokerReportRecord( + entity_type=BrokerEntityType.SAMPLE, + entity_id=entity_id, + status="completed", + accession="ERS000001", + secondary_accession="SAMEA000001", + receipt_path=None, + message=None, + errors=[], + ) + ], + ), + current_user=_broker_user(), + db=db, + ) + + assert response.accepted is True + assert response.message == "reported" + assert row.status == "accepted" + assert row.attempt_id is None + + +def test_reports_attempt_rejection_creates_new_draft_submission(monkeypatch): + db = FakeSession() + attempt_id = uuid4() + entity_id = uuid4() + project_id = uuid4() + prepared_payload = {"alias": "sample-1"} + row = SimpleNamespace( + id=uuid4(), + sample_id=entity_id, + project_id=project_id, + status="submitting", + attempt_id=attempt_id, + authority="ENA", + accession=None, + prepared_payload=prepared_payload, + response_payload={"receipt_path": None}, + lock_acquired_at=datetime.now(timezone.utc), + lock_expires_at=datetime.now(timezone.utc), + finalised_attempt_id=None, + biosample_accession=None, + ) + + monkeypatch.setattr( + broker, + "_find_submission_for_attempt", + lambda db_arg, entity_type_arg, entity_id_arg, attempt_id_arg: row, + ) + monkeypatch.setattr( + broker, + "_register_submission_accession", + lambda db_arg, entity_type_arg, row_arg, accession_arg, secondary_accession=None: None, + ) + + response = broker.report_submission_outcomes( + attempt_id=attempt_id, + payload=BrokerReportRequest( + tax_id=9606, + results=[ + BrokerReportRecord( + entity_type=BrokerEntityType.SAMPLE, + entity_id=entity_id, + status="rejected", + receipt_path=None, + message="bad", + errors=[{"x": "y"}], + ) + ], + ), + current_user=_broker_user(), + db=db, + ) + + assert response.accepted is True + assert row.status == "rejected" + assert row.attempt_id is None + + new_rows = [obj for obj in db.added if isinstance(obj, SampleSubmission)] + assert len(new_rows) == 1 + new_row = new_rows[0] + assert new_row.sample_id == entity_id + assert new_row.project_id == project_id + assert new_row.status == "draft" + assert new_row.response_payload is None + + +def test_claim_ready_entities_no_claimable_entities_returns_empty_response(monkeypatch): + """Test that when organism exists but has no claimable entities, returns empty response instead of 404.""" + db = FakeSession() + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + monkeypatch.setattr( + broker, + "_get_organism_by_tax_id", + lambda db_arg, tax_id: SimpleNamespace(grouping_key="org-1"), + ) + # Mock all query functions to return empty lists + monkeypatch.setattr(broker, "_query_ready_project_submissions", lambda db_arg, tax_id: []) + monkeypatch.setattr(broker, "_query_ready_sample_submissions", lambda db_arg, tax_id: []) + monkeypatch.setattr(broker, "_query_ready_experiment_submissions", lambda db_arg, tax_id: []) + monkeypatch.setattr(broker, "_query_ready_run_submissions", lambda db_arg, tax_id: []) + + response = broker.claim_ready_entities( + payload=BrokerReadyClaimRequest(tax_id="9606"), + current_user=_broker_user(), + db=db, + ) + + assert response.tax_id == "9606" + assert response.scope == "full" + assert response.entities == [] + assert response.attempt_id is None + assert db.committed is False # No commit needed when returning early with empty response + + +def test_claim_ready_entities_organism_not_found_returns_404(monkeypatch): + """Test that when organism doesn't exist, returns 404.""" + db = FakeSession() + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + monkeypatch.setattr(broker, "_get_organism_by_tax_id", lambda db_arg, tax_id: None) + + with pytest.raises(HTTPException) as exc_info: + broker.claim_ready_entities( + payload=BrokerReadyClaimRequest(tax_id="99999"), + current_user=_broker_user(), + db=db, + ) + + assert exc_info.value.status_code == 404 + assert "Organism with tax_id 99999 not found" in exc_info.value.detail + + +def test_claims_batch_single_entity(monkeypatch): + """Test batch endpoint with a single entity (simplest case).""" + db = FakeSession() + sample_id = uuid4() + row = SimpleNamespace( + id=uuid4(), + sample_id=sample_id, + status="ready", + prepared_payload={"alias": "sample-1"}, + accession=None, + project_accession=None, + sample_accession=None, + experiment_accession=None, + ) + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + monkeypatch.setattr( + broker, + "_find_latest_claimable_entity_submission", + lambda db_arg, entity_type_arg, entity_id_arg: row, + ) + monkeypatch.setattr( + broker, + "_lookup_taxonomy_for_entity", + lambda db_arg, entity_type_arg, entity_id_arg: ("org-1", 9606), + ) + + response = broker.claim_batch_entities( + payload=BrokerBatchClaimRequest(sample_ids=[sample_id]), + current_user=_broker_user(), + db=db, + ) + + assert response.tax_id == "9606" + assert len(response.entities) == 1 + assert response.entities[0].type == BrokerEntityType.SAMPLE + assert response.entities[0].id == sample_id + assert row.status == "submitting" + assert db.committed is True + + +def test_claims_batch_multiple_entities_same_type(monkeypatch): + """Test batch endpoint with multiple entities of the same type.""" + db = FakeSession() + sample_id_1 = uuid4() + sample_id_2 = uuid4() + sample_id_3 = uuid4() + + rows = { + sample_id_1: SimpleNamespace( + id=uuid4(), + sample_id=sample_id_1, + status="ready", + prepared_payload={"alias": "sample-1"}, + accession=None, + project_accession=None, + sample_accession=None, + experiment_accession=None, + ), + sample_id_2: SimpleNamespace( + id=uuid4(), + sample_id=sample_id_2, + status="ready", + prepared_payload={"alias": "sample-2"}, + accession=None, + project_accession=None, + sample_accession=None, + experiment_accession=None, + ), + sample_id_3: SimpleNamespace( + id=uuid4(), + sample_id=sample_id_3, + status="ready", + prepared_payload={"alias": "sample-3"}, + accession=None, + project_accession=None, + sample_accession=None, + experiment_accession=None, + ), + } + + def mock_find_submission(db_arg, entity_type_arg, entity_id_arg): + return rows.get(entity_id_arg) + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + monkeypatch.setattr(broker, "_find_latest_claimable_entity_submission", mock_find_submission) + monkeypatch.setattr( + broker, + "_lookup_taxonomy_for_entity", + lambda db_arg, entity_type_arg, entity_id_arg: ("org-1", 9606), + ) + + response = broker.claim_batch_entities( + payload=BrokerBatchClaimRequest(sample_ids=[sample_id_1, sample_id_2, sample_id_3]), + current_user=_broker_user(), + db=db, + ) + + assert response.tax_id == "9606" + assert len(response.entities) == 3 + assert all(e.type == BrokerEntityType.SAMPLE for e in response.entities) + assert {e.id for e in response.entities} == {sample_id_1, sample_id_2, sample_id_3} + assert all(rows[sid].status == "submitting" for sid in [sample_id_1, sample_id_2, sample_id_3]) + assert db.committed is True + + +def test_claims_batch_multiple_entity_types(monkeypatch): + """Test batch endpoint with multiple entity types in one request.""" + db = FakeSession() + project_id = uuid4() + sample_id = uuid4() + experiment_id = uuid4() + run_id = uuid4() + + rows = { + (BrokerEntityType.PROJECT, project_id): SimpleNamespace( + id=uuid4(), + project_id=project_id, + status="ready", + prepared_payload={"alias": "project-1"}, + accession=None, + project_accession="PRJ000001", + ), + (BrokerEntityType.SAMPLE, sample_id): SimpleNamespace( + id=uuid4(), + sample_id=sample_id, + status="ready", + prepared_payload={"alias": "sample-1"}, + accession=None, + project_accession=None, + sample_accession=None, + experiment_accession=None, + ), + (BrokerEntityType.EXPERIMENT, experiment_id): SimpleNamespace( + id=uuid4(), + experiment_id=experiment_id, + status="ready", + prepared_payload={"alias": "experiment-1"}, + accession=None, + sample_accession="SAMEA000001", + project_accession=None, + experiment_accession=None, + ), + (BrokerEntityType.RUN, run_id): SimpleNamespace( + id=uuid4(), + read_id=run_id, + status="ready", + prepared_payload={ + "alias": "run-1", + "file_name": "reads.fastq.gz", + "file_format": "fastq", + }, + accession=None, + experiment_accession=None, + ), + } + + def mock_find_submission(db_arg, entity_type_arg, entity_id_arg): + return rows.get((entity_type_arg, entity_id_arg)) + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + monkeypatch.setattr(broker, "_find_latest_claimable_entity_submission", mock_find_submission) + monkeypatch.setattr( + broker, + "_lookup_taxonomy_for_entity", + lambda db_arg, entity_type_arg, entity_id_arg: ("org-1", 9606), + ) + + response = broker.claim_batch_entities( + payload=BrokerBatchClaimRequest( + project_ids=[project_id], + sample_ids=[sample_id], + experiment_ids=[experiment_id], + run_ids=[run_id], + ), + current_user=_broker_user(), + db=db, + ) + + assert response.tax_id == "9606" + assert len(response.entities) == 4 + entity_types = [e.type for e in response.entities] + assert BrokerEntityType.PROJECT in entity_types + assert BrokerEntityType.SAMPLE in entity_types + assert BrokerEntityType.EXPERIMENT in entity_types + assert BrokerEntityType.RUN in entity_types + assert db.committed is True + + +def test_claims_batch_multi_organism_returns_null_tax_id(monkeypatch): + """Test batch endpoint with entities from different organisms returns null tax_id.""" + db = FakeSession() + sample_id_1 = uuid4() + sample_id_2 = uuid4() + + rows = { + sample_id_1: SimpleNamespace( + id=uuid4(), + sample_id=sample_id_1, + status="ready", + prepared_payload={"alias": "sample-1"}, + accession=None, + project_accession=None, + sample_accession=None, + experiment_accession=None, + ), + sample_id_2: SimpleNamespace( + id=uuid4(), + sample_id=sample_id_2, + status="ready", + prepared_payload={"alias": "sample-2"}, + accession=None, + project_accession=None, + sample_accession=None, + experiment_accession=None, + ), + } + + def mock_find_submission(db_arg, entity_type_arg, entity_id_arg): + return rows.get(entity_id_arg) + + def mock_lookup_taxonomy(db_arg, entity_type_arg, entity_id_arg): + # Different organisms for different samples + if entity_id_arg == sample_id_1: + return ("org-1", 9606) + else: + return ("org-2", 9685) + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + monkeypatch.setattr(broker, "_find_latest_claimable_entity_submission", mock_find_submission) + monkeypatch.setattr(broker, "_lookup_taxonomy_for_entity", mock_lookup_taxonomy) + + response = broker.claim_batch_entities( + payload=BrokerBatchClaimRequest(sample_ids=[sample_id_1, sample_id_2]), + current_user=_broker_user(), + db=db, + ) + + assert response.tax_id is None # Multi-organism batch + assert len(response.entities) == 2 + assert db.committed is True + + +def test_claims_batch_empty_request_returns_400(monkeypatch): + """Test batch endpoint with no entity IDs returns 400.""" + db = FakeSession() + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + + with pytest.raises(HTTPException) as exc_info: + broker.claim_batch_entities( + payload=BrokerBatchClaimRequest(), + current_user=_broker_user(), + db=db, + ) + + assert exc_info.value.status_code == 400 + assert "At least one entity ID must be provided" in exc_info.value.detail + + +def test_claims_batch_entity_not_found_returns_404(monkeypatch): + """Test batch endpoint with non-existent entity returns 404.""" + db = FakeSession() + sample_id = uuid4() + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + monkeypatch.setattr( + broker, + "_find_latest_claimable_entity_submission", + lambda db_arg, entity_type_arg, entity_id_arg: None, + ) + + with pytest.raises(HTTPException) as exc_info: + broker.claim_batch_entities( + payload=BrokerBatchClaimRequest(sample_ids=[sample_id]), + current_user=_broker_user(), + db=db, + ) + + assert exc_info.value.status_code == 404 + assert "No claimable submission found" in exc_info.value.detail + + +def test_claims_batch_entity_without_taxonomy_returns_404(monkeypatch): + """Test batch endpoint with entity that has no taxonomy returns 404.""" + db = FakeSession() + sample_id = uuid4() + row = SimpleNamespace( + id=uuid4(), + sample_id=sample_id, + status="ready", + prepared_payload={"alias": "sample-1"}, + accession=None, + project_accession=None, + sample_accession=None, + experiment_accession=None, + ) + + monkeypatch.setattr(broker, "expire_stale_leases", lambda db_arg: {}) + monkeypatch.setattr( + broker, + "_find_latest_claimable_entity_submission", + lambda db_arg, entity_type_arg, entity_id_arg: row, + ) + monkeypatch.setattr( + broker, + "_lookup_taxonomy_for_entity", + lambda db_arg, entity_type_arg, entity_id_arg: (None, None), + ) + + with pytest.raises(HTTPException) as exc_info: + broker.claim_batch_entities( + payload=BrokerBatchClaimRequest(sample_ids=[sample_id]), + current_user=_broker_user(), + db=db, + ) + + assert exc_info.value.status_code == 404 + assert "Entity not found" in exc_info.value.detail