From 0a4a92f07c4ceb23cf3993f0525e7868f390bec2 Mon Sep 17 00:00:00 2001 From: Jamie Hendrickson Date: Wed, 8 Oct 2025 11:50:12 -0700 Subject: [PATCH 1/4] feat: add support for adding QC evaluations via MetadataDbClient - Added _add_qc_evaluation_url property for QC evaluation endpoint - Implemented add_qc_evaluation() method to sign and send POST requests - Added unit tests for success and failure cases - Updates import lines to follow isort standards --- src/aind_data_access_api/document_db.py | 28 ++++++++ .../helpers/data_schema.py | 4 +- tests/helpers/test_data_schema.py | 11 +-- tests/test_document_db.py | 72 +++++++++++++++++++ 4 files changed, 108 insertions(+), 7 deletions(-) diff --git a/src/aind_data_access_api/document_db.py b/src/aind_data_access_api/document_db.py index a1c6fe1..45fcd99 100644 --- a/src/aind_data_access_api/document_db.py +++ b/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) 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,29 @@ 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_eval: Dict[str, Any] + ) -> Dict[str, Any]: + """Add one or more QC evaluations to a data asset.""" + + post_request_content = { + "data_asset_id": data_asset_id, + "qc_evaluation": qc_eval, + } + + data = json.dumps(post_request_content) + + 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""" diff --git a/src/aind_data_access_api/helpers/data_schema.py b/src/aind_data_access_api/helpers/data_schema.py index 184a661..a5fc244 100644 --- a/src/aind_data_access_api/helpers/data_schema.py +++ b/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 diff --git a/tests/helpers/test_data_schema.py b/tests/helpers/test_data_schema.py index dafc386..90edf83 100644 --- a/tests/helpers/test_data_schema.py +++ b/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 diff --git a/tests/test_document_db.py b/tests/test_document_db.py index 475f7d3..eee0689 100644 --- a/tests/test_document_db.py +++ b/tests/test_document_db.py @@ -1016,6 +1016,78 @@ 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_eval = { + "modality": {"name": "ecephys", "abbreviation": "ecephys"}, + "stage": "Raw data", + "name": "Test QC", + "metrics": [], + } + response = client.add_qc_evaluation("fake-uuid", qc_eval) + 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( + {"data_asset_id": "fake-uuid", "qc_evaluation": qc_eval} + ), + ) + 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_eval = {"foo": "bar"} + + with self.assertRaises(requests.exceptions.HTTPError) as e: + client.add_qc_evaluation("fake-uuid", qc_eval) + self.assertIn("400 Client Error", str(e.exception)) + class TestAnalysisDbClient(unittest.TestCase): """Test methods in AnalysisDbClient class.""" From c9f785a3ca10fe87fe3666e7f2cd720ab7bdbb19 Mon Sep 17 00:00:00 2001 From: Jamie Hendrickson Date: Wed, 8 Oct 2025 13:26:14 -0700 Subject: [PATCH 2/4] fix: revert import line changes --- src/aind_data_access_api/helpers/data_schema.py | 4 ++-- tests/helpers/test_data_schema.py | 11 +++++------ 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/src/aind_data_access_api/helpers/data_schema.py b/src/aind_data_access_api/helpers/data_schema.py index a5fc244..184a661 100644 --- a/src/aind_data_access_api/helpers/data_schema.py +++ b/src/aind_data_access_api/helpers/data_schema.py @@ -1,10 +1,10 @@ """Module for convenience functions for the data access API.""" import json -from datetime import datetime, timezone +import pandas as pd from typing import List, Optional, Union +from datetime import datetime, timezone -import pandas as pd from aind_data_schema.core.quality_control import QualityControl from aind_data_access_api.document_db import MetadataDbClient diff --git a/tests/helpers/test_data_schema.py b/tests/helpers/test_data_schema.py index 90edf83..dafc386 100644 --- a/tests/helpers/test_data_schema.py +++ b/tests/helpers/test_data_schema.py @@ -3,27 +3,26 @@ 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, - QualityControl, - Stage, Status, + Stage, ) 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_by_names, - get_quality_control_status_df, get_quality_control_value_df, + get_quality_control_status_df, + get_quality_control_by_names, ) TEST_DIR = Path(os.path.dirname(os.path.realpath(__file__))).parent From d19b91f9b72e21fcdde024ae3bc1eea065ec8acc Mon Sep 17 00:00:00 2001 From: Jamie Hendrickson Date: Wed, 8 Oct 2025 14:31:56 -0700 Subject: [PATCH 3/4] feat: add upper bound handling and allow QC evaluation with qc_contents - Allow users to provide `qc_contents` directly when adding a QC evaluation - Include `data_asset_id` in the `qc_contents` dict for post requests - Update pydantic upper bound handling - Restore isort import recommendations --- pyproject.toml | 2 +- src/aind_data_access_api/document_db.py | 16 ++++++------- .../helpers/data_schema.py | 4 ++-- tests/helpers/test_data_schema.py | 11 +++++---- tests/test_document_db.py | 23 +++++++++---------- 5 files changed, 27 insertions(+), 29 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index db5470e..c26d208 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ dynamic = ["version"] dependencies = [ "requests", "boto3", - "pydantic>=2.0", + "pydantic>=2.0,<2.12", "pydantic-settings>=2.0", ] diff --git a/src/aind_data_access_api/document_db.py b/src/aind_data_access_api/document_db.py index 45fcd99..067b760 100644 --- a/src/aind_data_access_api/document_db.py +++ b/src/aind_data_access_api/document_db.py @@ -627,7 +627,7 @@ def _deregister_asset_url(self) -> str: @property def _add_qc_evaluation_url(self) -> str: - """Url to add QC evaluation(s) to a data asset""" + """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]: @@ -672,17 +672,15 @@ def deregister_asset(self, s3_location: str) -> Dict[str, Any]: return response.json() def add_qc_evaluation( - self, data_asset_id: str, qc_eval: Dict[str, Any] + self, data_asset_id: str, qc_contents: Dict[str, Any] ) -> Dict[str, Any]: - """Add one or more QC evaluations to a data asset.""" + """Add one or more QC evaluations (or other QC content) + to a data asset.""" - post_request_content = { - "data_asset_id": data_asset_id, - "qc_evaluation": qc_eval, - } - - data = json.dumps(post_request_content) + qc_contents_with_id = dict(qc_contents) + qc_contents_with_id["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 ) diff --git a/src/aind_data_access_api/helpers/data_schema.py b/src/aind_data_access_api/helpers/data_schema.py index 184a661..a5fc244 100644 --- a/src/aind_data_access_api/helpers/data_schema.py +++ b/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 diff --git a/tests/helpers/test_data_schema.py b/tests/helpers/test_data_schema.py index dafc386..90edf83 100644 --- a/tests/helpers/test_data_schema.py +++ b/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 diff --git a/tests/test_document_db.py b/tests/test_document_db.py index eee0689..a98cc95 100644 --- a/tests/test_document_db.py +++ b/tests/test_document_db.py @@ -1038,20 +1038,20 @@ def test_add_qc_evaluation_success( mock_post.return_value = mock_response client = MetadataDbClient(**self.example_client_args) - qc_eval = { - "modality": {"name": "ecephys", "abbreviation": "ecephys"}, - "stage": "Raw data", - "name": "Test QC", - "metrics": [], + qc_contents = { + "qc_evaluation": { + "modality": {"name": "ecephys", "abbreviation": "ecephys"}, + "stage": "Raw data", + "name": "Test QC", + "metrics": [], + } } - response = client.add_qc_evaluation("fake-uuid", qc_eval) + 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( - {"data_asset_id": "fake-uuid", "qc_evaluation": qc_eval} - ), + data=json.dumps({**qc_contents, "data_asset_id": "fake-uuid"}), ) self.assertEqual(response, response_message) @@ -1082,10 +1082,9 @@ def test_add_qc_evaluation_failure( mock_post.return_value = mock_response client = MetadataDbClient(**self.example_client_args) - qc_eval = {"foo": "bar"} - + qc_contents = {"foo": "bar"} with self.assertRaises(requests.exceptions.HTTPError) as e: - client.add_qc_evaluation("fake-uuid", qc_eval) + client.add_qc_evaluation("fake-uuid", qc_contents) self.assertIn("400 Client Error", str(e.exception)) From 23d6ae330e50b39a45b84ff21981d8565af6b17f Mon Sep 17 00:00:00 2001 From: Jamie Hendrickson Date: Fri, 10 Oct 2025 09:06:06 -0700 Subject: [PATCH 4/4] refactor: build qc_contents_with_id using dict unpacking --- src/aind_data_access_api/document_db.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/aind_data_access_api/document_db.py b/src/aind_data_access_api/document_db.py index 067b760..7bd8fb0 100644 --- a/src/aind_data_access_api/document_db.py +++ b/src/aind_data_access_api/document_db.py @@ -677,9 +677,7 @@ def add_qc_evaluation( """Add one or more QC evaluations (or other QC content) to a data asset.""" - qc_contents_with_id = dict(qc_contents) - qc_contents_with_id["data_asset_id"] = data_asset_id - + 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