SERVER-104999 port motor to pymongo async (#36697)

GitOrigin-RevId: 4f0df8ed64d03912ef2ef44ecced31b6e40b0c38

buildscripts/cost_model/ce_generate_data.py

@@ -140,7 +140,7 @@ async def dump_collections_to_json(db, dump_path, database_name, collections):
     load_file = open(Path(dump_path) / f"{database_name}.data", "w")
     load_file.write("// This is a generated file.\n")
     # Create an array named 'chunkNames' with all chunk file names to be loaded.
-    load_file.write(f'const chunkNames = [{",".join(all_chunk_names)}];')
+    load_file.write(f"const chunkNames = [{','.join(all_chunk_names)}];")
 
 
 async def generate_histograms(coll_template, coll, dump_path):
@@ -184,7 +184,7 @@ async def main():
         old_db_collections = await database_instance.database.list_collection_names()
         for coll_name in old_db_collections:
             collection = database_instance.database[coll_name]
-            collection.drop()
+            await collection.drop()
 
         generator = DataGenerator(database_instance, data_generator_config)
         await generator.populate_collections()

buildscripts/cost_model/data_generator.py

@@ -37,8 +37,8 @@ from typing import Sequence
 import pymongo
 from config import DataGeneratorConfig, WriteMode
 from database_instance import DatabaseInstance
-from motor.motor_asyncio import AsyncIOMotorCollection
 from pymongo import IndexModel
+from pymongo.asynchronous.collection import AsyncCollection
 from random_generator import DataType, RandomDistribution
 
 __all__ = ["DataGenerator"]
@@ -106,7 +106,7 @@ class DataGenerator:
             await task
 
         t1 = time.time()
-        print(f"\npopulate Collections took {t1-t0} s.")
+        print(f"\npopulate Collections took {t1 - t0} s.")
 
     def _generate_collection_infos(self):
         for coll_template in self.config.collection_templates:
@@ -130,9 +130,7 @@ class DataGenerator:
                     compound_indexes=coll_template.compound_indexes,
                 )
 
-    async def _populate_collection(
-        self, coll: AsyncIOMotorCollection, coll_info: CollectionInfo
-    ) -> None:
+    async def _populate_collection(self, coll: AsyncCollection, coll_info: CollectionInfo) -> None:
         print(f"\nGenerating ${coll_info.name} ...")
         batch_size = self.config.batch_size
         tasks = []
@@ -150,7 +148,7 @@ class DataGenerator:
 
 
 async def populate_batch(
-    coll: AsyncIOMotorCollection, documents_count: int, fields: Sequence[FieldInfo]
+    coll: AsyncCollection, documents_count: int, fields: Sequence[FieldInfo]
 ) -> None:
     """Generate collection data and write it to the collection."""
 
@@ -167,9 +165,7 @@ def generate_collection_data(documents_count: int, fields: Sequence[FieldInfo]):
     return documents
 
 
-async def create_single_field_indexes(
-    coll: AsyncIOMotorCollection, fields: Sequence[FieldInfo]
-) -> None:
+async def create_single_field_indexes(coll: AsyncCollection, fields: Sequence[FieldInfo]) -> None:
     """Create single-fields indexes on the given collection."""
 
     indexes = [IndexModel([(field.name, pymongo.ASCENDING)]) for field in fields if field.indexed]
@@ -178,7 +174,7 @@ async def create_single_field_indexes(
         print(f"create_single_field_indexes done. {[index.document for index in indexes]}")
 
 
-async def create_compound_indexes(coll: AsyncIOMotorCollection, coll_info: CollectionInfo) -> None:
+async def create_compound_indexes(coll: AsyncCollection, coll_info: CollectionInfo) -> None:
     """Create a coumpound indexes on the given collection."""
 
     indexes_spec = []

buildscripts/cost_model/database_instance.py

@@ -34,7 +34,7 @@ from contextlib import asynccontextmanager
 from typing import Any, Mapping, NewType, Sequence
 
 from config import DatabaseConfig, RestoreMode
-from motor.motor_asyncio import AsyncIOMotorClient
+from pymongo import AsyncMongoClient
 
 __all__ = ["DatabaseInstance", "Pipeline"]
 """MongoDB Aggregate's Pipeline"""
@@ -47,7 +47,7 @@ class DatabaseInstance:
     def __init__(self, config: DatabaseConfig) -> None:
         """Initialize wrapper."""
         self.config = config
-        self.client = AsyncIOMotorClient(config.connection_string)
+        self.client = AsyncMongoClient(config.connection_string)
         self.database = self.client[config.database_name]
 
     def __enter__(self):

buildscripts/cost_model/requirements.txt

@@ -5,13 +5,12 @@ fonttools==4.43.0
 joblib==1.3.2
 kiwisolver==1.4.5
 matplotlib==3.8.0
-motor==3.3.1
 numpy==1.26.0
 packaging==23.2
 pandas==2.1.1
 patsy==0.5.3
 Pillow==10.0.1
-pymongo==4.5.0
+pymongo==4.13.0
 pyparsing==3.1.1
 python-dateutil==2.8.2
 pytz==2023.3.post1