Custom cohorts in metamist

README.md

@@ -36,9 +36,9 @@ It comprises three key components:
 
 As of Jan 15, 2024 this schema should reflect the data structure on the tables:
 
-![Database Structure](resources/2024-01-15_db-diagram.png)
+![Database Structure](resources/schemav7.7.png.png)
 
-You can also find this at [DbDiagram](https://dbdiagram.io/d/Metamist-Schema-v6-6-2-65a48ac7ac844320aee60d16).
+You can also find this at [DbDiagram](https://dbdiagram.io/d/Metamist-Schema-v7-7-6600c875ae072629ced6a1fc).
 
 The codebase contains the following modules worth noting:
 

api/graphql/filters.py

@@ -18,6 +18,8 @@ class GraphQLFilter(Generic[T]):
     gte: T | None = None
     lt: T | None = None
     lte: T | None = None
+    contains: T | None = None
+    icontains: T | None = None
 
     def all_values(self):
         """
@@ -38,6 +40,10 @@ def all_values(self):
             v.append(self.lt)
         if self.lte:
             v.append(self.lte)
+        if self.contains:
+            v.append(self.contains)
+        if self.icontains:
+            v.append(self.icontains)
 
         return v
 
@@ -53,6 +59,8 @@ def to_internal_filter(self, f: Callable[[T], Any] = None):
                 gte=f(self.gte) if self.gte else None,
                 lt=f(self.lt) if self.lt else None,
                 lte=f(self.lte) if self.lte else None,
+                contains=f(self.contains) if self.contains else None,
+                icontains=f(self.icontains) if self.icontains else None,
             )
 
         return GenericFilter(
@@ -63,6 +71,8 @@ def to_internal_filter(self, f: Callable[[T], Any] = None):
             gte=self.gte,
             lt=self.lt,
             lte=self.lte,
+            contains=self.contains,
+            icontains=self.icontains,
         )
 
 

api/graphql/loaders.py

@@ -364,10 +364,13 @@ async def load_projects_for_ids(project_ids: list[int], connection) -> list[Proj
     """
     pttable = ProjectPermissionsTable(connection)
     projects = await pttable.get_and_check_access_to_projects_for_ids(
-        user=connection.user, project_ids=project_ids, readonly=True
+        user=connection.author, project_ids=project_ids, readonly=True
     )
+
     p_by_id = {p.id: p for p in projects}
-    return [p_by_id.get(p) for p in project_ids]
+    projects = [p_by_id.get(p) for p in project_ids]
+
+    return [p for p in projects if p is not None]
 
 
 @connected_data_loader(LoaderKeys.FAMILIES_FOR_PARTICIPANTS)

api/graphql/schema.py

@@ -1,6 +1,6 @@
 # type: ignore
 # flake8: noqa
-# pylint: disable=no-value-for-parameter,redefined-builtin,missing-function-docstring,unused-argument
+# pylint: disable=no-value-for-parameter,redefined-builtin,missing-function-docstring,unused-argument,too-many-lines
 """
 Schema for GraphQL.
 
@@ -22,13 +22,15 @@
     AnalysisLayer,
     AnalysisRunnerLayer,
     AssayLayer,
+    CohortLayer,
     FamilyLayer,
     SampleLayer,
     SequencingGroupLayer,
 )
 from db.python.tables.analysis import AnalysisFilter
 from db.python.tables.analysis_runner import AnalysisRunnerFilter
 from db.python.tables.assay import AssayFilter
+from db.python.tables.cohort import CohortFilter, CohortTemplateFilter
 from db.python.tables.project import ProjectPermissionsTable
 from db.python.tables.sample import SampleFilter
 from db.python.tables.sequencing_group import SequencingGroupFilter
@@ -38,6 +40,8 @@
     AnalysisInternal,
     AssayInternal,
     AuditLogInternal,
+    CohortInternal,
+    CohortTemplateInternal,
     FamilyInternal,
     ParticipantInternal,
     Project,
@@ -46,6 +50,11 @@
 )
 from models.models.analysis_runner import AnalysisRunnerInternal
 from models.models.sample import sample_id_transform_to_raw
+from models.utils.cohort_id_format import cohort_id_format, cohort_id_transform_to_raw
+from models.utils.cohort_template_id_format import (
+    cohort_template_id_format,
+    cohort_template_id_transform_to_raw,
+)
 from models.utils.sample_id_format import sample_id_format
 from models.utils.sequencing_group_id_format import (
     sequencing_group_id_format,
@@ -73,6 +82,87 @@ async def m(info: Info) -> list[str]:
 GraphQLEnum = strawberry.type(type('GraphQLEnum', (object,), enum_methods))
 
 
+# Create cohort GraphQL model
+@strawberry.type
+class GraphQLCohort:
+    """Cohort GraphQL model"""
+
+    id: str
+    name: str
+    description: str
+    author: str
+
+    @staticmethod
+    def from_internal(internal: CohortInternal) -> 'GraphQLCohort':
+        return GraphQLCohort(
+            id=cohort_id_format(internal.id),
+            name=internal.name,
+            description=internal.description,
+            author=internal.author,
+        )
+
+    @strawberry.field()
+    async def template(self, info: Info, root: 'Cohort') -> 'GraphQLCohortTemplate':
+        connection = info.context['connection']
+        template = await CohortLayer(connection).get_template_by_cohort_id(
+            cohort_id_transform_to_raw(root.id)
+        )
+
+        return GraphQLCohortTemplate.from_internal(template)
+
+    @strawberry.field()
+    async def sequencing_groups(
+        self, info: Info, root: 'Cohort'
+    ) -> list['GraphQLSequencingGroup']:
+        connection = info.context['connection']
+        cohort_layer = CohortLayer(connection)
+        sg_ids = await cohort_layer.get_cohort_sequencing_group_ids(
+            cohort_id_transform_to_raw(root.id)
+        )
+
+        sg_layer = SequencingGroupLayer(connection)
+        sequencing_groups = await sg_layer.get_sequencing_groups_by_ids(sg_ids)
+        return [GraphQLSequencingGroup.from_internal(sg) for sg in sequencing_groups]
+
+    @strawberry.field()
+    async def analyses(self, info: Info, root: 'Cohort') -> list['GraphQLAnalysis']:
+        connection = info.context['connection']
+        connection.project = root.project
+        internal_analysis = await AnalysisLayer(connection).query(
+            AnalysisFilter(
+                cohort_id=GenericFilter(in_=[cohort_id_transform_to_raw(root.id)]),
+            )
+        )
+        return [GraphQLAnalysis.from_internal(a) for a in internal_analysis]
+
+    @strawberry.field()
+    async def project(self, info: Info, root: 'Cohort') -> 'GraphQLProject':
+        loader = info.context[LoaderKeys.PROJECTS_FOR_IDS]
+        project = await loader.load(root.project)
+        return GraphQLProject.from_internal(project)
+
+
+# Create cohort template GraphQL model
+@strawberry.type
+class GraphQLCohortTemplate:
+    """CohortTemplate GraphQL model"""
+
+    id: str
+    name: str
+    description: str
+    criteria: strawberry.scalars.JSON
+
+    @staticmethod
+    def from_internal(internal: CohortTemplateInternal) -> 'GraphQLCohortTemplate':
+        # At this point, the object that comes in doesn't have an ID field.
+        return GraphQLCohortTemplate(
+            id=cohort_template_id_format(internal.id),
+            name=internal.name,
+            description=internal.description,
+            criteria=internal.criteria,
+        )
+
+
 @strawberry.type
 class GraphQLProject:
     """Project GraphQL model"""
@@ -243,6 +333,35 @@ async def analyses(
         )
         return [GraphQLAnalysis.from_internal(a) for a in internal_analysis]
 
+    @strawberry.field()
+    async def cohorts(
+        self,
+        info: Info,
+        root: 'Project',
+        id: GraphQLFilter[int] | None = None,
+        name: GraphQLFilter[str] | None = None,
+        author: GraphQLFilter[str] | None = None,
+        template_id: GraphQLFilter[int] | None = None,
+        timestamp: GraphQLFilter[datetime.datetime] | None = None,
+    ) -> list['GraphQLCohort']:
+        connection = info.context['connection']
+        connection.project = root.id
+
+        c_filter = CohortFilter(
+            id=id.to_internal_filter(cohort_id_transform_to_raw) if id else None,
+            name=name.to_internal_filter() if name else None,
+            author=author.to_internal_filter() if author else None,
+            template_id=(
+                template_id.to_internal_filter(cohort_template_id_transform_to_raw)
+                if template_id
+                else None
+            ),
+            timestamp=timestamp.to_internal_filter() if timestamp else None,
+        )
+
+        cohorts = await CohortLayer(connection).query(c_filter)
+        return [GraphQLCohort.from_internal(c) for c in cohorts]
+
 
 @strawberry.type
 class GraphQLAuditLog:
@@ -472,11 +591,16 @@ async def participant(
 
     @strawberry.field
     async def assays(
-        self, info: Info, root: 'GraphQLSample', type: GraphQLFilter[str] | None = None
+        self,
+        info: Info,
+        root: 'GraphQLSample',
+        type: GraphQLFilter[str] | None = None,
+        meta: GraphQLMetaFilter | None = None,
     ) -> list['GraphQLAssay']:
         loader_assays_for_sample_ids = info.context[LoaderKeys.ASSAYS_FOR_SAMPLES]
         filter_ = AssayFilter(
             type=type.to_internal_filter() if type else None,
+            meta=meta,
         )
         assays = await loader_assays_for_sample_ids.load(
             {'id': root.internal_id, 'filter': filter_}
@@ -607,7 +731,8 @@ async def assays(
         self, info: Info, root: 'GraphQLSequencingGroup'
     ) -> list['GraphQLAssay']:
         loader = info.context[LoaderKeys.ASSAYS_FOR_SEQUENCING_GROUPS]
-        return await loader.load(root.internal_id)
+        assays = await loader.load(root.internal_id)
+        return [GraphQLAssay.from_internal(assay) for assay in assays]
 
 
 @strawberry.type
@@ -696,13 +821,92 @@ async def project(
 
 
 @strawberry.type
-class Query:
+class Query:  # entry point to graphql.
     """GraphQL Queries"""
 
     @strawberry.field()
     def enum(self, info: Info) -> GraphQLEnum:
         return GraphQLEnum()
 
+    @strawberry.field()
+    async def cohort_templates(
+        self,
+        info: Info,
+        id: GraphQLFilter[str] | None = None,
+        project: GraphQLFilter[str] | None = None,
+    ) -> list[GraphQLCohortTemplate]:
+        connection = info.context['connection']
+        cohort_layer = CohortLayer(connection)
+
+        ptable = ProjectPermissionsTable(connection)
+        project_name_map: dict[str, int] = {}
+        project_filter = None
+        if project:
+            project_names = project.all_values()
+            projects = await ptable.get_and_check_access_to_projects_for_names(
+                user=connection.author, project_names=project_names, readonly=True
+            )
+            project_name_map = {p.name: p.id for p in projects}
+            project_filter = project.to_internal_filter(
+                lambda pname: project_name_map[pname]
+            )
+
+        filter_ = CohortTemplateFilter(
+            id=(
+                id.to_internal_filter(cohort_template_id_transform_to_raw)
+                if id
+                else None
+            ),
+            project=project_filter,
+        )
+
+        cohort_templates = await cohort_layer.query_cohort_templates(filter_)
+        return [
+            GraphQLCohortTemplate.from_internal(cohort_template)
+            for cohort_template in cohort_templates
+        ]
+
+    @strawberry.field()
+    async def cohorts(
+        self,
+        info: Info,
+        id: GraphQLFilter[str] | None = None,
+        project: GraphQLFilter[str] | None = None,
+        name: GraphQLFilter[str] | None = None,
+        author: GraphQLFilter[str] | None = None,
+        template_id: GraphQLFilter[int] | None = None,
+    ) -> list[GraphQLCohort]:
+        connection = info.context['connection']
+        cohort_layer = CohortLayer(connection)
+
+        ptable = ProjectPermissionsTable(connection)
+        project_name_map: dict[str, int] = {}
+        project_filter = None
+        if project:
+            project_names = project.all_values()
+            projects = await ptable.get_and_check_access_to_projects_for_names(
+                user=connection.author, project_names=project_names, readonly=True
+            )
+            project_name_map = {p.name: p.id for p in projects}
+            project_filter = project.to_internal_filter(
+                lambda pname: project_name_map[pname]
+            )
+
+        filter_ = CohortFilter(
+            id=id.to_internal_filter(cohort_id_transform_to_raw) if id else None,
+            name=name.to_internal_filter() if name else None,
+            project=project_filter,
+            author=author.to_internal_filter() if author else None,
+            template_id=(
+                template_id.to_internal_filter(cohort_template_id_transform_to_raw)
+                if template_id
+                else None
+            ),
+        )
+
+        cohorts = await cohort_layer.query(filter_)
+        return [GraphQLCohort.from_internal(cohort) for cohort in cohorts]
+
     @strawberry.field()
     async def project(self, info: Info, name: str) -> GraphQLProject:
         connection = info.context['connection']
@@ -761,6 +965,7 @@ async def sample(
         samples = await slayer.query(filter_)
         return [GraphQLSample.from_internal(sample) for sample in samples]
 
+    # pylint: disable=too-many-arguments
     @strawberry.field
     async def sequencing_groups(
         self,
@@ -772,6 +977,10 @@ async def sequencing_groups(
         technology: GraphQLFilter[str] | None = None,
         platform: GraphQLFilter[str] | None = None,
         active_only: GraphQLFilter[bool] | None = None,
+        created_on: GraphQLFilter[datetime.date] | None = None,
+        assay_meta: GraphQLMetaFilter | None = None,
+        has_cram: bool | None = None,
+        has_gvcf: bool | None = None,
     ) -> list[GraphQLSequencingGroup]:
         connection = info.context['connection']
         sglayer = SequencingGroupLayer(connection)
@@ -812,6 +1021,10 @@ async def sequencing_groups(
                 if active_only
                 else GenericFilter(eq=True)
             ),
+            created_on=created_on.to_internal_filter() if created_on else None,
+            assay_meta=assay_meta,
+            has_cram=has_cram,
+            has_gvcf=has_gvcf,
         )
         sgs = await sglayer.query(filter_)
         return [GraphQLSequencingGroup.from_internal(sg) for sg in sgs]

api/routes/__init__.py

@@ -2,6 +2,7 @@
 from api.routes.analysis_runner import router as analysis_runner_router
 from api.routes.assay import router as assay_router
 from api.routes.billing import router as billing_router
+from api.routes.cohort import router as cohort_router
 from api.routes.enum import router as enum_router
 from api.routes.family import router as family_router
 from api.routes.imports import router as import_router