Merge pull request #115 from AllenNeuralDynamics/release-v0.14.0

Release v0.14.0

Author: jtyoung84

.github/workflows/publish_dev.yml

@@ -8,7 +8,7 @@ jobs:
   publish:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Docker Buildx
         id: buildx
         uses: docker/setup-buildx-action@v2

.github/workflows/publish_main.yml

@@ -11,7 +11,7 @@ jobs:
     outputs:
       pkg_version: ${{ steps.output_version.outputs.pkg_version }}
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Get version from file
       run: |
         pkg_name=$(grep -P 'version = \{attr = .*\}' pyproject.toml | grep -oP '\w+.__version__')

.github/workflows/run_dev_tests.yml

@@ -10,11 +10,11 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ '3.8', '3.9', '3.10' ]
+        python-version: [ '3.9', '3.10', '3.11' ]
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies

.github/workflows/run_main_tests.yml

@@ -11,11 +11,11 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ '3.8', '3.9', '3.10' ]
+        python-version: [ '3.9', '3.10', '3.11' ]
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
@@ -28,7 +28,7 @@ jobs:
   verify_version:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Check version incremented
         run: |
           pkg_name=$(grep -P 'version = \{attr = .*\}' pyproject.toml | grep -oP '\w+.__version__')

README.md

@@ -12,7 +12,7 @@ kept in sync:
 1. **S3 buckets** store raw metadata files, including the ``metadata.nd.json``.
 2. A **document database (DocDB)** contains unstructured json
    documents describing the ``metadata.nd.json`` for a data asset.
-3. **Code Ocean**: data assets are mounted as CodeOcean data asssets.
+3. **Code Ocean**: data assets are mounted as CodeOcean data assets.
    Processed results are also stored in an internal Code Ocean bucket.
 
 We have automated jobs to keep changes in DocDB and S3 in sync.

docs/source/conf.py

@@ -1,4 +1,5 @@
 """Configuration file for the Sphinx documentation builder."""
+
 #
 # For the full list of built-in configuration values, see the documentation:
 # https://www.sphinx-doc.org/en/master/usage/configuration.html

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "aind-data-asset-indexer"
 description = "Service Capsule to write data asset metadata to document store"
 license = {text = "MIT"}
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 authors = [
     {name = "AIND"}
 ]
@@ -24,7 +24,7 @@ dependencies = [
     "pymongo==4.3.3",
     "dask==2023.5.0",
     "aind-data-schema==1.2.0",
-    "aind-codeocean-api==0.5.0",
+    "codeocean==0.3.0",
 ]
 
 [project.optional-dependencies]

src/aind_data_asset_indexer/__init__.py

@@ -1,3 +1,3 @@
 """Package"""
 
-__version__ = "0.13.0"
+__version__ = "0.14.0"

src/aind_data_asset_indexer/aind_bucket_indexer.py

