diff --git a/alembic/versions/0002_update_taxon_cols.py b/alembic/versions/0002_update_taxon_cols.py new file mode 100644 index 0000000..67dd0fd --- /dev/null +++ b/alembic/versions/0002_update_taxon_cols.py @@ -0,0 +1,52 @@ +"""Remove outdated columns from taxonomy_info table. + +Revision ID: 0002_update_taxon_cols +Revises: 0001_initial_schema +Create Date: 2026-05-11 + +Removing columns "defined_class" (redundant with "ncbi_class"), "mito_ref" (redundant with "mitohifi_reference_species"), and +"busco_dataset_name" (replaced by "busco_odb10_dataset_name" and "busco_odb12_dataset_name"). Also adding "hic_specimen_sample_ids" +column to assembly table, to replace "hic_specimen_sample_id". +""" + +import sqlalchemy as sa +from sqlalchemy.dialects.postgresql import JSONB + +from alembic import op + +revision = "0002_update_taxon_cols" +down_revision = "0001_initial_schema" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.drop_column("taxonomy_info", "defined_class") + op.execute( + sa.text( + """ + UPDATE taxonomy_info + SET mitohifi_reference_species = COALESCE(mitohifi_reference_species, mito_ref) + WHERE mito_ref IS NOT NULL + """ + ) + ) + op.drop_column("taxonomy_info", "mito_ref") + op.drop_column("taxonomy_info", "busco_dataset_name") + op.add_column("assembly", sa.Column("hic_specimen_sample_ids", JSONB, nullable=True)) + + +def downgrade() -> None: + op.add_column( + "taxonomy_info", + sa.Column("defined_class", sa.Text(), nullable=True), + ) + op.add_column( + "taxonomy_info", + sa.Column("mito_ref", sa.Text(), nullable=True), + ) + op.add_column( + "taxonomy_info", + sa.Column("busco_dataset_name", sa.Text(), nullable=True), + ) + op.drop_column("assembly", "hic_specimen_sample_ids") diff --git a/app/api/v1/endpoints/assemblies.py b/app/api/v1/endpoints/assemblies.py index 76802f7..2ddc99c 100644 --- a/app/api/v1/endpoints/assemblies.py +++ b/app/api/v1/endpoints/assemblies.py @@ -323,7 +323,13 @@ def get_assembly_manifest( detail="No manifest stored for this assembly. Re-submit an intent to generate one.", ) - return JSONResponse(content=assembly.manifest_json) + return JSONResponse( + content={ + "assembly_id": str(assembly.id), + "version": assembly.version, + "manifest": assembly.manifest_json, + } + ) @router.post("/intent/{taxon_id}") @@ -360,12 +366,15 @@ def create_assembly_intent( db, intent_in.long_read_specimen_sample_id, org_taxon_id, "long_read_specimen_sample_id" ) - # 3. Validate hic_specimen_sample_id if provided - hic_sample = None - if intent_in.hic_specimen_sample_id: - hic_sample = _validate_specimen_sample( - db, intent_in.hic_specimen_sample_id, org_taxon_id, "hic_specimen_sample_id" - ) + # 3. Validate hic_specimen_sample_ids if provided + hic_samples: List[Sample] = [] + if intent_in.hic_specimen_sample_ids: + for idx, hic_sid in enumerate(intent_in.hic_specimen_sample_ids): + hic_samples.append( + _validate_specimen_sample( + db, hic_sid, org_taxon_id, f"hic_specimen_sample_ids[{idx}]" + ) + ) # 4. Fetch long-read experiments (PacBio or ONT only) and their reads long_read_lineage_sample_ids = _get_lineage_sample_ids_for_specimen(db, long_read_sample) @@ -391,16 +400,16 @@ def create_assembly_intent( long_read_exp_ids = [e.id for e in long_read_experiments] long_reads = db.query(Read).filter(Read.experiment_id.in_(long_read_exp_ids)).all() - # 5. Fetch Hi-C experiments and reads if hic_sample is provided + # 5. Fetch Hi-C experiments and reads for each hic_sample hic_experiments: List[Experiment] = [] hic_reads: List[Read] = [] hic_sample_id_map: Dict[str, str] = {} - if hic_sample: + for hic_sample in hic_samples: hic_lineage_sample_ids = _get_lineage_sample_ids_for_specimen(db, hic_sample) - hic_sample_id_map = { - str(sample_id): str(hic_sample.id) for sample_id in hic_lineage_sample_ids - } - hic_experiments = ( + hic_sample_id_map.update( + {str(sample_id): str(hic_sample.id) for sample_id in hic_lineage_sample_ids} + ) + hic_exps = ( db.query(Experiment) .filter( Experiment.sample_id.in_(hic_lineage_sample_ids), @@ -408,18 +417,17 @@ def create_assembly_intent( ) .all() ) - hic_experiments = [ - e for e in hic_experiments if (e.library_strategy or "").upper() == "HI-C" - ] - if not hic_experiments: + hic_exps = [e for e in hic_exps if (e.library_strategy or "").upper() == "HI-C"] + if not hic_exps: raise AppError( status_code=422, code="no_hic_experiments", - message="No Hi-C experiments (ILLUMINA + Hi-C) found for hic_specimen_sample_id", + message="No Hi-C experiments (ILLUMINA + Hi-C) found for hic_specimen_sample_ids", details={"hic_specimen_sample_id": str(hic_sample.id)}, ) - hic_exp_ids = [e.id for e in hic_experiments] - hic_reads = db.query(Read).filter(Read.experiment_id.in_(hic_exp_ids)).all() + hic_exp_ids = [e.id for e in hic_exps] + hic_experiments.extend(hic_exps) + hic_reads.extend(db.query(Read).filter(Read.experiment_id.in_(hic_exp_ids)).all()) # 6. Determine data_types from the relevant experiments # TODO review and remove this step now that we pass in the specimen_id @@ -439,7 +447,7 @@ def create_assembly_intent( db, taxon_id=org_taxon_id, long_read_specimen_sample_id=long_read_sample.id, - hic_specimen_sample_id=hic_sample.id if hic_sample else None, + hic_specimen_sample_ids=[s.id for s in hic_samples] if hic_samples else None, # TODO review if we need data_types when creating experiment data_types=data_types, tol_id=intent_in.tol_id, @@ -450,7 +458,7 @@ def create_assembly_intent( try: # 8. Build sample metadata and generate JSON manifest all_reads = long_reads + hic_reads - sample_metadata_by_id = _build_specimen_metadata(long_read_sample, hic_sample) + sample_metadata_by_id = _build_specimen_metadata(long_read_sample, *hic_samples) sequencing_sample_to_specimen_sample_id = { **long_read_sample_id_map, **hic_sample_id_map, @@ -462,16 +470,19 @@ def create_assembly_intent( experiments=all_experiments, # TODO tolid from sample (reported by broker) not input from caller tol_id=assembly.tol_id, + assembly_id=str(assembly.id), version=assembly.version, long_read_sample_id=long_read_sample.id, - hic_sample_id=hic_sample.id if hic_sample else None, + hic_sample_ids=[s.id for s in hic_samples] if hic_samples else None, sample_metadata_by_id=sample_metadata_by_id, sequencing_sample_to_specimen_sample_id=sequencing_sample_to_specimen_sample_id, ) # 9. Validate that the manifest contains eligible reads then persist it + # TODO decide whether to keep this validation. + """ long_read_keys = {"PACBIO_SMRT", "OXFORD_NANOPORE"} - if not any(k in manifest_data["reads"] for k in long_read_keys): + if not any(k in manifest_data["read_files"] for k in long_read_keys): raise AppError( status_code=422, code="no_eligible_long_reads", @@ -482,14 +493,14 @@ def create_assembly_intent( details={"long_read_specimen_sample_id": str(long_read_sample.id)}, ) - if hic_sample and "Hi-C" not in manifest_data["reads"]: + if hic_samples and "Hi-C" not in manifest_data["read_files"]: raise AppError( status_code=422, code="no_eligible_hic_reads", - message="No eligible Hi-C reads found for hic_specimen_sample_id", - details={"hic_specimen_sample_id": str(hic_sample.id)}, + message="No eligible Hi-C reads found for hic_specimen_sample_ids", + details={"hic_specimen_sample_ids": [str(s.id) for s in hic_samples]}, ) - + """ assembly.manifest_json = manifest_data db.add(assembly) db.commit() @@ -503,7 +514,6 @@ def create_assembly_intent( content={ "assembly_id": str(assembly.id), "version": assembly.version, - "status": assembly.status, "manifest": manifest_data, } ) @@ -556,9 +566,7 @@ def cancel_assembly_intent( "long_read_specimen_sample_id": str(assembly.long_read_specimen_sample_id) if assembly.long_read_specimen_sample_id else None, - "hic_specimen_sample_id": str(assembly.hic_specimen_sample_id) - if assembly.hic_specimen_sample_id - else None, + "hic_specimen_sample_ids": assembly.hic_specimen_sample_ids or [], "version": assembly.version, "status": assembly.status, } @@ -635,7 +643,13 @@ def get_manifest_by_assembly_id( detail="No manifest stored for this assembly. Re-submit an intent to generate one.", ) - return JSONResponse(content=assembly.manifest_json) + return JSONResponse( + content={ + "assembly_id": str(assembly.id), + "version": assembly.version, + "manifest": assembly.manifest_json, + } + ) @router.get("/{assembly_id}", response_model=AssemblySchema) diff --git a/app/api/v1/endpoints/taxonomy_info.py b/app/api/v1/endpoints/taxonomy_info.py index 8639694..b64e968 100644 --- a/app/api/v1/endpoints/taxonomy_info.py +++ b/app/api/v1/endpoints/taxonomy_info.py @@ -8,7 +8,11 @@ from app.core.policy import policy from app.models.taxonomy_info import TaxonomyInfo from app.models.user import User -from app.schemas.bulk_import import BulkImportResponse, BulkTaxonomyInfoImport +from app.schemas.bulk_import import ( + BulkImportResponse, + BulkNcbiRefreshRequest, + BulkTaxonomyInfoImport, +) from app.schemas.taxonomy_info import ( TaxonomyInfo as TaxonomyInfoSchema, ) @@ -50,6 +54,40 @@ def bulk_import_taxonomy_info( return taxonomy_info_service.bulk_import(db, data=data.root) +@router.post("/bulk-upsert", response_model=BulkImportResponse) +@policy("taxonomy_info:bulk_upsert") +def bulk_upsert_taxonomy_info( + *, + db: Session = Depends(get_db), + data: BulkTaxonomyInfoImport, + current_user: User = Depends(get_current_active_user), +) -> Any: + """ + Bulk upsert taxonomy info from a dictionary keyed by taxon_id. + + Inserts new rows and updates existing ones. Re-fetches NCBI data for all rows. + The taxon_id key must reference an existing organism. + """ + return taxonomy_info_service.bulk_upsert(db, data=data.root) + + +@router.post("/bulk-ncbi-refresh", response_model=BulkImportResponse) +@policy("taxonomy_info:bulk_ncbi_refresh") +def bulk_ncbi_refresh_taxonomy_info( + *, + db: Session = Depends(get_db), + data: BulkNcbiRefreshRequest, + current_user: User = Depends(get_current_active_user), +) -> Any: + """ + Re-sync NCBI taxonomy fields for existing taxonomy_info rows. + + Only updates ncbi_* columns — upstream fields are left unchanged. + Rows that do not yet have a taxonomy_info record are skipped. + """ + return taxonomy_info_service.bulk_ncbi_refresh(db, taxon_ids=data.taxon_ids) + + @router.get("/{taxon_id}", response_model=TaxonomyInfoSchema) def get_taxonomy_info( *, diff --git a/app/core/policy.py b/app/core/policy.py index 238fbb7..22524ad 100644 --- a/app/core/policy.py +++ b/app/core/policy.py @@ -66,6 +66,8 @@ "taxonomy_info:update": ["curator", "admin"], "taxonomy_info:delete": ["admin", "superuser"], "taxonomy_info:bulk_import": ["curator", "admin"], + "taxonomy_info:bulk_upsert": ["curator", "admin"], + "taxonomy_info:bulk_ncbi_refresh": ["curator", "admin"], } diff --git a/app/models/assembly.py b/app/models/assembly.py index 4198d2a..7ff016a 100644 --- a/app/models/assembly.py +++ b/app/models/assembly.py @@ -74,6 +74,7 @@ class Assembly(Base): UUID(as_uuid=True), ForeignKey("sample.id"), nullable=True ) hic_specimen_sample_id = Column(UUID(as_uuid=True), ForeignKey("sample.id"), nullable=True) + hic_specimen_sample_ids = Column(JSONB, nullable=True) manifest_json = Column(JSONB, nullable=True) # Relationships diff --git a/app/models/taxonomy_info.py b/app/models/taxonomy_info.py index 0263908..6e024c2 100644 --- a/app/models/taxonomy_info.py +++ b/app/models/taxonomy_info.py @@ -23,8 +23,6 @@ class TaxonomyInfo(Base): ncbi_tax_string = Column(Text, nullable=True) ncbi_full_lineage = Column(Text, nullable=True) ncbi_last_synced_at = Column(DateTime(timezone=True), nullable=True) - mito_ref = Column(Text, nullable=True) - busco_dataset_name = Column(Text, nullable=True) busco_odb10_dataset_name = Column(Text, nullable=True) busco_odb12_dataset_name = Column(Text, nullable=True) find_plastid = Column(Boolean, nullable=True) @@ -32,7 +30,6 @@ class TaxonomyInfo(Base): mitochondrial_genetic_code_id = Column(Integer, nullable=True) mitohifi_reference_species = Column(Text, nullable=True) oatk_hmm_name = Column(Text, nullable=True) - defined_class = Column(Text, nullable=True) augustus_dataset_name = Column(Text, nullable=True) genetic_code_id = Column(Integer, nullable=True) diff --git a/app/schemas/assembly.py b/app/schemas/assembly.py index 4eb5120..7e4653f 100644 --- a/app/schemas/assembly.py +++ b/app/schemas/assembly.py @@ -86,7 +86,7 @@ class AssemblyIntent(BaseModel): tol_id: Optional[str] = None long_read_specimen_sample_id: UUID - hic_specimen_sample_id: Optional[UUID] = None + hic_specimen_sample_ids: Optional[List[UUID]] = None class AssemblyIntentResponse(BaseModel): diff --git a/app/schemas/bulk_import.py b/app/schemas/bulk_import.py index 7e99051..fb7a3ff 100644 --- a/app/schemas/bulk_import.py +++ b/app/schemas/bulk_import.py @@ -42,10 +42,17 @@ class BulkTaxonomyInfoImport(RootModel[Dict[int, TaxonomyInfoUpdate]]): """Bulk taxonomy info import keyed by taxon_id.""" +class BulkNcbiRefreshRequest(BaseModel): + """Request body for bulk NCBI taxonomy refresh.""" + + taxon_ids: List[int] + + class BulkImportResponse(BaseModel): """Schema for bulk import response.""" created_count: int + updated_count: Optional[int] = 0 skipped_count: Optional[int] = 0 message: str errors: Optional[List[str]] = None # List of all errors with context diff --git a/app/schemas/taxonomy_info.py b/app/schemas/taxonomy_info.py index 43468eb..f5c8602 100644 --- a/app/schemas/taxonomy_info.py +++ b/app/schemas/taxonomy_info.py @@ -17,8 +17,6 @@ class TaxonomyInfoBase(BaseModel): ncbi_tax_string: Optional[str] = None ncbi_full_lineage: Optional[str] = None ncbi_last_synced_at: Optional[datetime] = None - mito_ref: Optional[str] = None - busco_dataset_name: Optional[str] = None busco_odb10_dataset_name: Optional[str] = None busco_odb12_dataset_name: Optional[str] = None find_plastid: Optional[bool] = None @@ -26,21 +24,17 @@ class TaxonomyInfoBase(BaseModel): mitochondrial_genetic_code_id: Optional[int] = None mitohifi_reference_species: Optional[str] = None oatk_hmm_name: Optional[str] = None - defined_class: Optional[str] = None augustus_dataset_name: Optional[str] = None genetic_code_id: Optional[int] = None class TaxonomyInfoWriteBase(BaseModel): - busco_dataset_name: Optional[str] = None busco_odb10_dataset_name: Optional[str] = None busco_odb12_dataset_name: Optional[str] = None find_plastid: Optional[bool] = None hic_motif: Optional[str] = None mitochondrial_genetic_code_id: Optional[int] = None - mitohifi_reference_species: Optional[str] = None oatk_hmm_name: Optional[str] = None - defined_class: Optional[str] = None augustus_dataset_name: Optional[str] = None genetic_code_id: Optional[int] = None diff --git a/app/services/assembly_helper.py b/app/services/assembly_helper.py index 9bc5698..7794444 100644 --- a/app/services/assembly_helper.py +++ b/app/services/assembly_helper.py @@ -128,9 +128,10 @@ def generate_assembly_manifest_json( reads: List[Read], experiments: List[Experiment], tol_id: str | None, + assembly_id: str, version: int, long_read_sample_id: UUID, - hic_sample_id: Optional[UUID] = None, + hic_sample_ids: Optional[List[UUID]] = None, sample_metadata_by_id: Dict[str, Dict[str, Any]] | None = None, sequencing_sample_to_specimen_sample_id: Dict[str, str] | None = None, ) -> Dict[str, Any]: @@ -138,7 +139,7 @@ def generate_assembly_manifest_json( Reads are routed to sections based on which specimen sample they belong to: - Reads from long_read_sample_id experiments → PACBIO_SMRT or OXFORD_NANOPORE section - - Reads from hic_sample_id experiments → Hi-C section (omitted when hic_sample_id is None) + - Reads from hic_sample_ids experiments → Hi-C section (omitted when hic_sample_ids is None) PacBio filtering: only files ending in .ccs.bam or hifi_reads.bam are included. ONT: all reads are included. @@ -147,12 +148,12 @@ def generate_assembly_manifest_json( Args: organism: Organism object taxonomy_info: TaxonomyInfo object related to the organism when available - reads: Combined list of Read objects from both specimen samples - experiments: Combined list of Experiment objects from both specimen samples + reads: Combined list of Read objects from all specimen samples + experiments: Combined list of Experiment objects from all specimen samples tol_id: ToL ID for the assembly (optional) version: Assembly version number long_read_sample_id: sample.id of the long-read specimen sample - hic_sample_id: sample.id of the Hi-C specimen sample (optional) + hic_sample_ids: sample.ids of the Hi-C specimen samples (optional) sample_metadata_by_id: Sample metadata keyed by sample.id as string (optional) Returns: @@ -166,7 +167,7 @@ def generate_assembly_manifest_json( logger.info("Total experiments: %d, Total reads: %d", len(experiments), len(reads)) long_read_sample_str = str(long_read_sample_id) - hic_sample_str = str(hic_sample_id) if hic_sample_id else None + hic_sample_id_strs = {str(sid) for sid in hic_sample_ids} if hic_sample_ids else set() # Build experiment info map: experiment.id → metadata exp_info_map: Dict[Any, Dict[str, Any]] = {} @@ -208,30 +209,32 @@ def generate_assembly_manifest_json( # Route by specimen sample, then by platform is_long_read_sample = resolved_sample_id == long_read_sample_str - is_hic_sample = hic_sample_str is not None and resolved_sample_id == hic_sample_str + is_hic_sample = bool(hic_sample_id_strs) and resolved_sample_id in hic_sample_id_strs if is_long_read_sample: if platform == "PACBIO_SMRT" and library_strategy in ("WGS", "WGA") and read.file_name: - if read.file_name.endswith(".ccs.bam") or read.file_name.endswith("hifi_reads.bam"): - logger.info("Adding PacBio read: %s", read.file_name) - if bpa_package_id not in pacbio_by_package: - entry: Dict[str, Any] = { - "sample_id": resolved_sample_id, - "bpa_sample_id": sample_meta.get("bpa_sample_id"), - "specimen_id": sample_meta.get("specimen_id"), - "resources": [], - } - if exp_info["bioplatforms_base_url"]: - entry["bioplatforms_base_url"] = exp_info["bioplatforms_base_url"] - pacbio_by_package[bpa_package_id] = entry - pacbio_by_package[bpa_package_id]["resources"].append( - {"md5sum": read.file_checksum, "url": read.bioplatforms_url} - ) + # if read.file_name.endswith(".ccs.bam") or read.file_name.endswith("hifi_reads.bam"): + logger.info("Adding PacBio read: %s", read.file_name) + if bpa_package_id not in pacbio_by_package: + entry: Dict[str, Any] = { + "sample_id": resolved_sample_id, + "bpa_sample_id": sample_meta.get("bpa_sample_id"), + "specimen_id": sample_meta.get("specimen_id"), + "single_end": [], + } + if exp_info["bioplatforms_base_url"]: + entry["bioplatforms_base_url"] = exp_info["bioplatforms_base_url"] + pacbio_by_package[bpa_package_id] = entry + pacbio_by_package[bpa_package_id]["single_end"].append( + {"md5sum": read.file_checksum, "url": read.bioplatforms_url} + ) + """ else: logger.debug( "Skipping PacBio read %s — not .ccs.bam or hifi_reads.bam", read.file_name, ) + """ elif platform == "OXFORD_NANOPORE" and library_strategy in ("WGS", "WGA"): logger.info("Adding ONT read: %s", read.file_name) if bpa_package_id not in ont_by_package: @@ -239,12 +242,12 @@ def generate_assembly_manifest_json( "sample_id": resolved_sample_id, "bpa_sample_id": sample_meta.get("bpa_sample_id"), "specimen_id": sample_meta.get("specimen_id"), - "resources": [], + "single_end": [], } if exp_info["bioplatforms_base_url"]: entry["bioplatforms_base_url"] = exp_info["bioplatforms_base_url"] ont_by_package[bpa_package_id] = entry - ont_by_package[bpa_package_id]["resources"].append( + ont_by_package[bpa_package_id]["single_end"].append( {"md5sum": read.file_checksum, "url": read.bioplatforms_url} ) @@ -265,11 +268,12 @@ def generate_assembly_manifest_json( "sample_id": resolved_sample_id, "bpa_sample_id": sample_meta.get("bpa_sample_id"), "specimen_id": sample_meta.get("specimen_id"), - "resources": {"r1": [], "r2": []}, + "r1": [], + "r2": [], } rkey = _normalize_read_number(read.read_number) if rkey in ("r1", "r2"): - hic_by_package[bpa_package_id]["resources"][rkey].append( + hic_by_package[bpa_package_id][rkey].append( { "url": read.bioplatforms_url, "md5sum": read.file_checksum, @@ -314,7 +318,8 @@ def generate_assembly_manifest_json( manifest = { "scientific_name": organism.scientific_name, "taxon_id": organism.taxon_id, - "tolid": tol_id, + "dataset_id": tol_id, + "assembly_id": assembly_id, "version": version, "busco_odb10_dataset_name": getattr(taxonomy_info, "busco_odb10_dataset_name", None), "busco_odb12_dataset_name": getattr(taxonomy_info, "busco_odb12_dataset_name", None), @@ -327,7 +332,7 @@ def generate_assembly_manifest_json( "oatk_hmm_name": getattr(taxonomy_info, "oatk_hmm_name", None), "augustus_dataset_name": getattr(taxonomy_info, "augustus_dataset_name", None), "genetic_code_id": getattr(taxonomy_info, "genetic_code_id", None), - "defined_class": getattr(taxonomy_info, "defined_class", None), - "reads": reads_section, + "ncbi_class": getattr(taxonomy_info, "ncbi_class", None), + "read_files": reads_section, } return manifest diff --git a/app/services/assembly_service.py b/app/services/assembly_service.py index a72ee17..432d8d3 100644 --- a/app/services/assembly_service.py +++ b/app/services/assembly_service.py @@ -212,7 +212,7 @@ def create_from_intent( *, taxon_id: int, long_read_specimen_sample_id: UUID, - hic_specimen_sample_id: Optional[UUID], + hic_specimen_sample_ids: Optional[List[UUID]], data_types: str, tol_id: Optional[str], project_id: Optional[UUID], @@ -221,7 +221,7 @@ def create_from_intent( """Create an Assembly at manifest-request time with status='requested'. Versioning is scoped by (taxon_id, long_read_specimen_sample_id) only — - data_types and hic_specimen_sample_id are excluded from the version key. + data_types and hic_specimen_sample_ids are excluded from the version key. sample_id is set to long_read_specimen_sample_id for backward compatibility. """ version = self.get_next_version_for_intent( @@ -229,11 +229,16 @@ def create_from_intent( taxon_id=taxon_id, long_read_specimen_sample_id=long_read_specimen_sample_id, ) + # Keep singular FK column pointing to the first HiC sample for backwards compatibility + hic_specimen_sample_id = hic_specimen_sample_ids[0] if hic_specimen_sample_ids else None assembly = Assembly( taxon_id=taxon_id, sample_id=long_read_specimen_sample_id, long_read_specimen_sample_id=long_read_specimen_sample_id, hic_specimen_sample_id=hic_specimen_sample_id, + hic_specimen_sample_ids=[str(sid) for sid in hic_specimen_sample_ids] + if hic_specimen_sample_ids + else None, data_types=data_types, version=version, tol_id=tol_id, diff --git a/app/services/ncbi_taxonomy_service.py b/app/services/ncbi_taxonomy_service.py index 998a9ce..32e303b 100644 --- a/app/services/ncbi_taxonomy_service.py +++ b/app/services/ncbi_taxonomy_service.py @@ -303,7 +303,7 @@ def extract_taxonomy_fields(report: dict[str, Any]) -> dict[str, Any]: full_lineage_future = executor.submit(process_parents, parents) mito_future = executor.submit(organelle_ref_lookup, parents, species_taxid, "Mitochondrion") ncbi_full_lineage = full_lineage_future.result() - mito_ref = mito_future.result() + mitohifi_reference_species = mito_future.result() return { "taxon_id": species_taxid, @@ -318,7 +318,7 @@ def extract_taxonomy_fields(report: dict[str, Any]) -> dict[str, Any]: "ncbi_lineage": lineage, "ncbi_tax_string": lineage_to_string(lineage), "ncbi_full_lineage": ncbi_full_lineage, - "mito_ref": mito_ref, + "mitohifi_reference_species": mitohifi_reference_species, } diff --git a/app/services/organism_service.py b/app/services/organism_service.py index dc24c3e..91e618d 100644 --- a/app/services/organism_service.py +++ b/app/services/organism_service.py @@ -288,12 +288,7 @@ def update_organism( if not organism: return None non_bpa_fields = [ - "ncbi_order", - "ncbi_family", - "busco_dataset_name", - "common_name", - "common_name_source", - "tax_string", + "scientific_name", ] new_bpa_json: Dict[str, Any] = organism.bpa_json or {} update_data = organism_in.model_dump(exclude_unset=True) diff --git a/app/services/taxonomy_info_service.py b/app/services/taxonomy_info_service.py index f40acd9..88cf9b0 100644 --- a/app/services/taxonomy_info_service.py +++ b/app/services/taxonomy_info_service.py @@ -160,13 +160,75 @@ def delete(self, db: Session, *, taxon_id: int) -> Optional[TaxonomyInfo]: if """ + def _bulk_process_rows( + self, + db: Session, + *, + candidates: List[ + tuple[int, Optional[TaxonomyInfoUpdate], Organism, Optional[TaxonomyInfo]] + ], + ncbi_by_taxon_id: Dict[int, Any], + apply_payload: bool = True, + ) -> tuple[int, int, int, List[str]]: + """Apply NCBI data (and optionally payload values) to a list of validated candidates. + + Returns (created_count, updated_count, skipped_count, errors). + """ + created_count = 0 + updated_count = 0 + skipped_count = 0 + errors: List[str] = [] + + for taxon_id, row, organism, existing_ti in candidates: + try: + is_new = existing_ti is None + ti = existing_ti if existing_ti else TaxonomyInfo(taxon_id=taxon_id) + if is_new: + db.add(ti) + + mapped = ncbi_by_taxon_id.get(taxon_id) + if mapped: + applied_fields = self._apply_ncbi_values(ti, mapped) + ti.ncbi_last_synced_at = datetime.now(timezone.utc) + logger.info( + "NCBI taxonomy enrichment %s taxonomy_info for taxon_id=%s; applied_fields=%s", + "created" if is_new else "updated", + taxon_id, + applied_fields, + ) + else: + logger.warning( + "NCBI enrichment returned no mapped taxonomy for taxon_id=%s", + taxon_id, + ) + + if apply_payload and row is not None: + self._apply_payload_values(ti, row.model_dump(exclude_unset=True)) + + sync_organism_scientific_name( + organism, + ncbi_scientific_name=ti.ncbi_scientific_name, + ) + db.add(organism) + db.commit() + + if is_new: + created_count += 1 + else: + updated_count += 1 + except Exception as e: + errors.append(f"{taxon_id}: {str(e)}") + db.rollback() + skipped_count += 1 + + return created_count, updated_count, skipped_count, errors + def bulk_import( self, db: Session, *, data: Dict[int, TaxonomyInfoUpdate] ) -> BulkImportResponse: - created_count = 0 skipped_count = 0 errors: List[str] = [] - candidates: List[tuple[int, TaxonomyInfoUpdate, Organism]] = [] + candidates: List[tuple[int, TaxonomyInfoUpdate, Organism, None]] = [] for taxon_id, row in data.items(): try: @@ -182,61 +244,126 @@ def bulk_import( ) skipped_count += 1 continue - except Exception as e: errors.append(f"{taxon_id}: {str(e)}") db.rollback() skipped_count += 1 continue - candidates.append((taxon_id, row, organism)) + candidates.append((taxon_id, row, organism, None)) scientific_names_by_taxon_id = { taxon_id: getattr(organism, "bpa_scientific_name", None) - for taxon_id, _, organism in candidates + for taxon_id, _, organism, _ in candidates } ncbi_by_taxon_id, unmapped = fetch_taxonomy_for_taxon_ids(scientific_names_by_taxon_id) if unmapped: logger.warning("NCBI bulk enrichment returned unmapped taxon_ids: %s", unmapped) - for taxon_id, row, organism in candidates: - try: - ti = TaxonomyInfo(taxon_id=taxon_id) - db.add(ti) + created_count, _, row_skipped, row_errors = self._bulk_process_rows( + db, candidates=candidates, ncbi_by_taxon_id=ncbi_by_taxon_id + ) + skipped_count += row_skipped + errors.extend(row_errors) - mapped = ncbi_by_taxon_id.get(taxon_id) - if mapped: - applied_fields = self._apply_ncbi_values(ti, mapped) - ti.ncbi_last_synced_at = datetime.now(timezone.utc) - logger.info( - "NCBI taxonomy enrichment %s taxonomy_info for taxon_id=%s; applied_fields=%s", - "created", - taxon_id, - applied_fields, - ) - else: - logger.warning( - "NCBI enrichment returned no mapped taxonomy during bulk import for taxon_id=%s", - taxon_id, - ) + return BulkImportResponse( + created_count=created_count, + skipped_count=skipped_count, + message=f"TaxonomyInfo import complete. Created: {created_count}, Skipped: {skipped_count}", + errors=errors if errors else None, + ) - self._apply_payload_values(ti, row.model_dump(exclude_unset=True)) - sync_organism_scientific_name( - organism, - ncbi_scientific_name=ti.ncbi_scientific_name, - ) - db.add(organism) - db.commit() - created_count += 1 + def bulk_upsert( + self, db: Session, *, data: Dict[int, TaxonomyInfoUpdate] + ) -> BulkImportResponse: + skipped_count = 0 + errors: List[str] = [] + candidates: List[tuple[int, TaxonomyInfoUpdate, Organism, Optional[TaxonomyInfo]]] = [] + + for taxon_id, row in data.items(): + try: + organism = db.query(Organism).filter(Organism.taxon_id == taxon_id).first() + if not organism: + errors.append(f"{taxon_id}: organism with taxon_id {taxon_id} does not exist") + skipped_count += 1 + continue + existing = db.query(TaxonomyInfo).filter(TaxonomyInfo.taxon_id == taxon_id).first() except Exception as e: errors.append(f"{taxon_id}: {str(e)}") db.rollback() skipped_count += 1 + continue + + candidates.append((taxon_id, row, organism, existing)) + + scientific_names_by_taxon_id = { + taxon_id: getattr(organism, "bpa_scientific_name", None) + for taxon_id, _, organism, _ in candidates + } + ncbi_by_taxon_id, unmapped = fetch_taxonomy_for_taxon_ids(scientific_names_by_taxon_id) + if unmapped: + logger.warning("NCBI bulk enrichment returned unmapped taxon_ids: %s", unmapped) + + created_count, updated_count, row_skipped, row_errors = self._bulk_process_rows( + db, candidates=candidates, ncbi_by_taxon_id=ncbi_by_taxon_id + ) + skipped_count += row_skipped + errors.extend(row_errors) return BulkImportResponse( created_count=created_count, + updated_count=updated_count, skipped_count=skipped_count, - message=f"TaxonomyInfo import complete. Created: {created_count}, Skipped: {skipped_count}", + message=f"TaxonomyInfo upsert complete. Created: {created_count}, Updated: {updated_count}, Skipped: {skipped_count}", + errors=errors if errors else None, + ) + + def bulk_ncbi_refresh(self, db: Session, *, taxon_ids: List[int]) -> BulkImportResponse: + skipped_count = 0 + errors: List[str] = [] + candidates: List[tuple[int, None, Organism, TaxonomyInfo]] = [] + + for taxon_id in taxon_ids: + try: + organism = db.query(Organism).filter(Organism.taxon_id == taxon_id).first() + if not organism: + errors.append(f"{taxon_id}: organism with taxon_id {taxon_id} does not exist") + skipped_count += 1 + continue + existing = db.query(TaxonomyInfo).filter(TaxonomyInfo.taxon_id == taxon_id).first() + if not existing: + errors.append( + f"{taxon_id}: taxonomy_info for taxon_id {taxon_id} does not exist" + ) + skipped_count += 1 + continue + except Exception as e: + errors.append(f"{taxon_id}: {str(e)}") + db.rollback() + skipped_count += 1 + continue + + candidates.append((taxon_id, None, organism, existing)) + + scientific_names_by_taxon_id = { + taxon_id: getattr(organism, "bpa_scientific_name", None) + for taxon_id, _, organism, _ in candidates + } + ncbi_by_taxon_id, unmapped = fetch_taxonomy_for_taxon_ids(scientific_names_by_taxon_id) + if unmapped: + logger.warning("NCBI bulk refresh returned unmapped taxon_ids: %s", unmapped) + + _, updated_count, row_skipped, row_errors = self._bulk_process_rows( + db, candidates=candidates, ncbi_by_taxon_id=ncbi_by_taxon_id, apply_payload=False + ) + skipped_count += row_skipped + errors.extend(row_errors) + + return BulkImportResponse( + created_count=0, + updated_count=updated_count, + skipped_count=skipped_count, + message=f"NCBI taxonomy refresh complete. Updated: {updated_count}, Skipped: {skipped_count}", errors=errors if errors else None, ) diff --git a/schema.sql b/schema.sql index ed60329..468a385 100644 --- a/schema.sql +++ b/schema.sql @@ -105,8 +105,6 @@ CREATE TABLE taxonomy_info ( ncbi_tax_string TEXT, ncbi_full_lineage TEXT, ncbi_last_synced_at TIMESTAMPTZ, - mito_ref TEXT, - busco_dataset_name TEXT, busco_odb10_dataset_name TEXT, busco_odb12_dataset_name TEXT, find_plastid BOOLEAN, @@ -114,7 +112,6 @@ CREATE TABLE taxonomy_info ( mitochondrial_genetic_code_id INTEGER, mitohifi_reference_species TEXT, oatk_hmm_name TEXT, - defined_class TEXT, augustus_dataset_name TEXT, genetic_code_id INTEGER ); diff --git a/tests/unit/endpoints/test_endpoints_assemblies.py b/tests/unit/endpoints/test_endpoints_assemblies.py index a86a764..92b94f3 100644 --- a/tests/unit/endpoints/test_endpoints_assemblies.py +++ b/tests/unit/endpoints/test_endpoints_assemblies.py @@ -808,7 +808,7 @@ def test_cancel_assembly_intent_success(monkeypatch): id=run_id, taxon_id=172942, long_read_specimen_sample_id=long_read_sample_id, - hic_specimen_sample_id=None, + hic_specimen_sample_ids=None, version=2, status="requested", ) diff --git a/tests/unit/endpoints/test_endpoints_taxonomy_info.py b/tests/unit/endpoints/test_endpoints_taxonomy_info.py index 50fcb27..236881e 100644 --- a/tests/unit/endpoints/test_endpoints_taxonomy_info.py +++ b/tests/unit/endpoints/test_endpoints_taxonomy_info.py @@ -29,9 +29,7 @@ def _make_ti(taxon_id=1, **kwargs): "find_plastid": None, "hic_motif": None, "mitochondrial_genetic_code_id": None, - "mitohifi_reference_species": None, "oatk_hmm_name": None, - "defined_class": None, "augustus_dataset_name": None, "genetic_code_id": None, } diff --git a/tests/unit/services/test_assembly_helper.py b/tests/unit/services/test_assembly_helper.py index 5b9afa2..35627b6 100644 --- a/tests/unit/services/test_assembly_helper.py +++ b/tests/unit/services/test_assembly_helper.py @@ -31,11 +31,9 @@ def _make_mock_taxonomy_info(): find_plastid=False, hic_motif="GATC", mitochondrial_genetic_code_id=1, - mitohifi_reference_species="test_species", oatk_hmm_name="test_hmm", augustus_dataset_name="test_augustus", genetic_code_id=1, - defined_class="test_class", ) @@ -753,11 +751,9 @@ def test_taxonomy_info_fields_included(self): assert result["find_plastid"] is False assert result["hic_motif"] == "GATC" assert result["mitochondrial_genetic_code_id"] == 1 - assert result["mitohifi_reference_species"] == "test_species" assert result["oatk_hmm_name"] == "test_hmm" assert result["augustus_dataset_name"] == "test_augustus" assert result["genetic_code_id"] == 1 - assert result["defined_class"] == "test_class" def test_sample_metadata_included_at_package_level(self): """Sample metadata (bpa_sample_id, specimen_id) appears at the package level.""" diff --git a/tests/unit/services/test_assembly_service.py b/tests/unit/services/test_assembly_service.py index 7d55de0..5281677 100644 --- a/tests/unit/services/test_assembly_service.py +++ b/tests/unit/services/test_assembly_service.py @@ -283,7 +283,7 @@ def test_creates_assembly_with_requested_status(self, mock_db, assembly_service) mock_db, taxon_id=172942, long_read_specimen_sample_id=LONG_READ_SAMPLE_ID, - hic_specimen_sample_id=None, + hic_specimen_sample_ids=None, data_types="PACBIO_SMRT", tol_id="tol-999", project_id=None, @@ -314,7 +314,7 @@ def capture_init(self, **kwargs): mock_db, taxon_id=172942, long_read_specimen_sample_id=LONG_READ_SAMPLE_ID, - hic_specimen_sample_id=None, + hic_specimen_sample_ids=None, data_types="PACBIO_SMRT", tol_id=None, project_id=None, @@ -342,7 +342,7 @@ def capture_init(self, **kwargs): mock_db, taxon_id=172942, long_read_specimen_sample_id=LONG_READ_SAMPLE_ID, - hic_specimen_sample_id=None, + hic_specimen_sample_ids=None, data_types="PACBIO_SMRT", tol_id=None, project_id=None, @@ -351,8 +351,8 @@ def capture_init(self, **kwargs): assert stored_manifests == [sample_manifest] - def test_hic_specimen_sample_id_stored(self, mock_db, assembly_service): - """hic_specimen_sample_id is stored on the Assembly.""" + def test_hic_specimen_sample_ids_stored(self, mock_db, assembly_service): + """hic_specimen_sample_ids is stored as a JSON list; first entry sets hic_specimen_sample_id FK.""" mock_query = Mock() mock_query.filter.return_value.scalar.return_value = None mock_db.query.return_value = mock_query @@ -360,23 +360,29 @@ def test_hic_specimen_sample_id_stored(self, mock_db, assembly_service): mock_db.commit = Mock() mock_db.refresh = Mock() - hic_ids = [] + captured = [] def capture_init(self, **kwargs): - hic_ids.append(kwargs.get("hic_specimen_sample_id")) + captured.append( + { + "hic_specimen_sample_id": kwargs.get("hic_specimen_sample_id"), + "hic_specimen_sample_ids": kwargs.get("hic_specimen_sample_ids"), + } + ) with patch.object(Assembly, "__init__", capture_init): assembly_service.create_from_intent( mock_db, taxon_id=172942, long_read_specimen_sample_id=LONG_READ_SAMPLE_ID, - hic_specimen_sample_id=HIC_SAMPLE_ID, + hic_specimen_sample_ids=[HIC_SAMPLE_ID], data_types="PACBIO_SMRT_HIC", tol_id=None, project_id=None, ) - assert hic_ids == [HIC_SAMPLE_ID] + assert captured[0]["hic_specimen_sample_id"] == HIC_SAMPLE_ID + assert captured[0]["hic_specimen_sample_ids"] == [str(HIC_SAMPLE_ID)] def test_sample_id_set_to_long_read_specimen_sample_id(self, mock_db, assembly_service): """sample_id is set to long_read_specimen_sample_id for backward compatibility.""" @@ -397,7 +403,7 @@ def capture_init(self, **kwargs): mock_db, taxon_id=172942, long_read_specimen_sample_id=LONG_READ_SAMPLE_ID, - hic_specimen_sample_id=None, + hic_specimen_sample_ids=None, data_types="PACBIO_SMRT", tol_id=None, project_id=None, diff --git a/tests/unit/services/test_taxonomy_info_service.py b/tests/unit/services/test_taxonomy_info_service.py index 3439dc3..8af7d39 100644 --- a/tests/unit/services/test_taxonomy_info_service.py +++ b/tests/unit/services/test_taxonomy_info_service.py @@ -77,7 +77,7 @@ def test_populate_from_ncbi_lookup_creates_taxonomy_info(monkeypatch): "ncbi_rank": "species", "ncbi_scientific_name": "Penicillium test", "ncbi_order": "Eurotiales", - "mito_ref": "Penicillium chrysogenum", + "mitohifi_reference_species": "Penicillium chrysogenum", } }, [], @@ -97,7 +97,7 @@ def test_populate_from_ncbi_lookup_creates_taxonomy_info(monkeypatch): assert ti.ncbi_taxon_id == 5077 assert ti.ncbi_rank == "species" assert ti.ncbi_order == "Eurotiales" - assert ti.mito_ref == "Penicillium chrysogenum" + assert ti.mitohifi_reference_species == "Penicillium chrysogenum" assert ti.ncbi_last_synced_at is not None assert organism.scientific_name == "Penicillium test" assert db.data[TaxonomyInfo][5077] is ti @@ -176,14 +176,14 @@ def test_bulk_import_batches_ncbi_lookup_and_creates_taxonomy_info(monkeypatch): "ncbi_taxon_id": 9612, "ncbi_rank": "species", "ncbi_scientific_name": "Canis lupus familiaris", - "mito_ref": "Canis lupus familiaris", + "mitohifi_reference_species": "Canis lupus familiaris", }, 9685: { "taxon_id": 9685, "ncbi_taxon_id": 9685, "ncbi_rank": "species", "ncbi_scientific_name": "Felis silvestris catus", - "mito_ref": "Felis silvestris catus", + "mitohifi_reference_species": "Felis silvestris catus", }, }, [], @@ -207,12 +207,12 @@ def test_bulk_import_batches_ncbi_lookup_and_creates_taxonomy_info(monkeypatch): saved_cat = db.data[TaxonomyInfo][9685] assert saved_dog.ncbi_taxon_id == 9612 assert saved_dog.ncbi_rank == "species" - assert saved_dog.mito_ref == "Canis lupus familiaris" + assert saved_dog.mitohifi_reference_species == "Canis lupus familiaris" assert saved_dog.ncbi_last_synced_at is not None assert saved_dog.genetic_code_id == 2 assert saved_cat.ncbi_taxon_id == 9685 assert saved_cat.ncbi_rank == "species" - assert saved_cat.mito_ref == "Felis silvestris catus" + assert saved_cat.mitohifi_reference_species == "Felis silvestris catus" assert saved_cat.ncbi_last_synced_at is not None assert saved_cat.genetic_code_id == 1 assert organisms[9612].scientific_name == "Canis lupus familiaris"