Merge pull request #1313 from microbiomedata/1307-metat-model-updates

Update ingest and data model to account for Metatranscriptomics

Author: naglepuff

nmdc_server/crud.py

@@ -51,6 +51,12 @@ def get_database_summary(db: Session) -> schemas.DatabaseSummary:
         reads_qc=aggregations.get_table_summary(db, models.ReadsQC),
         metagenome_assembly=aggregations.get_table_summary(db, models.MetagenomeAssembly),
         metagenome_annotation=aggregations.get_table_summary(db, models.MetagenomeAnnotation),
+        metatranscriptome_assembly=aggregations.get_table_summary(
+            db, models.MetatranscriptomeAssembly
+        ),
+        metatranscriptome_annotation=aggregations.get_table_summary(
+            db, models.MetatranscriptomeAnnotation
+        ),
         metaproteomic_analysis=aggregations.get_table_summary(db, models.MetaproteomicAnalysis),
         mags_analysis=aggregations.get_table_summary(db, models.MAGsAnalysis),
         read_based_analysis=aggregations.get_table_summary(db, models.ReadBasedAnalysis),

nmdc_server/data_object_filters.py

@@ -31,6 +31,8 @@ class WorkflowActivityTypeEnum(Enum):
     reads_qc = "nmdc:ReadQCAnalysisActivity"
     metagenome_assembly = "nmdc:MetagenomeAssembly"
     metagenome_annotation = "nmdc:MetagenomeAnnotation"  # TODO name out of date, fix
+    metatranscriptome_assembly = "nmdc:MetatranscriptomeAssembly"
+    metatranscriptome_annotation = "nmdc:MetatranscriptomeAnnotation"  # TODO name out of date, fix
     metaproteomic_analysis = "nmdc:MetaProteomicAnalysis"
     mags_analysis = "nmdc:MAGsAnalysisActivity"
     read_based_analysis = "nmdc:ReadbasedAnalysis"  # TODO name out of date, fix
@@ -52,6 +54,8 @@ def output_association(self):
     WorkflowActivityTypeEnum.reads_qc: models.ReadsQC,
     WorkflowActivityTypeEnum.metagenome_assembly: models.MetagenomeAssembly,
     WorkflowActivityTypeEnum.metagenome_annotation: models.MetagenomeAnnotation,
+    WorkflowActivityTypeEnum.metatranscriptome_assembly: models.MetatranscriptomeAssembly,
+    WorkflowActivityTypeEnum.metatranscriptome_annotation: models.MetatranscriptomeAnnotation,
     WorkflowActivityTypeEnum.metaproteomic_analysis: models.MetaproteomicAnalysis,
     WorkflowActivityTypeEnum.mags_analysis: models.MAGsAnalysis,
     WorkflowActivityTypeEnum.read_based_analysis: models.ReadBasedAnalysis,
@@ -67,6 +71,8 @@ def output_association(self):
     WorkflowActivityTypeEnum.reads_qc: models.reads_qc_output_association,
     WorkflowActivityTypeEnum.metagenome_assembly: models.metagenome_assembly_output_association,
     WorkflowActivityTypeEnum.metagenome_annotation: models.metagenome_annotation_output_association,
+    WorkflowActivityTypeEnum.metatranscriptome_assembly: models.metatranscriptome_assembly_output_association,  # noqa: E501
+    WorkflowActivityTypeEnum.metatranscriptome_annotation: models.metatranscriptome_annotation_output_association,  # noqa: E501
     _mpa: models.metaproteomic_analysis_output_association,
     WorkflowActivityTypeEnum.mags_analysis: models.mags_analysis_output_association,
     WorkflowActivityTypeEnum.read_based_analysis: models.read_based_analysis_output_association,

nmdc_server/filters.py

@@ -33,6 +33,7 @@
     EnvMediumAncestor,
     EnvMediumTerm,
     MetaPGeneFunction,
+    MetaTGeneFunction,
     Table,
     workflow_execution_tables,
 )