@@ -272,14 +272,12 @@ def _resolve_schema_information(
                 object_key = create_object_key(
                     prefix=prefix, filename=core_schema_file_name
                 )
-                common_kwargs[
-                    "core_schema_info_in_root"
-                ] = get_dict_of_file_info(
-                    s3_client=s3_client,
-                    bucket=self.job_settings.s3_bucket,
-                    keys=[object_key],
-                ).get(
-                    object_key
+                common_kwargs["core_schema_info_in_root"] = (
+                    get_dict_of_file_info(
+                        s3_client=s3_client,
+                        bucket=self.job_settings.s3_bucket,
+                        keys=[object_key],
+                    ).get(object_key)
                 )
                 self._copy_file_from_root_to_subdir(**common_kwargs)
             # If field is null, a file exists in the root folder, and
@@ -424,9 +422,9 @@ def _process_docdb_record(
                     )
                     db = docdb_client[self.job_settings.doc_db_db_name]
                     collection = db[self.job_settings.doc_db_collection_name]
-                    fields_to_update[
-                        "last_modified"
-                    ] = datetime.utcnow().isoformat()
+                    fields_to_update["last_modified"] = (
+                        datetime.utcnow().isoformat()
+                    )
                     response = collection.update_one(
                         {"_id": docdb_record["_id"]},
                         {"$set": fields_to_update},

src/aind_data_asset_indexer/codeocean_bucket_indexer.py

@@ -12,13 +12,13 @@
 
 import boto3
 import dask.bag as dask_bag
-import requests
-from aind_codeocean_api.codeocean import CodeOceanClient
 from aind_data_schema.core.metadata import ExternalPlatforms
+from codeocean import CodeOcean
+from codeocean.data_asset import DataAssetSearchOrigin, DataAssetSearchParams
 from mypy_boto3_s3 import S3Client
 from pymongo import MongoClient
 from pymongo.operations import UpdateOne
-from requests.exceptions import ReadTimeout
+from urllib3.util import Retry
 
 from aind_data_asset_indexer.models import CodeOceanIndexBucketJobSettings
 from aind_data_asset_indexer.utils import (
@@ -52,30 +52,51 @@ def __init__(self, job_settings: CodeOceanIndexBucketJobSettings):
         """Class constructor."""
         self.job_settings = job_settings
 
-    def _get_external_data_asset_records(self) -> Optional[List[dict]]:
+    @staticmethod
+    def _get_external_data_asset_records(
+        co_client: CodeOcean,
+    ) -> Optional[List[dict]]:
         """
         Retrieves list of code ocean ids and locations for external data
         assets. The timeout is set to 600 seconds.
+
+        Parameters
+        ----------
+        co_client : CodeOcean
+
         Returns
         -------
         List[dict] | None
           List items have shape {"id": str, "location": str}. If error occurs,
           return None.
+
         """
         try:
-            response = requests.get(
-                self.job_settings.temp_codeocean_endpoint,
-                timeout=600,
+            search_params = DataAssetSearchParams(
+                archived=False,
+                origin=DataAssetSearchOrigin.External,
+                limit=1000,
             )
-            if response.status_code == 200:
-                return response.json()
-            else:
-                return None
-        except ReadTimeout:
-            logging.error(
-                f"Read timed out at "
-                f"{self.job_settings.temp_codeocean_endpoint}"
+            data_assets = co_client.data_assets.search_data_assets_iterator(
+                search_params=search_params
             )
+            external_records = []
+            for data_asset in data_assets:
+                data_asset_source = data_asset.source_bucket
+                if (
+                    data_asset_source is not None
+                    and data_asset_source.bucket is not None
+                    and data_asset_source.prefix is not None
+                ):
+                    bucket = data_asset_source.bucket
+                    prefix = data_asset_source.prefix
+                    location = f"s3://{bucket}/{prefix}"
+                    external_records.append(
+                        {"id": data_asset.id, "location": location}
+                    )
+            return external_records
+        except Exception as e:
+            logging.exception(e)
             return None
 
     @staticmethod
@@ -97,7 +118,7 @@ def _map_external_list_to_dict(external_recs: List[dict]) -> dict:
         """
         new_records = dict()
         for r in external_recs:
-            location = r.get("source")
+            location = r.get("location")
             rec_id = r["id"]
             if location is not None and new_records.get(location) is not None:
                 old_id_set = new_records.get(location)
@@ -140,7 +161,7 @@ def _get_co_links_from_record(
         return external_links
 
     def _update_external_links_in_docdb(
-        self, docdb_client: MongoClient
+        self, docdb_client: MongoClient, co_client: CodeOcean
     ) -> None:
         """
         This method will:
@@ -159,7 +180,9 @@ def _update_external_links_in_docdb(
 
         """
         # Should return a list like [{"id": co_id, "location": "s3://..."},]
-        list_of_co_ids_and_locations = self._get_external_data_asset_records()
+        list_of_co_ids_and_locations = self._get_external_data_asset_records(
+            co_client=co_client
+        )
         db = docdb_client[self.job_settings.doc_db_db_name]
         collection = db[self.job_settings.doc_db_collection_name]
         if list_of_co_ids_and_locations is not None:
@@ -394,9 +417,16 @@ def _delete_records_from_docdb(self, record_list: List[str]):
     def run_job(self):
         """Main method to run."""
         logging.info("Starting to scan through CodeOcean.")
-        co_client = CodeOceanClient(
+        retry = Retry(
+            total=5,
+            backoff_factor=1,
+            status_forcelist=[429, 500, 502, 503, 504],
+            allowed_methods=["GET", "POST"],
+        )
+        co_client = CodeOcean(
             domain=self.job_settings.codeocean_domain,
             token=self.job_settings.codeocean_token.get_secret_value(),
+            retries=retry,
         )
         code_ocean_records = get_all_processed_codeocean_asset_records(
             co_client=co_client,
@@ -416,7 +446,7 @@ def run_job(self):
         # Use existing client to add external links to fields
         logging.info("Adding links to records.")
         self._update_external_links_in_docdb(
-            docdb_client=iterator_docdb_client
+            docdb_client=iterator_docdb_client, co_client=co_client
         )
         logging.info("Finished adding links to records")
         all_docdb_records = dict()