feat: refactor models to use NamedTuple for ArcBuildData and RelatedDataBatch, enhancing performance and memory efficiency

Author: Zalfsten

middleware/sql_to_arc/src/middleware/sql_to_arc/models.py

@@ -2,7 +2,7 @@
 
 import concurrent.futures
 from datetime import datetime
-from typing import Any
+from typing import Any, NamedTuple
 
 from pydantic import BaseModel, ConfigDict
 
@@ -83,7 +83,7 @@ class ContactRow(BaseModel):
     model_config = ConfigDict(extra="allow", coerce_numbers_to_str=True, from_attributes=True)
 
 
-class ArcBuildData(BaseModel):
+class ArcBuildData(NamedTuple):
     """Data bundle for building a single ARC."""
 
     investigation_row: InvestigationRow
@@ -93,8 +93,6 @@ class ArcBuildData(BaseModel):
     publications: list[PublicationRow]
     annotations: list[dict[str, Any]]
 
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
 
 class WorkerContext(BaseModel):
     """Context data for a worker process."""
@@ -114,7 +112,7 @@ class WorkerContext(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
-class RelatedDataBatch(BaseModel):
+class RelatedDataBatch(NamedTuple):
     """Batch of related data grouped by investigation ID."""
 
     studies_by_inv: dict[str, list[StudyRow]]
@@ -124,5 +122,3 @@ class RelatedDataBatch(BaseModel):
     anns_by_inv: dict[str, list[dict[str, Any]]]
     study_count: int
     assay_count: int
-
-    model_config = ConfigDict(arbitrary_types_allowed=True)

middleware/sql_to_arc/src/middleware/sql_to_arc/processor.py

@@ -7,10 +7,11 @@
 import multiprocessing
 from collections import defaultdict
 from collections.abc import AsyncGenerator
+from dataclasses import dataclass
 from typing import Any, TypeVar
 
 from opentelemetry import trace
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel
 
 from middleware.api_client import ApiClient, ApiClientError
 from middleware.sql_to_arc.builder import build_single_arc_task
@@ -171,7 +172,8 @@ async def group_stream(
     )
 
 
-class WorkerResources(BaseModel):
+@dataclass(slots=True)
+class WorkerResources:
     """Orchestration resources shared across investigation tasks."""
 
     client: ApiClient
@@ -180,8 +182,6 @@ class WorkerResources(BaseModel):
     executor: concurrent.futures.Executor
     semaphore: asyncio.Semaphore
 
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
 
 def _spawn_investigation_task(
     investigation: InvestigationRow,

middleware/sql_to_arc/tests/unit/test_builder.py

@@ -6,7 +6,14 @@
 import pytest
 
 from middleware.sql_to_arc.builder import build_single_arc_task
-from middleware.sql_to_arc.models import ArcBuildData
+from middleware.sql_to_arc.models import (
+    ArcBuildData,
+    AssayRow,
+    ContactRow,
+    InvestigationRow,
+    PublicationRow,
+    StudyRow,
+)
 
 
 @pytest.fixture
@@ -97,7 +104,12 @@ def sample_publications() -> list[dict[str, Any]]:
 def test_build_simple_arc(sample_investigation: dict[str, Any]) -> None:
     """Test building a basic ARC structure from investigation data."""
     arc_data = ArcBuildData(
-        investigation_row=sample_investigation, studies=[], assays=[], contacts=[], publications=[], annotations=[]
+        investigation_row=InvestigationRow.model_validate(sample_investigation),
+        studies=[],
+        assays=[],
+        contacts=[],
+        publications=[],
+        annotations=[],
     )
     arc_json = build_single_arc_task(arc_data)
     assert isinstance(arc_json, str)
@@ -115,9 +127,9 @@ def test_build_arc_with_study_and_assay(
 ) -> None:
     """Test building an ARC with nested study and assay structures."""
     arc_data = ArcBuildData(
-        investigation_row=sample_investigation,
-        studies=sample_studies,
-        assays=sample_assays,
+        investigation_row=InvestigationRow.model_validate(sample_investigation),
+        studies=[StudyRow.model_validate(s) for s in sample_studies],
+        assays=[AssayRow.model_validate(a) for a in sample_assays],
         contacts=[],
         publications=[],
         annotations=[],
@@ -142,11 +154,11 @@ def test_build_arc_with_contacts_and_pubs(
 ) -> None:
     """Test building an ARC with contacts and publications at both investigation and study levels."""
     arc_data = ArcBuildData(
-        investigation_row=sample_investigation,
-        studies=sample_studies,
+        investigation_row=InvestigationRow.model_validate(sample_investigation),
+        studies=[StudyRow.model_validate(s) for s in sample_studies],
         assays=[],
-        contacts=sample_contacts,
-        publications=sample_publications,
+        contacts=[ContactRow.model_validate(c) for c in sample_contacts],
+        publications=[PublicationRow.model_validate(p) for p in sample_publications],
         annotations=[],
     )
     arc_json = build_single_arc_task(arc_data)
@@ -167,8 +179,8 @@ def test_build_ignores_irrelevant_data(sample_investigation: dict[str, Any]) ->
     other_study = {"identifier": "styX", "investigation_ref": "inv2"}
 
     arc_data = ArcBuildData(
-        investigation_row=sample_investigation,
-        studies=[other_study],
+        investigation_row=InvestigationRow.model_validate(sample_investigation),
+        studies=[StudyRow.model_validate(other_study)],
         assays=[],
         contacts=[],
         publications=[],