Migrate from MongoClient to AsyncMongoClient

README.md

@@ -137,6 +137,7 @@ of like poetry dev.dependencies stanza. main.txt and dev.txt are kind of like po
 versions of dependencies to use. main.txt and dev.txt are combined in the docker compose build process to create the 
 final requirements.txt file and import the dependencies into the Docker image.
 
+
 ## Local Testing
 
 Tests can be found in `tests` and are run with the following commands:

nmdc_runtime/api/core/idgen.py

@@ -73,9 +73,9 @@ def generate_ids(
     shoulder: str = "fk4",
 ) -> List[str]:
     collection = mdb.get_collection(collection_name(naa, shoulder))
-    existing_count = collection.count_documents({})
+    initial_count = collection.count_documents({})
     n_chars = next(
-        (n for n, t in SPING_SIZE_THRESHOLDS if (number + existing_count) < t),
+        (n for n, t in SPING_SIZE_THRESHOLDS if (number + initial_count) < t),
         12,
     )
     collected = []

nmdc_runtime/api/db/mongo.py

@@ -10,7 +10,6 @@
 from jsonschema import Draft7Validator
 from nmdc_schema.nmdc import Database as NMDCDatabase
 from pymongo.errors import AutoReconnect, OperationFailure
-from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
 from refscan.lib.Finder import Finder
 from refscan.scanner import scan_outgoing_references
 from tenacity import wait_random_exponential, retry, retry_if_exception_type
@@ -27,15 +26,16 @@
     nmdc_database_collection_names,
     get_allowed_references,
 )
-from pymongo import MongoClient
+from pymongo import AsyncMongoClient, MongoClient
 from pymongo.database import Database as MongoDatabase
+from pymongo.asynchronous.database import AsyncDatabase
 
 
 @retry(
     retry=retry_if_exception_type(AutoReconnect),
     wait=wait_random_exponential(multiplier=0.5, max=60),
 )
-def check_mongo_ok_autoreconnect(mdb: MongoDatabase):
+def check_mongo_ok_autoreconnect(mdb: MongoDatabase) -> bool:
     mdb["_runtime.healthcheck"].insert_one({"_id": "ok"})
     mdb["_runtime.healthcheck"].delete_one({"_id": "ok"})
     return True
@@ -55,14 +55,29 @@ def get_mongo_client() -> MongoClient:
     )
 
 
+@lru_cache
+def get_async_mongo_client() -> AsyncMongoClient:
+    r"""
+    Returns an `AsyncMongoClient` instance you can use to access the MongoDB server specified via environment variables.
+    """
+    return AsyncMongoClient(
+        host=os.getenv("MONGO_HOST"),
+        username=os.getenv("MONGO_USERNAME"),
+        password=os.getenv("MONGO_PASSWORD"),
+        directConnection=True,
+    )
+
+
 @lru_cache
 def get_mongo_db() -> MongoDatabase:
     r"""
     Returns a `Database` instance you can use to access the MongoDB database specified via an environment variable.
     Reference: https://pymongo.readthedocs.io/en/stable/api/pymongo/database.html#pymongo.database.Database
     """
     _client = get_mongo_client()
-    mdb = _client[os.getenv("MONGO_DBNAME")]
+    database_name = os.getenv("MONGO_DBNAME")
+    assert database_name is not None, "MONGO_DBNAME is None"
+    mdb = _client[database_name]
     check_mongo_ok_autoreconnect(mdb)
     return mdb
 