@@ -327,6 +328,40 @@ def join_self(self, query: Query, parent: Table) -> Query:
         return query
 
 
+class MetaTGeneFunctionFilter(OmicsProcessingFilter):
+    table = Table.metat_gene_function
+
+    def join(self, target_table: Table, query: Query) -> Query:
+        if target_table == Table.metatranscriptome_annotation:
+            return query.join(
+                models.MetaTGeneFunctionAggregation,
+                models.MetaTGeneFunctionAggregation.metatranscriptome_annotation_id
+                == models.MetatranscriptomeAnnotation.id,
+            ).join(
+                MetaTGeneFunction,
+                MetaTGeneFunction.id == models.MetaTGeneFunctionAggregation.gene_function_id,
+            )
+        query = super().join(target_table, query)
+        return (
+            query.join(
+                models.MetatranscriptomeAnnotation,
+                models.MetatranscriptomeAnnotation.omics_processing_id == models.OmicsProcessing.id,
+            )
+            .join(
+                models.MetaTGeneFunctionAggregation,
+                models.MetaTGeneFunctionAggregation.metatranscriptome_annotation_id
+                == models.MetatranscriptomeAnnotation.id,
+            )
+            .join(
+                MetaTGeneFunction,
+                MetaTGeneFunction.id == models.MetaTGeneFunctionAggregation.gene_function_id,
+            )
+        )
+
+    def join_self(self, query: Query, parent: Table) -> Query:
+        return query
+
+
 def _get_all_subclasses(cls: Type[BaseFilter]) -> List[Type[BaseFilter]]:
     all_subclasses: List[Type[BaseFilter]] = []
     for subclass in cls.__subclasses__():

nmdc_server/ingest/all.py

@@ -120,14 +120,22 @@ def load(db: Session, function_limit=None, skip_annotation=False):
     )
     db.commit()
 
-    logger.info("Loading metatranscriptome activities...")
+    logger.info("Loading metatranscriptome expression analyses...")
     pipeline.load(
         db,
-        mongodb["metatranscriptome_activity_set"].find(),
+        mongodb["metatranscriptome_expression_analysis_set"].find(),
         pipeline.load_metatranscriptome,
         WorkflowActivityTypeEnum.metatranscriptome.value,
     )
 
+    logger.info("Loading metatranscriptome assemblies...")
+    pipeline.load(
+        db,
+        mongodb["metatranscriptome_assembly_set"].find(),
+        pipeline.load_mt_assembly,
+        WorkflowActivityTypeEnum.metatranscriptome_assembly.value,
+    )
+
     logger.info("Loading NOM analysis...")
     pipeline.load(
         db,
@@ -170,13 +178,29 @@ def load(db: Session, function_limit=None, skip_annotation=False):
                     annotations=mongodb["functional_annotation_agg"],
                     function_limit=function_limit,
                 )
+
         except Exception:
-            logger.exception("Failed during metag ingest.")
+            logger.exception("Failed during annotation ingest.")
         finally:
             db.commit()
 
     else:
