Merge pull request #1701 from microbiomedata/1693-wfe-dgn-many-to-many

Update data model for multivalued `was_informed_by`

Author: eecavanna

nmdc_server/fakes.py

@@ -203,7 +203,7 @@ class PipelineStepBase(SQLAlchemyModelFactory):
     started_at_time: datetime = Faker("date_time")
     ended_at_time: datetime = Faker("date_time")
     execution_resource: str = Faker("word")
-    omics_processing: models.OmicsProcessing = SubFactory(OmicsProcessingFactory)
+    was_informed_by: List[models.OmicsProcessing] = []
 
 
 class ReadsQCFactory(PipelineStepBase):

nmdc_server/filters.py

@@ -35,6 +35,7 @@
     MetaPGeneFunction,
     MetaTGeneFunction,
     Table,
+    _table_model_map,
     workflow_execution_tables,
 )
 
@@ -107,7 +108,13 @@ def join(self, target_table: Table, query: Query) -> Query:
         raise NotImplementedError()
 
     def _join_omics_processing_related_tables(self, target_table: Table, query: Query) -> Query:
-        if target_table != Table.omics_processing:
+        if target_table in workflow_execution_tables:
+            association_table = models.workflow_activity_to_data_generation_map[target_table.value]
+            query = query.join(association_table).join(
+                models.OmicsProcessing,
+                models.OmicsProcessing.id == association_table.c.data_generation_id,
+            )
+        elif target_table != Table.omics_processing:
             query = query.join(models.OmicsProcessing)
 
         return self.join_omics_processing(query)
@@ -252,10 +259,25 @@ def join_study(self, query: Query) -> Query:
         )
 
 
-workflow_filter_classes: List[Type[OmicsProcessingFilter]] = []
+class WorkflowExecutionFilter(OmicsProcessingFilter):
+    table = Table.reads_qc
+
+    def join_omics_processing(self, query: Query) -> Query:
+        association_table = models.workflow_activity_to_data_generation_map[self.table.value]
+        model = _table_model_map[self.table]
+        q = query.join(
+            association_table,
+            association_table.c.data_generation_id == models.OmicsProcessing.id,
+        ).join(
+            model, model.id == association_table.c[f"{self.table.value}_id"]  # type: ignore
+        )
+        return q
+
+
+workflow_filter_classes: List[Type[WorkflowExecutionFilter]] = []
 for table in workflow_execution_tables:
     workflow_filter_classes.append(
-        type(f"{table.value}_filter", (OmicsProcessingFilter,), {"table": table})
+        type(f"{table.value}_filter", (WorkflowExecutionFilter,), {"table": table})
     )
 
 
@@ -274,10 +296,19 @@ def join(self, target_table: Table, query: Query) -> Query:
             )
 
         query = super().join(target_table, query)
