pr suggestions

Author: korikuzma

src/anyvlm/storage/__init__.py

@@ -2,6 +2,4 @@
 
 from .base_storage import Storage
 
-DEFAULT_STORAGE_URI = "postgresql://postgres@localhost:5432/anyvlm"
-
-__all__ = ["DEFAULT_STORAGE_URI", "Storage"]
+__all__ = ["Storage"]

src/anyvlm/storage/base_storage.py

@@ -9,10 +9,6 @@ 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."""
 
@@ -29,7 +25,7 @@ def wipe_db(self) -> None:
         """Wipe all data from the storage backend."""
 
     @abstractmethod
-    def add_allele_frequency(self, caf: CohortAlleleFrequencyStudyResult) -> None:
+    def add_allele_frequencies(self, caf: CohortAlleleFrequencyStudyResult) -> None:
         """Add allele frequency data to the database. Will skip conflicts.
 
         :param caf: Cohort allele frequency study result object to insert into the DB

src/anyvlm/storage/mappers.py

@@ -55,8 +55,8 @@ def from_db_entity(
                 "heterozygotes": heterozygotes,
                 "hemizygotes": hemizygotes,
             },
-            cohort=StudyGroup(name=db_entity.cohort),
-        )
+            cohort=StudyGroup(name=db_entity.cohort),  # type: ignore
+        )  # type: ignore
 
     def to_db_entity(
         self, va_model: CohortAlleleFrequencyStudyResult
@@ -66,7 +66,8 @@ def to_db_entity(
         :param va_model: VA-Spec Cohort Allele Frequency Study Result instance
         :return: ORM Allele Frequency Data instance
         """
-        ancillary_results = va_model.ancillaryResults
+        ancillary_results = va_model.ancillaryResults or {}
+        quality_filters = va_model.qualityMeasures or {}
         focus_allele = va_model.focusAllele
 
         if isinstance(focus_allele, iriReference):
@@ -77,9 +78,10 @@ def to_db_entity(
         return orm.AlleleFrequencyData(
             vrs_id=vrs_id,
             an=va_model.locusAlleleCount,
-            ac_het=ancillary_results["heterozygotes"],
-            ac_hom=ancillary_results["homozygotes"],
-            ac_hemi=ancillary_results["hemizygotes"],
-            filter=va_model.qualityMeasures["qcFilters"],
+            ac=va_model.focusAlleleCount,
+            ac_het=ancillary_results.get("heterozygotes"),
+            ac_hom=ancillary_results.get("homozygotes"),
+            ac_hemi=ancillary_results.get("hemizygotes"),
+            filter=quality_filters.get("qcFilters"),
             cohort=va_model.cohort.name,
         )

src/anyvlm/storage/orm.py

@@ -41,12 +41,13 @@ 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)
-    ac_hemi: Mapped[int] = mapped_column(Integer)
-    filter: Mapped[list[str]] = mapped_column(ARRAY(String))
+    cohort: Mapped[str] = mapped_column(String, index=True, nullable=False)
+    an: Mapped[int] = mapped_column(Integer, nullable=False)
+    ac: Mapped[int] = mapped_column(Integer, nullable=False)
+    ac_het: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    ac_hom: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    ac_hemi: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    filter: Mapped[list[str] | None] = mapped_column(ARRAY(String), nullable=True)
 
 
 def create_tables(db_url: str) -> None:

src/anyvlm/storage/postgres.py

@@ -36,7 +36,7 @@ def wipe_db(self) -> None:
         with self.session_factory() as session, session.begin():
             session.execute(delete(orm.AlleleFrequencyData))
 
-    def add_allele_frequency(self, caf: CohortAlleleFrequencyStudyResult) -> None:
+    def add_allele_frequencies(self, caf: CohortAlleleFrequencyStudyResult) -> None:
         """Add allele frequency data to the database. Will skip conflicts.
 
         :param caf: Cohort allele frequency study result object to insert into the DB

tests/conftest.py

@@ -9,8 +9,6 @@
 from ga4gh.vrs import models
 from pydantic import BaseModel
 
-from anyvlm.storage import orm
-from anyvlm.storage.base_storage import Storage
 from anyvlm.storage.postgres import PostgresObjectStore
 
 
@@ -79,13 +77,13 @@ def postgres_uri():
 def postgres_storage(postgres_uri: str):
     """Reset storage state after each test case"""
     storage = PostgresObjectStore(postgres_uri)
-    yield storage
     storage.wipe_db()
+    return storage
 
 
 @pytest.fixture
 def caf_iri():
-    """Create test fixture for CAF object that uses iriRference for focusAllele
+    """Create test fixture for CAF object that uses iriReference for focusAllele
 
     This is a GREGoR example from issue #23 description
     """
@@ -100,11 +98,5 @@ def caf_iri():
             "heterozygotes": 1,
             "hemizygotes": 0,
         },
-        cohort=StudyGroup(name="rare disease"),
-    )
-
-
-def return_cafs(storage: Storage):
-    """Return allele frequency data in db"""
-    with storage.session_factory() as session:
-        return session.query(orm.AlleleFrequencyData).all()
+        cohort=StudyGroup(name="rare disease"),  # type: ignore
+    )  # type: ignore

tests/integration/storage/test_postgres_integration.py

@@ -1,19 +1,26 @@
 """Test postgres storage integration methods"""
 
 from ga4gh.va_spec.base import CohortAlleleFrequencyStudyResult
-from tests.conftest import return_cafs
+from sqlalchemy import select
 
+from anyvlm.storage import orm
 from anyvlm.storage.postgres import PostgresObjectStore
 
 
+def return_cafs(storage: PostgresObjectStore):
+    """Integration-test helper: Get all CAF rows from the DB"""
+    with storage.session_factory() as session:
+        return session.execute(select(orm.AlleleFrequencyData)).scalars().all()
+
+
 def test_db_lifecycle(postgres_uri: str, caf_iri: CohortAlleleFrequencyStudyResult):
     """Test that DB lifecycle works correctly"""
     # set up and populate DB
     storage = PostgresObjectStore(postgres_uri)
     caf_rows = return_cafs(storage)
     assert caf_rows == []
 
-    storage.add_allele_frequency(caf_iri)
+    storage.add_allele_frequencies(caf_iri)
     caf_rows = return_cafs(storage)
     assert len(caf_rows) == 1
 

tests/unit/storage/test_postgres_unit.py

@@ -2,9 +2,9 @@
 
 import pytest
 from ga4gh.core.models import iriReference
-from ga4gh.va_spec.base import CohortAlleleFrequencyStudyResult
-from ga4gh.vrs.models import Allele, LiteralSequenceExpression
-from tests.conftest import return_cafs
+from ga4gh.va_spec.base import CohortAlleleFrequencyStudyResult, StudyGroup
+from ga4gh.vrs.models import Allele, LiteralSequenceExpression, sequenceString
+from sqlalchemy.exc import IntegrityError
 
 from anyvlm.storage.postgres import PostgresObjectStore
 
@@ -17,37 +17,48 @@ def caf_allele(caf_iri: CohortAlleleFrequencyStudyResult):
     """
     caf_allele = caf_iri.model_copy(deep=True)
     caf_allele.focusAllele = Allele(
-        id=caf_iri.focusAllele.root,
+        id=caf_iri.focusAllele.root,  # type: ignore
         location=iriReference("locations.json#/1"),
-        state=LiteralSequenceExpression(sequence="A"),
+        state=LiteralSequenceExpression(sequence=sequenceString("A")),
     )
     return caf_allele
 
 
+@pytest.fixture
+def caf_empty_cohort(caf_iri: CohortAlleleFrequencyStudyResult):
+    """Create test fixture for CAF object that uses empty cohort"""
+    caf = caf_iri.model_copy(deep=True)
+    caf.cohort = StudyGroup()  # type: ignore
+    return caf
+
+
 @pytest.mark.parametrize("caf_fixture_name", ["caf_iri", "caf_allele"])
-def test_add_allele_frequency(
+def test_add_allele_frequencies(
     request, caf_fixture_name: str, postgres_storage: PostgresObjectStore
 ):
-    """Test that add_allele_frequency method works correctly"""
+    """Test that add_allele_frequencies method works correctly"""
     caf = request.getfixturevalue(caf_fixture_name)
+    try:
+        postgres_storage.add_allele_frequencies(caf)
+    except Exception as e:  # noqa: BLE001
+        pytest.fail(f"add_allele_frequencies raised an exception: {e}")
 
-    def insert_and_verify_one_row():
-        postgres_storage.add_allele_frequency(caf)
+    caf = CohortAlleleFrequencyStudyResult(
+        focusAllele=iriReference("ga4gh:VA.J3Hi64dkKFKdnKIwB2419Qz3STDB2sJq"),
+        focusAlleleCount=1,
+        locusAlleleCount=6164,
+        focusAlleleFrequency=0.000162232,
+        qualityMeasures={"qcFilters": ["LowQual", "NO_HQ_GENOTYPES"]},
+        cohort=StudyGroup(name="rare disease"),  # type: ignore
+    )  # type: ignore
 
-        rows = return_cafs(postgres_storage)
-        assert len(rows) == 1
+    postgres_storage.add_allele_frequencies(caf)
 
-        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
-        assert db_record.filter == ["LowQual", "NO_HQ_GENOTYPES"]
-
-    # first insert
-    insert_and_verify_one_row()
-
-    # second insert
-    insert_and_verify_one_row()
+def test_add_allele_frequencies_failures(
+    postgres_storage: PostgresObjectStore,
+    caf_empty_cohort: CohortAlleleFrequencyStudyResult,
+):
+    """Test that add_allele_frequencies method fails correctly on bad input"""
+    with pytest.raises(IntegrityError, match='null value in column "cohort"'):
+        postgres_storage.add_allele_frequencies(caf_empty_cohort)