Merge pull request #167 from AllenNeuralDynamics/release-v0.20.1

v0.20.1

Author: helen-m-lin

src/aind_data_asset_indexer/__init__.py

@@ -1,3 +1,3 @@
 """Package"""
 
-__version__ = "0.20.0"
+__version__ = "0.20.1"

src/aind_data_asset_indexer/aind_bucket_indexer.py

@@ -14,6 +14,7 @@
 import requests
 from aind_data_access_api.document_db import MetadataDbClient
 from aind_data_access_api.utils import (
+    build_docdb_location_to_id_map,
     get_s3_bucket_and_prefix,
     get_s3_location,
     paginate_docdb,
@@ -26,7 +27,6 @@
 
 from aind_data_asset_indexer.models import AindIndexBucketJobSettings
 from aind_data_asset_indexer.utils import (
-    build_docdb_location_to_id_map,
     compute_md5_hash,
     core_schema_file_names,
     create_metadata_object_key,
@@ -83,8 +83,7 @@ def _create_docdb_client(self) -> MetadataDbClient:
         session.mount("https://", adapter)
         return MetadataDbClient(
             host=self.job_settings.doc_db_host,
-            database=self.job_settings.doc_db_db_name,
-            collection=self.job_settings.doc_db_collection_name,
+            version="v1",
             session=session,
         )
 
@@ -431,7 +430,6 @@ def _process_docdb_record(
                     # Pull record from docdb to get new last_modified as well
                     docdb_response = docdb_client.retrieve_docdb_records(
                         filter_query={"_id": docdb_record["_id"]},
-                        paginate=False,
                     )
                     docdb_record = docdb_response[0]
                 # Sync docdb record to metadata.nd.json in root folder

src/aind_data_asset_indexer/codeocean_bucket_indexer.py

@@ -64,8 +64,7 @@ def _create_docdb_client(self) -> MetadataDbClient:
         session.mount("https://", adapter)
         return MetadataDbClient(
             host=self.job_settings.doc_db_host,
-            database=self.job_settings.doc_db_db_name,
-            collection=self.job_settings.doc_db_collection_name,
+            version="v1",
             session=session,
         )
 

src/aind_data_asset_indexer/models.py

@@ -51,8 +51,8 @@ class AindIndexBucketJobSettings(IndexJobSettings):
     """Aind Index Bucket Job Settings"""
 
     doc_db_host: str
-    doc_db_db_name: str
-    doc_db_collection_name: str
+    doc_db_db_name: Optional[str]
+    doc_db_collection_name: Optional[str]
     run_docdb_sync: bool = Field(
         default=True,
         description="If true, then process DocDB records to sync to S3.",
@@ -84,8 +84,8 @@ class CodeOceanIndexBucketJobSettings(IndexJobSettings):
     """Code Ocean Index Bucket Job Settings"""
 
     doc_db_host: str
-    doc_db_db_name: str
-    doc_db_collection_name: str
+    doc_db_db_name: Optional[str]
+    doc_db_collection_name: Optional[str]
     codeocean_domain: str
     codeocean_token: SecretStr
     run_co_sync: bool = Field(

src/aind_data_asset_indexer/utils.py

@@ -9,7 +9,6 @@
 from typing import Dict, Iterator, List, Optional
 from urllib.parse import urlparse
 
-from aind_data_access_api.document_db import MetadataDbClient
 from aind_data_access_api.utils import get_s3_location
 from aind_data_schema.core.data_description import DataLevel, DataRegex
 from aind_data_schema.core.metadata import CORE_FILES as CORE_SCHEMAS
@@ -783,40 +782,6 @@ def sync_core_json_files(
                 )
 
 
-# TODO: replace with method from aind_data_access_api.utils once available
-def build_docdb_location_to_id_map(
-    docdb_api_client: MetadataDbClient,
-    bucket: str,
-    prefixes: List[str],
-) -> Dict[str, str]:
-    """
-    For a given s3 bucket and list of prefixes, return a dictionary that looks
-    like {'s3://bucket/prefix': 'abc-1234'} where the value is the id of the
-    record in DocDb. If the record does not exist, then there will be no key
-    in the dictionary.
-
-    Parameters
-    ----------
-    docdb_api_client : MetadataDbClient
-    bucket : str
-    prefixes : List[str]
-
-    Returns
-    -------
-    Dict[str, str]
-
-    """
-    locations = [get_s3_location(bucket=bucket, prefix=p) for p in prefixes]
-    # NOTE: use aggregation since filter too large for retrieve_docdb_records
-    agg_pipeline = [
-        {"$match": {"location": {"$in": locations}}},
-        {"$project": {"location": 1, "_id": 1}},
-    ]
-    results = docdb_api_client.aggregate_docdb_records(pipeline=agg_pipeline)
-    location_to_id_map = {r["location"]: r["_id"] for r in results}
-    return location_to_id_map
-
-
 def get_all_processed_codeocean_asset_records(
     co_client: CodeOcean, co_data_asset_bucket: str
 ) -> Dict[str, dict]:

tests/test_aind_bucket_indexer.py

@@ -921,7 +921,6 @@ def test_process_docdb_record_valid_metadata_nd_json_file(
         )
         mock_docdb_client.retrieve_docdb_records.assert_called_once_with(
             filter_query={"_id": self.example_md_record.get("_id")},
-            paginate=False,
         )
         mock_write_root_file_with_record_info.assert_called_once_with(
             s3_client=mock_s3_client,

tests/test_utils.py

@@ -13,7 +13,6 @@
 from codeocean.data_asset import DataAsset
 
 from aind_data_asset_indexer.utils import (
-    build_docdb_location_to_id_map,
     build_metadata_record_from_prefix,
     compute_md5_hash,
     cond_copy_then_sync_core_json_files,
@@ -1226,46 +1225,6 @@ def test_cond_copy_then_sync_core_json_files_mismatch(
             Bucket=bucket, Key=f"{pfx}/rig.json"
         )
 
-    @patch("aind_data_access_api.utils.MetadataDbClient")
-    def test_build_docdb_location_to_id_map(
-        self, mock_docdb_client: MagicMock
-    ):
-        """Tests build_docdb_location_to_id_map"""
-        bucket = "aind-ephys-data-dev-u5u0i5"
-        mock_docdb_client.aggregate_docdb_records.return_value = [
-            {
-                "_id": "70bcf356-985f-4a2a-8105-de900e35e788",
-                "location": (
-                    f"s3://{bucket}/ecephys_655019_2000-04-04_04-00-00"
-                ),
-            },
-            {
-                "_id": "5ca4a951-d374-4f4b-8279-d570a35b2286",
-                "location": (
-                    f"s3://{bucket}/ecephys_567890_2000-01-01_04-00-00"
-                ),
-            },
-        ]
-
-        actual_map = build_docdb_location_to_id_map(
-            docdb_api_client=mock_docdb_client,
-            bucket=bucket,
-            prefixes=[
-                "ecephys_655019_2000-04-04_04-00-00",
-                "ecephys_567890_2000-01-01_04-00-00/",
-                "missing_655019_2000-01-01_01-01-02",
-            ],
-        )
-        expected_map = {
-            f"s3://{bucket}/ecephys_655019_2000-04-04_04-00-00": (
-                "70bcf356-985f-4a2a-8105-de900e35e788"
-            ),
-            f"s3://{bucket}/ecephys_567890_2000-01-01_04-00-00": (
-                "5ca4a951-d374-4f4b-8279-d570a35b2286"
-            ),
-        }
-        self.assertEqual(expected_map, actual_map)
-
     @patch("codeocean.data_asset.DataAssets.search_data_assets_iterator")
     def test_get_all_processed_codeocean_asset_records(
         self, mock_search_all_data_assets: MagicMock