From 92ca51ecfe4a079a4b7fbab7e373975a0e388107 Mon Sep 17 00:00:00 2001 From: Helen Lin Date: Mon, 23 Jun 2025 17:55:36 -0700 Subject: [PATCH 1/3] feat: default database and collection for MetadataDbClient --- src/aind_data_access_api/document_db.py | 30 +++++++++++++++++++++++++ tests/test_document_db.py | 22 ++++++++++++++++-- 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/src/aind_data_access_api/document_db.py b/src/aind_data_access_api/document_db.py index d6b68ed..1771b16 100644 --- a/src/aind_data_access_api/document_db.py +++ b/src/aind_data_access_api/document_db.py @@ -362,6 +362,36 @@ def __exit__(self, exc_type, exc_val, exc_tb): class MetadataDbClient(Client): """Class to manage reading and writing to metadata db""" + def __init__( + self, + host: str, + database: str = "metadata_index", + collection: str = "data_assets", + version: str = "v1", + boto: Optional[BotoSession] = None, + session: Optional[Session] = None, + ): + """ + Instantiate a MetadataDbClient. + + Parameters + ---------- + host : str + database : str + collection : str + version : str + boto : Optional[BotoSession] + session : Optional[Session] + """ + super().__init__( + host=host, + database=database, + collection=collection, + version=version, + boto=boto, + session=session, + ) + def retrieve_docdb_records( self, filter_query: Optional[dict] = None, diff --git a/tests/test_document_db.py b/tests/test_document_db.py index d4be5a7..4f5dcae 100644 --- a/tests/test_document_db.py +++ b/tests/test_document_db.py @@ -437,8 +437,6 @@ class TestMetadataDbClient(unittest.TestCase): example_client_args = { "host": "example.com/", - "database": "metadata_db", - "collection": "data_assets", } example_record_list = [ @@ -452,6 +450,26 @@ class TestMetadataDbClient(unittest.TestCase): for id_num in range(0, 10) ] + def test_metadatadbclient_constructor(self): + """Tests class constructor""" + client = MetadataDbClient(**self.example_client_args) + + self.assertEqual("example.com", client.host) + self.assertEqual("metadata_index", client.database) + self.assertEqual("data_assets", client.collection) + self.assertEqual("v1", client.version) + self.assertEqual( + "https://example.com/v1/metadata_index/data_assets", + client._base_url, + ) + + client = MetadataDbClient(**self.example_client_args, version="v2") + self.assertEqual("v2", client.version) + self.assertEqual( + "https://example.com/v2/metadata_index/data_assets", + client._base_url, + ) + @patch("aind_data_access_api.document_db.Client._find_records") def test_retrieve_docdb_records( self, From daa9604ad1ef8dae9bccbe785eddd7eb0c48d57e Mon Sep 17 00:00:00 2001 From: Helen Lin Date: Mon, 23 Jun 2025 17:56:22 -0700 Subject: [PATCH 2/3] docs: update MetadataDbClient examples --- docs/source/ExamplesDocDBRestApi.rst | 13 +++++-------- docs/source/UserGuide.rst | 9 +++++---- 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/docs/source/ExamplesDocDBRestApi.rst b/docs/source/ExamplesDocDBRestApi.rst index 9c8ae3f..d6a6778 100644 --- a/docs/source/ExamplesDocDBRestApi.rst +++ b/docs/source/ExamplesDocDBRestApi.rst @@ -18,13 +18,14 @@ Count Example 1: Get # of records with a certain subject_id from aind_data_access_api.document_db import MetadataDbClient API_GATEWAY_HOST = "api.allenneuraldynamics.org" - DATABASE = "metadata_index" - COLLECTION = "data_assets" + # Default database and collection names are set in the client + # DATABASE = "metadata_index" + # COLLECTION = "data_assets" docdb_api_client = MetadataDbClient( host=API_GATEWAY_HOST, - database=DATABASE, - collection=COLLECTION, + # database=DATABASE, + # collection=COLLECTION, ) filter = {"subject.subject_id": "731015"} @@ -136,8 +137,6 @@ It's possible to attach a custom Session to retry certain requests errors: from aind_data_access_api.document_db import MetadataDbClient API_GATEWAY_HOST = "api.allenneuraldynamics.org" - DATABASE = "metadata_index" - COLLECTION = "data_assets" retry = Retry( total=5, @@ -151,8 +150,6 @@ It's possible to attach a custom Session to retry certain requests errors: with MetadataDbClient( host=API_GATEWAY_HOST, - database=DATABASE, - collection=COLLECTION, session=session, ) as docdb_api_client: records = docdb_api_client.retrieve_docdb_records(limit=10) diff --git a/docs/source/UserGuide.rst b/docs/source/UserGuide.rst index d03daf3..a20301c 100644 --- a/docs/source/UserGuide.rst +++ b/docs/source/UserGuide.rst @@ -54,13 +54,14 @@ REST API from aind_data_access_api.document_db import MetadataDbClient API_GATEWAY_HOST = "api.allenneuraldynamics.org" - DATABASE = "metadata_index" - COLLECTION = "data_assets" + # Default database and collection names are set in the client + # DATABASE = "metadata_index" + # COLLECTION = "data_assets" docdb_api_client = MetadataDbClient( host=API_GATEWAY_HOST, - database=DATABASE, - collection=COLLECTION, + # database=DATABASE, + # collection=COLLECTION, ) filter = {"subject.subject_id": "731015"} From b3691a37050c33a76464c6b67629bda73fbeca47 Mon Sep 17 00:00:00 2001 From: Helen Lin Date: Tue, 24 Jun 2025 14:31:52 -0700 Subject: [PATCH 3/3] docs: how to override defaults --- docs/source/ExamplesDocDBRestApi.rst | 6 ++---- docs/source/UserGuide.rst | 6 ++---- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/docs/source/ExamplesDocDBRestApi.rst b/docs/source/ExamplesDocDBRestApi.rst index d6a6778..38fa8d6 100644 --- a/docs/source/ExamplesDocDBRestApi.rst +++ b/docs/source/ExamplesDocDBRestApi.rst @@ -19,13 +19,11 @@ Count Example 1: Get # of records with a certain subject_id API_GATEWAY_HOST = "api.allenneuraldynamics.org" # Default database and collection names are set in the client - # DATABASE = "metadata_index" - # COLLECTION = "data_assets" + # To override the defaults, provide the database and collection + # parameters in the constructor docdb_api_client = MetadataDbClient( host=API_GATEWAY_HOST, - # database=DATABASE, - # collection=COLLECTION, ) filter = {"subject.subject_id": "731015"} diff --git a/docs/source/UserGuide.rst b/docs/source/UserGuide.rst index a20301c..b4a11f7 100644 --- a/docs/source/UserGuide.rst +++ b/docs/source/UserGuide.rst @@ -55,13 +55,11 @@ REST API API_GATEWAY_HOST = "api.allenneuraldynamics.org" # Default database and collection names are set in the client - # DATABASE = "metadata_index" - # COLLECTION = "data_assets" + # To override the defaults, provide the database and collection + # parameters in the constructor docdb_api_client = MetadataDbClient( host=API_GATEWAY_HOST, - # database=DATABASE, - # collection=COLLECTION, ) filter = {"subject.subject_id": "731015"}