Merge pull request #506 from VariantEffect/jstone-dev/clinvar-control-query-optimization

ClinVar control query optimization and control response size reduction

Author: jstone-dev

src/mavedb/routers/score_sets.py

@@ -12,7 +12,7 @@
 from ga4gh.va_spec.base.core import Statement, ExperimentalVariantFunctionalImpactStudyResult
 from sqlalchemy import null, or_, select
 from sqlalchemy.exc import MultipleResultsFound, NoResultFound
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import contains_eager, Session
 
 from mavedb import deps
 from mavedb.lib.annotation.exceptions import MappingDataDoesntExistException
@@ -1463,30 +1463,27 @@ async def get_clinical_controls_for_score_set(
 
     clinical_controls_query = (
         select(ClinicalControl)
-        .join(MappedVariant, ClinicalControl.mapped_variants)
-        .join(Variant)
-        .where(Variant.score_set_id == item.id)
+        .join(ClinicalControl.mapped_variants)
+        .join(MappedVariant.variant)
+        .options(
+            contains_eager(ClinicalControl.mapped_variants)
+            .contains_eager(MappedVariant.variant)
+        )
+        .filter(MappedVariant.current.is_(True))
+        .filter(Variant.score_set_id == item.id)
     )
 
     if db_name is not None:
         save_to_logging_context({"db_name": db_name})
-        clinical_controls_query = clinical_controls_query.where(ClinicalControl.db_name == db_name)
+        clinical_controls_query = clinical_controls_query.filter(ClinicalControl.db_name == db_name)
 
     if db_version is not None:
         save_to_logging_context({"db_version": db_version})
-        clinical_controls_query = clinical_controls_query.where(ClinicalControl.db_version == db_version)
-
-    clinical_controls_for_item: Sequence[ClinicalControl] = _db.scalars(clinical_controls_query).all()
-    clinical_controls_with_mapped_variant = []
-    for control_variant in clinical_controls_for_item:
-        control_variant.mapped_variants = [
-            mv for mv in control_variant.mapped_variants if mv.current and mv.variant.score_set_id == item.id
-        ]
+        clinical_controls_query = clinical_controls_query.filter(ClinicalControl.db_version == db_version)
 
-        if control_variant.mapped_variants:
-            clinical_controls_with_mapped_variant.append(control_variant)
+    clinical_controls: Sequence[ClinicalControl] = _db.scalars(clinical_controls_query).unique().all()
 
-    if not clinical_controls_with_mapped_variant:
+    if not clinical_controls:
         logger.info(
             msg="No clinical control variants matching the provided filters are associated with the requested score set.",
             extra=logging_context(),
@@ -1496,9 +1493,9 @@ async def get_clinical_controls_for_score_set(
             detail=f"No clinical control variants matching the provided filters associated with score set URN {urn} were found",
         )
 
-    save_to_logging_context({"resource_count": len(clinical_controls_for_item)})
+    save_to_logging_context({"resource_count": len(clinical_controls)})
 
-    return clinical_controls_for_item
+    return clinical_controls
 
 
 @router.get(

src/mavedb/view_models/clinical_control.py

@@ -39,7 +39,7 @@ class Config:
 
 
 class SavedClinicalControlWithMappedVariants(SavedClinicalControl):
-    mapped_variants: Sequence["SavedMappedVariant"]
+    mapped_variants: Sequence["MappedVariantForClinicalControl"]
 
 
 # Properties to return to non-admin clients
@@ -48,7 +48,7 @@ class ClinicalControl(SavedClinicalControl):
 
 
 class ClinicalControlWithMappedVariants(SavedClinicalControlWithMappedVariants):
-    mapped_variants: Sequence["MappedVariant"]
+    pass
 
 
 class ClinicalControlOptions(BaseModel):
@@ -57,7 +57,7 @@ class ClinicalControlOptions(BaseModel):
 
 
 # ruff: noqa: E402
-from mavedb.view_models.mapped_variant import MappedVariant, SavedMappedVariant, MappedVariantCreate
+from mavedb.view_models.mapped_variant import MappedVariantCreate, MappedVariantForClinicalControl
 
 # ClinicalControlUpdate.model_rebuild()
 SavedClinicalControlWithMappedVariants.model_rebuild()

src/mavedb/view_models/mapped_variant.py

@@ -80,6 +80,22 @@ class MappedVariantWithControls(SavedMappedVariantWithControls):
     gnomad_variants: Sequence["GnomADVariant"]
 
 
+class MappedVariantForClinicalControl(BaseModel):
+    variant_urn: str
+
+    class Config:
+        from_attributes = True
+
+    @model_validator(mode="before")
+    def generate_score_set_urn_list(cls, data: Any):
+        if not hasattr(data, "variant_urn") and hasattr(data, "variant"):
+            try:
+                data.__setattr__("variant_urn", None if not data.variant else data.variant.urn)
+            except AttributeError as exc:
+                raise ValidationError(f"Unable to create {cls.__name__} without attribute: {exc}.")  # type: ignore
+        return data
+
+
 # ruff: noqa: E402
 from mavedb.view_models.clinical_control import ClinicalControlBase, ClinicalControl, SavedClinicalControl
 from mavedb.view_models.gnomad_variant import GnomADVariantBase, GnomADVariant, SavedGnomADVariant