@@ -74,20 +89,20 @@ def get_session_bound_mongo_db(session=None) -> MongoDatabase:
     Reference: https://pymongo.readthedocs.io/en/stable/api/pymongo/database.html#pymongo.database.Database
     """
     _client = get_mongo_client()
-    mdb = _client[os.getenv("MONGO_DBNAME")]
+    database_name = os.getenv("MONGO_DBNAME")
+    assert database_name is not None, "MONGO_DBNAME is None"
+    mdb = _client[database_name]
     check_mongo_ok_autoreconnect(mdb)
     return SessionBoundDatabase(mdb, session) if session is not None else mdb
 
 
 @lru_cache
-def get_async_mongo_db() -> AsyncIOMotorDatabase:
-    _client = AsyncIOMotorClient(
-        host=os.getenv("MONGO_HOST"),
-        username=os.getenv("MONGO_USERNAME"),
-        password=os.getenv("MONGO_PASSWORD"),
-        directConnection=True,
-    )
-    return _client[os.getenv("MONGO_DBNAME")]
+def get_async_mongo_db() -> AsyncDatabase:
+    _client = get_async_mongo_client()
+    database_name = os.getenv("MONGO_DBNAME")
+    assert database_name is not None, "MONGO_DBNAME is None"
+    mdb = _client[database_name]
+    return mdb
 
 
 def get_nonempty_nmdc_schema_collection_names(mdb: MongoDatabase) -> Set[str]:

nmdc_runtime/api/endpoints/jobs.py

@@ -1,6 +1,7 @@
 import json
 from typing import Optional, Annotated
 
+from pymongo.asynchronous.database import AsyncDatabase
 from pymongo.database import Database
 from fastapi import APIRouter, Depends, Query, HTTPException, Path
 from pymongo.errors import ConnectionFailure, OperationFailure
@@ -9,7 +10,7 @@
 from nmdc_runtime.api.core.util import (
     raise404_if_none,
 )
-from nmdc_runtime.api.db.mongo import get_mongo_db
+from nmdc_runtime.api.db.mongo import get_async_mongo_db, get_mongo_db
 from nmdc_runtime.api.endpoints.util import list_resources, _claim_job
 from nmdc_runtime.api.models.job import Job, JobClaim
 from nmdc_runtime.api.models.operation import Operation, MetadataT
@@ -26,9 +27,9 @@
 @router.get(
     "/jobs", response_model=ListResponse[Job], response_model_exclude_unset=True
 )
-def list_jobs(
+async def list_jobs(
     req: Annotated[ListRequest, Query()],
-    mdb: Database = Depends(get_mongo_db),
+    adb: AsyncDatabase = Depends(get_async_mongo_db),
     maybe_site: Optional[Site] = Depends(maybe_get_current_client_site),
 ):
     """List pre-configured workflow jobs.
