diff --git a/app/api/v1/endpoints/broker.py b/app/api/v1/endpoints/broker.py index f2ffb61..754fc7d 100644 --- a/app/api/v1/endpoints/broker.py +++ b/app/api/v1/endpoints/broker.py @@ -181,6 +181,20 @@ def _lookup_scientific_name(db: Session, *, taxon_id: str | int) -> Optional[str return db.query(Organism.scientific_name).filter(Organism.taxon_id == tax_id_int).scalar() +def _lookup_project_id(db: Session, *, entity_id: UUID) -> Optional[UUID]: + project_id = db.query(Experiment.project_id).filter(Experiment.id == entity_id).scalar() + if not project_id: + return None + return project_id + + +def _lookup_sample_id(db: Session, *, entity_id: UUID) -> Optional[UUID]: + sample_id = db.query(Experiment.sample_id).filter(Experiment.id == entity_id).scalar() + if not sample_id: + return None + return sample_id + + def _create_new_draft_submission_after_rejection( db: Session, *, @@ -303,15 +317,26 @@ def _get_accession_for_entity( """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, + if entity_type == "sample": + return ( + db.query(AccessionRegistry.secondary_accession) + .filter( + AccessionRegistry.entity_type == entity_type, + AccessionRegistry.entity_id == entity_id, + AccessionRegistry.authority == authority, + ) + .scalar() + ) + else: + return ( + db.query(AccessionRegistry.accession) + .filter( + AccessionRegistry.entity_type == entity_type, + AccessionRegistry.entity_id == entity_id, + AccessionRegistry.authority == authority, + ) + .scalar() ) - .scalar() - ) def _extract_broker_prerequisites( @@ -384,6 +409,13 @@ def _build_contract_entity( files = _extract_run_files(prepared_payload) if entity_type == BrokerEntityType.RUN else None scientific_name = _lookup_scientific_name(db, taxon_id=taxon_id) + if entity_type == BrokerEntityType.EXPERIMENT: + project_id = _lookup_project_id(db, entity_id=entity_id) + sample_id = _lookup_sample_id(db, entity_id=entity_id) + else: + project_id = None + sample_id = None + if entity_type in (BrokerEntityType.EXPERIMENT, BrokerEntityType.SAMPLE): title = _derive_claim_title( db, entity_type=entity_type, entity_id=entity_id, taxon_id=taxon_id @@ -396,6 +428,8 @@ def _build_contract_entity( id=entity_id, taxon_id=str(taxon_id), scientific_name=scientific_name, + project_id=project_id, + sample_id=sample_id, payload=prepared_payload or None, prerequisites=prerequisites if _prerequisites_to_dict(prerequisites) else None, files=files, diff --git a/app/schemas/broker_contract.py b/app/schemas/broker_contract.py index 066f8c6..9c77af1 100644 --- a/app/schemas/broker_contract.py +++ b/app/schemas/broker_contract.py @@ -36,6 +36,8 @@ class BrokerClaimEntity(BaseModel): id: UUID taxon_id: str scientific_name: Optional[str] = None + project_id: Optional[UUID] = None + sample_id: Optional[UUID] = None payload: Optional[Dict[str, Any]] = None prerequisites: Optional[BrokerPrerequisites] = None files: Optional[List[BrokerFileMetadata]] = None