+        # Use the association table to join from OmicsProcessing/DataGeneration to
+        # MetagenomeAnnotation. Due to how the association table(s) are generated
+        # dynamically, mypy does not know what the columns are.
+        association_table = models.metagenome_annotation_data_generation_association
         return (
             query.join(
+                association_table,
+                association_table.data_generation_id == models.OmicsProcessing.id,  # type: ignore
+            )
+            .join(
                 models.MetagenomeAnnotation,
-                models.MetagenomeAnnotation.omics_processing_id == models.OmicsProcessing.id,
+                models.MetagenomeAnnotation.id
+                == association_table.metagenome_annotation_id,  # type: ignore
             )
             .join(
                 models.MGAGeneFunctionAggregation,
@@ -325,10 +356,16 @@ def join(self, target_table: Table, query: Query) -> Query:
             )
 
         query = super().join(target_table, query)
+        association_table = models.metaproteomic_analysis_data_generation_association
         return (
             query.join(
+                association_table,
+                association_table.data_generation_id == models.OmicsProcessing.id,  # type: ignore
+            )
+            .join(
                 models.MetaproteomicAnalysis,
-                models.MetaproteomicAnalysis.omics_processing_id == models.OmicsProcessing.id,
+                models.MetaproteomicAnalysis.id
+                == association_table.metaproteomic_analysis_id,  # type: ignore
             )
             .join(
                 models.MetaPGeneFunctionAggregation,
@@ -359,10 +396,16 @@ def join(self, target_table: Table, query: Query) -> Query:
                 MetaTGeneFunction.id == models.MetaTGeneFunctionAggregation.gene_function_id,
             )
         query = super().join(target_table, query)
+        association_table = models.metatranscriptome_annotation_data_generation_association
         return (
             query.join(
+                association_table,
+                association_table.data_generation_id == models.OmicsProcessing.id,  # type: ignore
+            )
+            .join(
                 models.MetatranscriptomeAnnotation,
-                models.MetatranscriptomeAnnotation.omics_processing_id == models.OmicsProcessing.id,
+                models.MetatranscriptomeAnnotation.id
+                == association_table.metatranscriptome_annotation_id,  # type: ignore
             )
             .join(
                 models.MetaTGeneFunctionAggregation,
@@ -383,7 +426,14 @@ class MetaproteomicAnalysisFilter(OmicsProcessingFilter):
     table = Table.metaproteomic_analysis
 
     def join_omics_processing(self, query: Query) -> Query:
-        return query.join(self.table.model)
+        association_table = models.metaproteomic_analysis_data_generation_association
+        return query.join(
+            association_table,
+            association_table.c.data_generation_id == models.OmicsProcessing.id,
+        ).join(
+            models.MetaproteomicAnalysis,
+            models.MetaproteomicAnalysis.id == association_table.c.metaproteomic_analysis_id,
+        )
 
     def join_biosample(self, query: Query) -> Query:
         return (

nmdc_server/ingest/all.py

@@ -101,6 +101,14 @@ def load(db: Session, function_limit=None, skip_annotation=False):
     )
     db.commit()
 
+    # Update the FK relationship from the data_object table to the
+    # omics_processing table.
+    logger.info("Updating foreign key relationship from Data Object to Data Generation")
+    data_object.update_data_generation_relation(
+        db,
+        mongodb["data_object_set"].find(),
+    )
+
     workflow_set = "workflow_execution_set"
 
     logger.info("Loading metabolomics analysis...")

nmdc_server/ingest/data_object.py

@@ -5,7 +5,7 @@
 from sqlalchemy.orm import Session
 
 from nmdc_server.logger import get_logger
-from nmdc_server.models import DataObject
+from nmdc_server.models import DataObject, OmicsProcessing
 from nmdc_server.schemas import DataObjectCreate
 
 file_type_map: Dict[str, Tuple[str, str]] = {}
@@ -43,3 +43,23 @@ def load(db: Session, cursor: Cursor, file_types: List[Dict[str, Any]]):
 
     if objects_without_type:
         logger.error(f"Encountered {objects_without_type} objects without data_object_type")
+
+
+def update_data_generation_relation(db: Session, cursor: Cursor):
+    """
+    Update DataObject's omics_processing_id FK.
+
+    This should run after ingesting all data objects and data generations (omics processing).
+    """
+    for data_object in cursor:
+        id = data_object["id"]
+        was_generated_by = data_object.pop("was_generated_by", None)
+        if not was_generated_by:
+            continue
+        # Mypy does not like db.get, and reports that "Session" has no attribute "get."
+        # See https://docs.sqlalchemy.org/en/14/orm/session_basics.html#get-by-primary-key
+        data_generation = db.get(OmicsProcessing, was_generated_by)  # type: ignore
+        row = db.get(DataObject, id)  # type: ignore
+        if row and data_generation:
+            row.omics_processing_id = was_generated_by
+            db.add(row)

nmdc_server/ingest/pipeline.py

@@ -238,7 +238,11 @@ def load(
             if reported_type != workflow_type:
                 logger.warning(f"Unexpected type {reported_type} (expected {workflow_type})")
 
-        obj["omics_processing_id"] = obj.pop("was_informed_by")
+        was_informed_by: str | list[str] = obj.pop("was_informed_by")
+        if isinstance(was_informed_by, str):
+            was_informed_by = [was_informed_by]
+        obj["omics_processing_id"] = was_informed_by[0]
+        obj["was_informed_by"] = was_informed_by
 
         # TODO: pydantic should parse datetime like this... need to look into it
         #   2021-01-26T21:36:26.759770Z+0000
@@ -267,6 +271,7 @@ def load(
 
         input_association = getattr(models, f"{table_name}_input_association")
         output_association = getattr(models, f"{table_name}_output_association")
+        was_informed_by_association = getattr(models, f"{table_name}_data_generation_association")
 
         # TODO: Find a different way to validate ref. integrity
         valid_inputs = [d for d in inputs if db.query(models.DataObject).get(d)]
@@ -296,12 +301,12 @@ def load(
                 .values([(id_, f) for f in outputs])
                 .on_conflict_do_nothing()
             )
-
-        db.execute(
-            models.DataObject.__table__.update()
-            .where(models.DataObject.id.in_(inputs + outputs))
-            .values({"omics_processing_id": pipeline.omics_processing_id})
-        )
+        if was_informed_by:
+            db.execute(
+                insert(was_informed_by_association)
+                .values([(id_, data_generation) for data_generation in was_informed_by])
+                .on_conflict_do_nothing()
+            )
 
         for id_ in outputs:
             output = db.query(models.DataObject).get(id_)