add cohort

Author: korikuzma

src/anyvlm/storage/base_storage.py

@@ -9,6 +9,10 @@ class StorageError(Exception):
     """Base AnyLM storage error."""
 
 
+class IncompleteVAObjectError(StorageError):
+    """Raise if provided VA object is missing fully-materialized properties required for storage"""
+
+
 class Storage(ABC):
     """Abstract base class for interacting with storage backends."""
 

src/anyvlm/storage/mapper_registry.py

@@ -34,16 +34,14 @@ def from_db_entity(self, db_entity):  # noqa: ANN201, ANN001
         mapper = self.get_mapper(type(db_entity))
         return mapper.from_db_entity(db_entity)
 
-    def to_db_entity(self, anyvlm_entity) -> orm.Base:  # noqa: ANN001
+    def to_db_entity(self, va_model) -> orm.Base:  # noqa: ANN001
         """Convert any VA-Spec model to its corresponding DB entity."""
-        db_type = self.va_model_to_db_mapping.get(type(anyvlm_entity))
+        db_type = self.va_model_to_db_mapping.get(type(va_model))
         if db_type is None:
-            raise ValueError(
-                f"No DB entity type mapped for VRS model: {type(anyvlm_entity)}"
-            )
+            raise ValueError(f"No DB entity type mapped for VA model: {type(va_model)}")
 
         mapper = self.get_mapper(db_type)
-        return mapper.to_db_entity(anyvlm_entity)
+        return mapper.to_db_entity(va_model)
 
 
 # Global registry instance

src/anyvlm/storage/mappers.py

@@ -24,7 +24,9 @@ def to_db_entity(self, va_model: V) -> D:
         """Convert VA-Spec model to DB entity."""
 
 
-class AlleleFrequencyMapper(BaseMapper):
+class AlleleFrequencyMapper(
+    BaseMapper[CohortAlleleFrequencyStudyResult, orm.AlleleFrequencyData]
+):
     """Maps between Allele Frequency Entities"""
 
     def from_db_entity(
@@ -53,7 +55,7 @@ def from_db_entity(
                 "heterozygotes": heterozygotes,
                 "hemizygotes": hemizygotes,
             },
-            cohort=StudyGroup(name="rare disease"),
+            cohort=StudyGroup(name=db_entity.cohort),
         )
 
     def to_db_entity(
@@ -79,4 +81,5 @@ def to_db_entity(
             ac_hom=ancillary_results["homozygotes"],
             ac_hemi=ancillary_results["hemizygotes"],
             filter=va_model.qualityMeasures["qcFilters"],
+            cohort=va_model.cohort.name,
         )

src/anyvlm/storage/orm.py

@@ -41,6 +41,7 @@ def __tablename__(self) -> str:
         return "allele_frequency_data"
 
     vrs_id: Mapped[str] = mapped_column(String, primary_key=True)
+    cohort: Mapped[str] = mapped_column(String, index=True)
     an: Mapped[int] = mapped_column(Integer)
     ac_het: Mapped[int] = mapped_column(Integer)
     ac_hom: Mapped[int] = mapped_column(Integer)

tests/unit/storage/test_postgres_unit.py

@@ -40,6 +40,7 @@ def insert_and_verify_one_row():
         db_record = rows[0]
 
         assert db_record.vrs_id == "ga4gh:VA.J3Hi64dkKFKdnKIwB2419Qz3STDB2sJq"
+        assert db_record.cohort == "rare disease"
         assert db_record.an == 6164
         assert db_record.ac_het == 1
         assert db_record.ac_hom == 0