feat: refactor data models and mapping functions to use Pydantic for improved type safety and validation

Author: Zalfsten

middleware/sql_to_arc/src/middleware/sql_to_arc/builder.py

@@ -8,6 +8,8 @@
 
 from arctrl import (  # type: ignore[import-untyped]
     ARC,
+    ArcAssay,
+    ArcStudy,
     ArcTable,
     CompositeCell,
     CompositeHeader,
@@ -22,57 +24,74 @@
     map_publication,
     map_study,
 )
-from middleware.sql_to_arc.models import ArcBuildData
+from middleware.sql_to_arc.models import (
+    ArcBuildData,
+    AssayRow,
+    ContactRow,
+    PublicationRow,
+    StudyRow,
+)
 
 logger = logging.getLogger(__name__)
 
 
-def _add_studies_to_arc(arc: ARC, study_rows: list[dict[str, Any]]) -> dict[str, Any]:
+def _add_studies_to_arc(arc: ARC, study_rows: list[StudyRow]) -> dict[str, ArcStudy]:
     """Add studies to ARC and return study map."""
-    study_map = {}
+    study_map: dict[str, ArcStudy] = {}
     for s_row in study_rows:
         study = map_study(s_row)
         arc.AddRegisteredStudy(study)
-        study_map[str(s_row["identifier"])] = study
+        study_map[str(s_row.identifier)] = study
     return study_map
 
 
-def _add_assays_to_arc(arc: ARC, assay_rows: list[dict[str, Any]], study_map: dict[str, Any]) -> dict[str, Any]:
+def _add_assays_to_arc(arc: ARC, assay_rows: list[AssayRow], study_map: dict[str, ArcStudy]) -> dict[str, ArcAssay]:
     """Add assays to ARC, link to studies, and return assay map."""
-    assay_map = {}
+    assay_map: dict[str, ArcAssay] = {}
     for a_row in assay_rows:
         assay = map_assay(a_row)
         arc.AddAssay(assay)
-        assay_map[str(a_row["identifier"])] = assay
+        assay_map[str(a_row.identifier)] = assay
 
         # Link Assay to Studies
-        study_ref_json = a_row.get("study_ref")
-        if not study_ref_json:
-            continue
+        _link_assay_to_studies(assay, a_row.study_ref, study_map)
+
+    return assay_map
+
 
+def _link_assay_to_studies(assay: ArcAssay, study_ref_val: Any, study_map: dict[str, ArcStudy]) -> None:
+    """Link an assay to one or more studies based on the study_ref value."""
+    if not study_ref_val:
+        return
+
+    if isinstance(study_ref_val, str):
         try:
-            study_refs = json.loads(study_ref_json)
+            study_refs = json.loads(study_ref_val)
             if isinstance(study_refs, list):
                 for s_ref in study_refs:
-                    if s_ref in study_map:
-                        study_map[s_ref].RegisterAssay(assay.Identifier)
+                    if str(s_ref) in study_map:
+                        study_map[str(s_ref)].RegisterAssay(assay.Identifier)
+                return
         except json.JSONDecodeError:
+            # Handle single ID if it's not JSON (fall through)
             pass
 
-    return assay_map
+    # Handle single ID (string or int)
+    if str(study_ref_val) in study_map:
+        study_map[str(study_ref_val)].RegisterAssay(assay.Identifier)
 
 
 def _add_contacts_to_arc(
     arc: ARC,
     inv_id: str,
-    contacts: list[dict[str, Any]],
-    study_map: dict[str, Any],
-    assay_map: dict[str, Any],
+    contacts: list[ContactRow],
+    study_map: dict[str, ArcStudy],
+    assay_map: dict[str, ArcAssay],
 ) -> None:
     """Add contacts to investigation, studies, and assays."""
     # Investigation contacts
     inv_contacts = [
-        c for c in contacts if c.get("investigation_ref") == inv_id and c.get("target_type") == "investigation"
+        c for c in contacts if str(c.investigation_ref) == inv_id and getattr(c, "target_type", None) == "investigation"
     ]
     for c_row in inv_contacts:
         arc.Contacts.append(map_contact(c_row))