-        logger.info("Skipping annotation ingest")
+        logger.info("Skipping mg annotation ingest")
+
+    try:
+        logger.info("Loading metatranscriptome annotation...")
+        pipeline.load(
+            db,
+            mongodb["metatranscriptome_annotation_set"].find(),
+            pipeline.load_mt_annotation,
+            WorkflowActivityTypeEnum.metatranscriptome_annotation.value,
+            annotations=mongodb["functional_annotation_agg"],
+            function_limit=function_limit,
+        )
+    except Exception:
+        logger.exception("Failed during metatranscriptome annotation ingest.")
+    finally:
+        db.commit()
 
     logger.info("Loading read qc...")
     pipeline.load(

nmdc_server/ingest/pipeline.py

@@ -130,6 +130,55 @@ def load_mp_analysis(db: Session, obj: Dict[str, Any], **kwargs) -> LoadObjectRe
     return pipeline
 
 
+def load_mt_annotation(db: Session, obj: Dict[str, Any], **kwargs) -> LoadObjectReturn:
+    # Ingest the MetatranscriptomeAnnotation record
+    pipeline = cast(models.MetatranscriptomeAnnotation, load_mt_annotation_base(db, obj, **kwargs))
+
+    annotations: Collection = kwargs["annotations"]
+
+    # Query gene function annotations from mongo and build the appropriate objects
+    query = annotations.find(
+        {
+            "metagenome_annotation_id": pipeline.id,
+            "gene_function_id": {
+                "$regex": ko_regex,
+            },
+        },
+        no_cursor_timeout=True,
+        projection={
+            "_id": False,
+            "metatranscriptome_annotation_id": True,
+            "count": True,
+            "gene_function_id": True,
+        },
+    )
+    if kwargs.get("function_limit"):
+        query = query.limit(kwargs["function_limit"])
+
+    gene_functions: Set[str] = set()
+    gene_function_aggregations: List[models.MetaTGeneFunctionAggregation] = []
+    for annotation in query:
+        function_id = annotation["gene_function_id"]
+        gene_functions.add(function_id)
+        gene_function_aggregations.append(
+            models.MetaTGeneFunctionAggregation(
+                metatranscriptome_annotation_id=pipeline.id,
+                gene_function_id=function_id,
+                count=annotation["count"],
+            )
+        )
+    # Save both newly encountered gene functions and the gene function aggregations
+    if gene_function_aggregations:
+        db.execute(
+            insert(models.GeneFunction)
+            .on_conflict_do_nothing()
+            .values([(gf,) for gf in gene_functions])
+        )
+        db.bulk_save_objects(gene_function_aggregations)
+
+    return pipeline
+
+
 # This is a loader for a generic workflow type that doesn't need any
 # additional processing.
 def generate_pipeline_loader(schema, model) -> LoadObject:
@@ -161,6 +210,12 @@ def loader(db: Session, obj: Dict[str, Any], **kwargs: Any) -> LoadObjectReturn:
 load_metatranscriptome = generate_pipeline_loader(
     schemas.MetatranscriptomeBase, models.Metatranscriptome
 )
+load_mt_assembly = generate_pipeline_loader(
+    schemas.MetatranscriptomeAssemblyBase, models.MetatranscriptomeAssembly
+)
+load_mt_annotation_base = generate_pipeline_loader(
+    schemas.MetatranscriptomeAnnotationBase, models.MetatranscriptomeAnnotation
+)
 
 
 # This is a generic function for load workflow execution objects.  Some workflow types require

nmdc_server/migrations/versions/c0b36f8dc4b8_add_metat_gene_function.py

@@ -0,0 +1,52 @@
+"""Add table for metaT gene functions
+
+Revision ID: c0b36f8dc4b8
+Revises: d3d563e01a74
+Create Date: 2024-07-25 14:24:16.657561
+
+"""
+
+from typing import Optional
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "c0b36f8dc4b8"
+down_revision: Optional[str] = "d3d563e01a74"
+branch_labels: Optional[str] = None
+depends_on: Optional[str] = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "metat_gene_function_aggregation",
+        sa.Column("metatranscriptome_annotation_id", sa.String(), nullable=False),
+        sa.Column("gene_function_id", sa.String(), nullable=False),
+        sa.Column("count", sa.BigInteger(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["gene_function_id"],
+            ["gene_function.id"],
+            name=op.f("fk_metat_gene_function_aggregation_gene_function_id_gene_function"),
+        ),
+        sa.ForeignKeyConstraint(
+            ["metatranscriptome_annotation_id"],
+            ["metatranscriptome_annotation.id"],
+            name=op.f(
+                "fk_metat_gene_function_aggregation_metatranscriptome_annotation_id_metatranscriptome_annotation"  # noqa: E501
+            ),
+        ),
+        sa.PrimaryKeyConstraint(
+            "metatranscriptome_annotation_id",
+            "gene_function_id",
+            name=op.f("pk_metat_gene_function_aggregation"),
+        ),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("metat_gene_function_aggregation")
+    # ### end Alembic commands ###