feat: add support for adding QC evaluations via MetadataDbClient

pyproject.toml

@@ -19,7 +19,7 @@ dynamic = ["version"]
 dependencies = [
     "requests",
     "boto3",
-    "pydantic>=2.0",
+    "pydantic>=2.0,<2.12",
     "pydantic-settings>=2.0",
 ]
 

src/aind_data_access_api/document_db.py

@@ -625,6 +625,11 @@ def _deregister_asset_url(self) -> str:
         """Url to deregister (delete) an asset in DocDB and Code Ocean"""
         return f"https://{self.host}/{self.version}/assets/deregister"
 
+    @property
+    def _add_qc_evaluation_url(self) -> str:
+        """Url to add QC evaluation(s) or other QC content to a data asset"""
+        return f"https://{self.host}/{self.version}/add_qc_evaluation"
+
     def generate_data_summary(self, record_id: str) -> Dict[str, Any]:
         """Get an LLM-generated summary for a data asset."""
         url = f"{self._data_summary_url}/{record_id}"
@@ -666,6 +671,25 @@ def deregister_asset(self, s3_location: str) -> Dict[str, Any]:
         response.raise_for_status()
         return response.json()
 
+    def add_qc_evaluation(
+        self, data_asset_id: str, qc_contents: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Add one or more QC evaluations (or other QC content)
+        to a data asset."""
+
+        qc_contents_with_id = {**qc_contents, "data_asset_id": data_asset_id}
+        data = json.dumps(qc_contents_with_id)
+        signed_header = self._signed_request(
+            method="POST", url=self._add_qc_evaluation_url, data=data
+        )
+        response = self.session.post(
+            url=self._add_qc_evaluation_url,
+            headers=dict(signed_header.headers),
+            data=data,
+        )
+        response.raise_for_status()
+        return response.json()
+
 
 class AnalysisDbClient(Client):
     """Class to manage reading and writing to analysis db"""

src/aind_data_access_api/helpers/data_schema.py

@@ -1,10 +1,10 @@
 """Module for convenience functions for the data access API."""
 
 import json
-import pandas as pd
-from typing import List, Optional, Union
 from datetime import datetime, timezone
+from typing import List, Optional, Union
 
+import pandas as pd
 from aind_data_schema.core.quality_control import QualityControl
 
 from aind_data_access_api.document_db import MetadataDbClient

tests/helpers/test_data_schema.py

@@ -3,26 +3,27 @@
 import json
 import os
 import unittest
-import pandas as pd
 from datetime import datetime
 from pathlib import Path
 from unittest.mock import MagicMock, patch
+
+import pandas as pd
 from aind_data_schema.core.quality_control import (
-    QualityControl,
     QCEvaluation,
     QCMetric,
     QCStatus,
-    Status,
+    QualityControl,
     Stage,
+    Status,
 )
 from aind_data_schema_models.modalities import Modality
 
 from aind_data_access_api.helpers.data_schema import (
     get_quality_control_by_id,
     get_quality_control_by_name,
-    get_quality_control_value_df,
-    get_quality_control_status_df,
     get_quality_control_by_names,
+    get_quality_control_status_df,
+    get_quality_control_value_df,
 )
 
 TEST_DIR = Path(os.path.dirname(os.path.realpath(__file__))).parent

tests/test_document_db.py

@@ -1016,6 +1016,77 @@ def test_deregister_asset(
         )
         self.assertEqual(response_message, response)
 
+    @patch("boto3.session.Session")
+    @patch("botocore.auth.SigV4Auth.add_auth")
+    @patch("requests.Session.post")
+    def test_add_qc_evaluation_success(
+        self,
+        mock_post: MagicMock,
+        mock_auth: MagicMock,
+        mock_session: MagicMock,
+    ):
+        """Tests add_qc_evaluation method success case"""
+        mock_creds = MagicMock()
+        mock_creds.access_key = "abc"
+        mock_creds.secret_key = "efg"
+        mock_session.return_value.region_name = "us-west-2"
+        mock_session.get_credentials.return_value = mock_creds
+        mock_response = Response()
+        mock_response.status_code = 200
+        response_message = {"acknowledged": True, "matchedCount": 1}
+        mock_response._content = json.dumps(response_message).encode("utf-8")
+        mock_post.return_value = mock_response
+
+        client = MetadataDbClient(**self.example_client_args)
+        qc_contents = {
+            "qc_evaluation": {
+                "modality": {"name": "ecephys", "abbreviation": "ecephys"},
+                "stage": "Raw data",
+                "name": "Test QC",
+                "metrics": [],
+            }
+        }
+        response = client.add_qc_evaluation("fake-uuid", qc_contents)
+        mock_auth.assert_called_once()
+        mock_post.assert_called_once_with(
+            url="https://example.com/v1/add_qc_evaluation",
+            headers={"Content-Type": "application/json"},
+            data=json.dumps({**qc_contents, "data_asset_id": "fake-uuid"}),
+        )
+        self.assertEqual(response, response_message)
+
+    @patch("boto3.session.Session")
+    @patch("botocore.auth.SigV4Auth.add_auth")
+    @patch("requests.Session.post")
+    def test_add_qc_evaluation_failure(
+        self,
+        mock_post: MagicMock,
+        mock_auth: MagicMock,
+        mock_session: MagicMock,
+    ):
+        """Tests add_qc_evaluation failure case"""
+        mock_creds = MagicMock()
+        mock_creds.access_key = "abc"
+        mock_creds.secret_key = "efg"
+        mock_session.return_value.region_name = "us-west-2"
+        mock_session.return_value.get_credentials.return_value = mock_creds
+        mock_response = Response()
+        mock_response.status_code = 400
+        mock_response._content = json.dumps(
+            {
+                "detail": "Invalid QC evaluation submission. "
+                "'qc_evaluation' field missing required keys.",
+                "error": "BadRequest",
+            }
+        ).encode("utf-8")
+        mock_post.return_value = mock_response
+
+        client = MetadataDbClient(**self.example_client_args)
+        qc_contents = {"foo": "bar"}
+        with self.assertRaises(requests.exceptions.HTTPError) as e:
+            client.add_qc_evaluation("fake-uuid", qc_contents)
+        self.assertIn("400 Client Error", str(e.exception))
+
 
 class TestAnalysisDbClient(unittest.TestCase):
     """Test methods in AnalysisDbClient class."""