@@ -82,7 +101,9 @@ def _add_contacts_to_arc(
         stu_contacts = [
             c
             for c in contacts
-            if c.get("investigation_ref") == inv_id and c.get("target_type") == "study" and c.get("target_ref") == s_id
+            if str(c.investigation_ref) == inv_id
+            and getattr(c, "target_type", None) == "study"
+            and str(getattr(c, "target_ref", None)) == s_id
         ]
         for c_row in stu_contacts:
             study.Contacts.append(map_contact(c_row))
@@ -92,19 +113,26 @@ def _add_contacts_to_arc(
         ass_contacts = [
             c
             for c in contacts
-            if c.get("investigation_ref") == inv_id and c.get("target_type") == "assay" and c.get("target_ref") == a_id
+            if str(c.investigation_ref) == inv_id
+            and getattr(c, "target_type", None) == "assay"
+            and str(getattr(c, "target_ref", None)) == a_id
         ]
         for c_row in ass_contacts:
             assay.Performers.append(map_contact(c_row))
 
 
 def _add_publications_to_arc(
-    arc: ARC, inv_id: str, publications: list[dict[str, Any]], study_map: dict[str, Any]
+    arc: ARC,
+    inv_id: str,
+    publications: list[PublicationRow],
+    study_map: dict[str, ArcStudy],
 ) -> None:
     """Add publications to investigation and studies."""
     # Investigation publications
     inv_pubs = [
-        p for p in publications if p.get("investigation_ref") == inv_id and p.get("target_type") == "investigation"
+        p
+        for p in publications
+        if str(p.investigation_ref) == inv_id and getattr(p, "target_type", None) == "investigation"
     ]
     for p_row in inv_pubs:
         arc.Publications.append(map_publication(p_row))
@@ -114,7 +142,9 @@ def _add_publications_to_arc(
         stu_pubs = [
             p
             for p in publications
-            if p.get("investigation_ref") == inv_id and p.get("target_type") == "study" and p.get("target_ref") == s_id
+            if str(p.investigation_ref) == inv_id
+            and getattr(p, "target_type", None) == "study"
+            and str(getattr(p, "target_ref", None)) == s_id
         ]
         for p_row in stu_pubs:
             study.Publications.append(map_publication(p_row))
@@ -271,16 +301,16 @@ def build_single_arc_task(data: ArcBuildData) -> str:
     This function is designed to run in a separate process.
     It returns the JSON representation to minimize memory footprint in the main process.
     """
-    inv_id = str(data.investigation_row["identifier"])
+    inv_id = str(data.investigation_row.identifier)
 
     try:
         # Map Investigation and create ARC
         arc_inv = map_investigation(data.investigation_row)
         arc = ARC.from_arc_investigation(arc_inv)
 
         # Identify relevant studies and assays
-        relevant_studies = [s for s in data.studies if s.get("investigation_ref") == inv_id]
-        relevant_assays = [a for a in data.assays if a.get("investigation_ref") == inv_id]
+        relevant_studies = [s for s in data.studies if str(s.investigation_ref) == inv_id]
+        relevant_assays = [a for a in data.assays if str(a.investigation_ref) == inv_id]
 
         # Add studies and assays
         study_map = _add_studies_to_arc(arc, relevant_studies)

middleware/sql_to_arc/src/middleware/sql_to_arc/database.py

@@ -15,6 +15,14 @@
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
 from sqlalchemy.sql import select
 
+from middleware.sql_to_arc.models import (
+    AssayRow,
+    ContactRow,
+    InvestigationRow,
+    PublicationRow,
+    StudyRow,
+)
+
 # Define metadata
 metadata = MetaData()
 
@@ -126,55 +134,55 @@ def __init__(self, connection_string: str) -> None:
         """Initialize database with connection string."""
         self.engine: AsyncEngine = create_async_engine(connection_string, echo=False)
 
-    async def stream_investigations(self, limit: int | None = None) -> AsyncGenerator[dict[str, Any], None]:
+    async def stream_investigations(self, limit: int | None = None) -> AsyncGenerator[InvestigationRow, None]:
         """Stream investigations using a server-side cursor."""
         async with self.engine.connect() as conn:
             stmt = select(v_investigation)
             if limit:
                 stmt = stmt.limit(limit)
             result = await conn.stream(stmt.execution_options(stream_results=True))
             async for row in result.mappings():
-                yield dict(row)
+                yield InvestigationRow.model_validate(row)
 
-    async def stream_studies(self, investigation_ids: list[str]) -> AsyncGenerator[dict[str, Any], None]:
+    async def stream_studies(self, investigation_ids: list[str]) -> AsyncGenerator[StudyRow, None]:
         """Stream studies for given investigations."""
         if not investigation_ids:
             return
         async with self.engine.connect() as conn:
             stmt = select(v_study).where(v_study.c.investigation_ref.in_(investigation_ids))
             result = await conn.stream(stmt.execution_options(stream_results=True))
             async for row in result.mappings():
-                yield dict(row)
+                yield StudyRow.model_validate(row)
 
-    async def stream_assays(self, investigation_ids: list[str]) -> AsyncGenerator[dict[str, Any], None]:
+    async def stream_assays(self, investigation_ids: list[str]) -> AsyncGenerator[AssayRow, None]:
         """Stream assays for given investigations."""
         if not investigation_ids:
             return
         async with self.engine.connect() as conn:
             stmt = select(v_assay).where(v_assay.c.investigation_ref.in_(investigation_ids))
             result = await conn.stream(stmt.execution_options(stream_results=True))
             async for row in result.mappings():
-                yield dict(row)
+                yield AssayRow.model_validate(row)
 
-    async def stream_contacts(self, investigation_ids: list[str]) -> AsyncGenerator[dict[str, Any], None]:
+    async def stream_contacts(self, investigation_ids: list[str]) -> AsyncGenerator[ContactRow, None]:
         """Stream contacts for given investigations."""
         if not investigation_ids:
             return
         async with self.engine.connect() as conn:
             stmt = select(v_contact).where(v_contact.c.investigation_ref.in_(investigation_ids))
             result = await conn.stream(stmt.execution_options(stream_results=True))
             async for row in result.mappings():
-                yield dict(row)
+                yield ContactRow.model_validate(row)
 
-    async def stream_publications(self, investigation_ids: list[str]) -> AsyncGenerator[dict[str, Any], None]:
+    async def stream_publications(self, investigation_ids: list[str]) -> AsyncGenerator[PublicationRow, None]:
         """Stream publications for given investigations."""
         if not investigation_ids:
             return
         async with self.engine.connect() as conn:
             stmt = select(v_publication).where(v_publication.c.investigation_ref.in_(investigation_ids))
             result = await conn.stream(stmt.execution_options(stream_results=True))
             async for row in result.mappings():
-                yield dict(row)
+                yield PublicationRow.model_validate(row)
 
     async def stream_annotation_tables(self, investigation_ids: list[str]) -> AsyncGenerator[dict[str, Any], None]:
         """Stream annotation tables for given investigations."""