@@ -39,7 +40,7 @@ def list_jobs(
     """
     if isinstance(maybe_site, Site) and req.filter is None:
         req.filter = json.dumps({"claims.site_id": {"$ne": maybe_site.id}})
-    return list_resources(req, mdb, "jobs")
+    return await list_resources(req, adb, "jobs")
 
 
 @router.get("/jobs/{job_id}", response_model=Job, response_model_exclude_unset=True)

nmdc_runtime/api/endpoints/nmdcschema.py

@@ -3,6 +3,7 @@
 from typing import List, Dict, Annotated
 
 import pymongo
+from pymongo.asynchronous.database import AsyncDatabase
 from fastapi import APIRouter, Depends, HTTPException, Path, Query
 from pydantic import AfterValidator
 from refscan.lib.helpers import (
@@ -24,6 +25,7 @@
 from nmdc_runtime.api.core.metadata import map_id_to_collection, get_collection_for_id
 from nmdc_runtime.api.core.util import raise404_if_none
 from nmdc_runtime.api.db.mongo import (
+    get_async_mongo_db,
     get_mongo_db,
 )
 from nmdc_runtime.api.endpoints.util import (
@@ -485,7 +487,7 @@ def get_collection_names():
     response_model=ListResponse[Doc],
     response_model_exclude_unset=True,
 )
-def list_from_collection(
+async def list_from_collection(
     collection_name: Annotated[
         str,
         Path(
@@ -495,7 +497,7 @@ def list_from_collection(
         ),
     ],
     req: Annotated[ListRequest, Query()],
-    mdb: MongoDatabase = Depends(get_mongo_db),
+    adb: AsyncDatabase = Depends(get_async_mongo_db),
 ):
     r"""
     Retrieves resources that match the specified filter criteria and reside in the specified collection.
@@ -521,7 +523,7 @@ def list_from_collection(
     # raise HTTP_400_BAD_REQUEST on invalid collection_name
     ensure_collection_name_is_known_to_schema(collection_name)
 
-    rv = list_resources(req, mdb, collection_name)
+    rv = await list_resources(req, adb, collection_name)
     rv["resources"] = [strip_oid(d) for d in rv["resources"]]
     return rv
 

nmdc_runtime/api/endpoints/objects.py

@@ -4,14 +4,15 @@
 from fastapi import APIRouter, status, Depends, HTTPException, Query
 from gridfs import GridFS
 from pymongo import ReturnDocument
+from pymongo.asynchronous.database import AsyncDatabase
 from pymongo.database import Database as MongoDatabase
 import requests
 from starlette.responses import RedirectResponse
 from toolz import merge
 
 from nmdc_runtime.api.core.idgen import decode_id, generate_one_id, local_part
 from nmdc_runtime.api.core.util import raise404_if_none, API_SITE_ID
-from nmdc_runtime.api.db.mongo import get_mongo_db
+from nmdc_runtime.api.db.mongo import get_async_mongo_db, get_mongo_db
 from nmdc_runtime.api.db.s3 import S3_ID_NS, presigned_url_to_get, get_s3_client
 from nmdc_runtime.api.endpoints.util import (
     list_resources,
@@ -90,11 +91,11 @@ def create_object(
 
 
 @router.get("/objects", response_model=ListResponse[DrsObject])
-def list_objects(
+async def list_objects(
     req: Annotated[ListRequest, Query()],
-    mdb: MongoDatabase = Depends(get_mongo_db),
+    adb: AsyncDatabase = Depends(get_async_mongo_db),
 ):
-    return list_resources(req, mdb, "objects")
+    return await list_resources(req, adb, "objects")
 
 
 @router.get(

nmdc_runtime/api/endpoints/operations.py

@@ -1,11 +1,12 @@
 from typing import Annotated
 
 import pymongo
+from pymongo.asynchronous.database import AsyncDatabase
 from fastapi import APIRouter, Depends, status, HTTPException, Query
 from toolz import get_in, merge, assoc
 
 from nmdc_runtime.api.core.util import raise404_if_none, pick
-from nmdc_runtime.api.db.mongo import get_mongo_db
+from nmdc_runtime.api.db.mongo import get_async_mongo_db, get_mongo_db
 from nmdc_runtime.api.endpoints.util import list_resources
 from nmdc_runtime.api.models.operation import (
     ListOperationsResponse,
@@ -21,11 +22,11 @@
 
 
 @router.get("/operations", response_model=ListOperationsResponse[ResultT, MetadataT])
-def list_operations(
+async def list_operations(
     req: Annotated[ListRequest, Query()],
-    mdb: pymongo.database.Database = Depends(get_mongo_db),
+    adb: AsyncDatabase = Depends(get_async_mongo_db),
 ):
-    return list_resources(req, mdb, "operations")
+    return await list_resources(req, adb, "operations")
 
 
 @router.get("/operations/{op_id}", response_model=Operation[ResultT, MetadataT])

nmdc_runtime/api/endpoints/search.py

@@ -1,9 +1,9 @@
 import json
 
 from fastapi import APIRouter, Depends
-from pymongo.database import Database as MongoDatabase
+from pymongo.asynchronous.database import AsyncDatabase
 
-from nmdc_runtime.api.db.mongo import get_mongo_db
+from nmdc_runtime.api.db.mongo import get_async_mongo_db
 from nmdc_runtime.api.endpoints.nmdcschema import strip_oid
 from nmdc_runtime.api.endpoints.util import list_resources
 from nmdc_runtime.api.models.nmdc_schema import (
@@ -21,9 +21,9 @@
     response_model=ListResponse[DataObject],
     response_model_exclude_unset=True,
 )
-def data_objects(
+async def data_objects(
     req: DataObjectListRequest = Depends(),
-    mdb: MongoDatabase = Depends(get_mongo_db),
+    adb: AsyncDatabase = Depends(get_async_mongo_db),
 ):
     filter_ = list_request_filter_to_mongo_filter(req.model_dump(exclude_unset=True))
     max_page_size = filter_.pop("max_page_size", None)
@@ -33,6 +33,6 @@ def data_objects(
         max_page_size=max_page_size,
         page_token=page_token,
     )
-    rv = list_resources(req, mdb, "data_objects")
+    rv = await list_resources(req, adb, "data_objects")
     rv["resources"] = [strip_oid(d) for d in rv["resources"]]
     return rv

nmdc_runtime/api/endpoints/sites.py

@@ -1,6 +1,7 @@
 from typing import List, Annotated
 
 import botocore
+from pymongo.asynchronous.database import AsyncDatabase
 import pymongo.database
 from fastapi import APIRouter, Depends, status, HTTPException, Path, Query
 from starlette.status import HTTP_403_FORBIDDEN
@@ -17,7 +18,7 @@
     generate_secret,
     API_SITE_ID,
 )
-from nmdc_runtime.api.db.mongo import get_mongo_db
+from nmdc_runtime.api.db.mongo import get_async_mongo_db, get_mongo_db
 from nmdc_runtime.api.db.s3 import (
     get_s3_client,
     presigned_url_to_put,
@@ -72,11 +73,11 @@ def create_site(
 @router.get(
     "/sites", response_model=ListResponse[Site], response_model_exclude_unset=True
 )
-def list_sites(
+async def list_sites(
     req: Annotated[ListRequest, Query()],
-    mdb: pymongo.database.Database = Depends(get_mongo_db),
+    adb: AsyncDatabase = Depends(get_async_mongo_db),
 ):
-    return list_resources(req, mdb, "sites")
+    return await list_resources(req, adb, "sites")
 
 
 @router.get("/sites/{site_id}", response_model=Site, response_model_exclude_unset=True)