bulk_write should be able to accept a generator

sleepyStick · sleepyStick · commit b93360180ded · 2025-04-03T15:30:01.000-07:00
diff --git a/pymongo/asynchronous/bulk.py b/pymongo/asynchronous/bulk.py
@@ -26,6 +26,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Generator,
     Iterator,
     Mapping,
     Optional,
@@ -72,7 +73,7 @@
 from pymongo.write_concern import WriteConcern
 
 if TYPE_CHECKING:
-    from pymongo.asynchronous.collection import AsyncCollection
+    from pymongo.asynchronous.collection import AsyncCollection, _WriteOp
     from pymongo.asynchronous.mongo_client import AsyncMongoClient
     from pymongo.asynchronous.pool import AsyncConnection
     from pymongo.typings import _DocumentOut, _DocumentType, _Pipeline
@@ -214,28 +215,45 @@ def add_delete(
             self.is_retryable = False
         self.ops.append((_DELETE, cmd))
 
-    def gen_ordered(self) -> Iterator[Optional[_Run]]:
+    def gen_ordered(self, requests) -> Iterator[Optional[_Run]]:
         """Generate batches of operations, batched by type of
         operation, in the order **provided**.
         """
         run = None
-        for idx, (op_type, operation) in enumerate(self.ops):
+        for idx, request in enumerate(requests):
+            try:
+                request._add_to_bulk(self)
+            except AttributeError:
+                raise TypeError(f"{request!r} is not a valid request") from None
+            (op_type, operation) = self.ops[idx]
             if run is None:
                 run = _Run(op_type)
             elif run.op_type != op_type:
                 yield run
                 run = _Run(op_type)
             run.add(idx, operation)
+        if run is None:
+            raise InvalidOperation("No operations to execute")
         yield run
 
-    def gen_unordered(self) -> Iterator[_Run]:
+    def gen_unordered(self, requests) -> Iterator[_Run]:
         """Generate batches of operations, batched by type of
         operation, in arbitrary order.
         """
         operations = [_Run(_INSERT), _Run(_UPDATE), _Run(_DELETE)]
-        for idx, (op_type, operation) in enumerate(self.ops):
+        for idx, request in enumerate(requests):
+            try:
+                request._add_to_bulk(self)
+            except AttributeError:
+                raise TypeError(f"{request!r} is not a valid request") from None
+            (op_type, operation) = self.ops[idx]
             operations[op_type].add(idx, operation)
-
+        if (
+            len(operations[_INSERT].ops) == 0
+            and len(operations[_UPDATE].ops) == 0
+            and len(operations[_DELETE].ops) == 0
+        ):
+            raise InvalidOperation("No operations to execute")
         for run in operations:
             if run.ops:
                 yield run
@@ -726,23 +744,22 @@ async def execute_no_results(
 
     async def execute(
         self,
+        generator: Generator[_WriteOp[_DocumentType]],
         write_concern: WriteConcern,
         session: Optional[AsyncClientSession],
         operation: str,
     ) -> Any:
         """Execute operations."""
-        if not self.ops:
-            raise InvalidOperation("No operations to execute")
         if self.executed:
             raise InvalidOperation("Bulk operations can only be executed once.")
         self.executed = True
         write_concern = write_concern or self.collection.write_concern
         session = _validate_session_write_concern(session, write_concern)
 
         if self.ordered:
-            generator = self.gen_ordered()
+            generator = self.gen_ordered(generator)
         else:
-            generator = self.gen_unordered()
+            generator = self.gen_unordered(generator)
 
         client = self.collection.database.client
         if not write_concern.acknowledged:
diff --git a/pymongo/asynchronous/collection.py b/pymongo/asynchronous/collection.py
@@ -23,6 +23,7 @@
     AsyncContextManager,
     Callable,
     Coroutine,
+    Generator,
     Generic,
     Iterable,
     Iterator,
@@ -699,7 +700,7 @@ async def _create(
     @_csot.apply
     async def bulk_write(
         self,
-        requests: Sequence[_WriteOp[_DocumentType]],
+        requests: Sequence[_WriteOp[_DocumentType]] | Generator[_WriteOp[_DocumentType]],
         ordered: bool = True,
         bypass_document_validation: Optional[bool] = None,
         session: Optional[AsyncClientSession] = None,
@@ -779,17 +780,12 @@ async def bulk_write(
 
         .. versionadded:: 3.0
         """
-        common.validate_list("requests", requests)
+        common.validate_list_or_generator("requests", requests)
 
         blk = _AsyncBulk(self, ordered, bypass_document_validation, comment=comment, let=let)
-        for request in requests:
-            try:
-                request._add_to_bulk(blk)
-            except AttributeError:
-                raise TypeError(f"{request!r} is not a valid request") from None
 
         write_concern = self._write_concern_for(session)
-        bulk_api_result = await blk.execute(write_concern, session, _Op.INSERT)
+        bulk_api_result = await blk.execute(requests, write_concern, session, _Op.INSERT)
         if bulk_api_result is not None:
             return BulkWriteResult(bulk_api_result, True)
         return BulkWriteResult({}, False)
diff --git a/pymongo/common.py b/pymongo/common.py
@@ -24,6 +24,7 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    Generator,
     Iterator,
     Mapping,
     MutableMapping,
@@ -530,6 +531,13 @@ def validate_list(option: str, value: Any) -> list:
     return value
 
 
+def validate_list_or_generator(option: str, value: Any) -> Union[list, Generator]:
+    """Validates that 'value' is a list or generator."""
+    if isinstance(value, Generator):
+        return value
+    return validate_list(option, value)
+
+
 def validate_list_or_none(option: Any, value: Any) -> Optional[list]:
     """Validates that 'value' is a list or None."""
     if value is None:
diff --git a/pymongo/synchronous/bulk.py b/pymongo/synchronous/bulk.py
@@ -26,6 +26,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Generator,
     Iterator,
     Mapping,
     Optional,
@@ -72,7 +73,7 @@
 from pymongo.write_concern import WriteConcern
 
 if TYPE_CHECKING:
-    from pymongo.synchronous.collection import Collection
+    from pymongo.synchronous.collection import Collection, _WriteOp
     from pymongo.synchronous.mongo_client import MongoClient
     from pymongo.synchronous.pool import Connection
     from pymongo.typings import _DocumentOut, _DocumentType, _Pipeline
@@ -214,28 +215,45 @@ def add_delete(
             self.is_retryable = False
         self.ops.append((_DELETE, cmd))
 
-    def gen_ordered(self) -> Iterator[Optional[_Run]]:
+    def gen_ordered(self, requests) -> Iterator[Optional[_Run]]:
         """Generate batches of operations, batched by type of
         operation, in the order **provided**.
         """
         run = None
-        for idx, (op_type, operation) in enumerate(self.ops):
+        for idx, request in enumerate(requests):
+            try:
+                request._add_to_bulk(self)
+            except AttributeError:
+                raise TypeError(f"{request!r} is not a valid request") from None
+            (op_type, operation) = self.ops[idx]
             if run is None:
                 run = _Run(op_type)
             elif run.op_type != op_type:
                 yield run
                 run = _Run(op_type)
             run.add(idx, operation)
+        if run is None:
+            raise InvalidOperation("No operations to execute")
         yield run
 
-    def gen_unordered(self) -> Iterator[_Run]:
+    def gen_unordered(self, requests) -> Iterator[_Run]:
         """Generate batches of operations, batched by type of
         operation, in arbitrary order.
         """
         operations = [_Run(_INSERT), _Run(_UPDATE), _Run(_DELETE)]
-        for idx, (op_type, operation) in enumerate(self.ops):
+        for idx, request in enumerate(requests):
+            try:
+                request._add_to_bulk(self)
+            except AttributeError:
+                raise TypeError(f"{request!r} is not a valid request") from None
+            (op_type, operation) = self.ops[idx]
             operations[op_type].add(idx, operation)
-
+        if (
+            len(operations[_INSERT].ops) == 0
+            and len(operations[_UPDATE].ops) == 0
+            and len(operations[_DELETE].ops) == 0
+        ):
+            raise InvalidOperation("No operations to execute")
         for run in operations:
             if run.ops:
                 yield run
@@ -724,23 +742,22 @@ def execute_no_results(
 
     def execute(
         self,
+        generator: Generator[_WriteOp[_DocumentType]],
         write_concern: WriteConcern,
         session: Optional[ClientSession],
         operation: str,
     ) -> Any:
         """Execute operations."""
-        if not self.ops:
-            raise InvalidOperation("No operations to execute")
         if self.executed:
             raise InvalidOperation("Bulk operations can only be executed once.")
         self.executed = True
         write_concern = write_concern or self.collection.write_concern
         session = _validate_session_write_concern(session, write_concern)
 
         if self.ordered:
-            generator = self.gen_ordered()
+            generator = self.gen_ordered(generator)
         else:
-            generator = self.gen_unordered()
+            generator = self.gen_unordered(generator)
 
         client = self.collection.database.client
         if not write_concern.acknowledged:
diff --git a/pymongo/synchronous/collection.py b/pymongo/synchronous/collection.py
@@ -22,6 +22,7 @@
     Any,
     Callable,
     ContextManager,
+    Generator,
     Generic,
     Iterable,
     Iterator,
@@ -698,7 +699,7 @@ def _create(
     @_csot.apply
     def bulk_write(
         self,
-        requests: Sequence[_WriteOp[_DocumentType]],
+        requests: Sequence[_WriteOp[_DocumentType]] | Generator[_WriteOp[_DocumentType]],
         ordered: bool = True,
         bypass_document_validation: Optional[bool] = None,
         session: Optional[ClientSession] = None,
@@ -778,17 +779,12 @@ def bulk_write(
 
         .. versionadded:: 3.0
         """
-        common.validate_list("requests", requests)
+        common.validate_list_or_generator("requests", requests)
 
         blk = _Bulk(self, ordered, bypass_document_validation, comment=comment, let=let)
-        for request in requests:
-            try:
-                request._add_to_bulk(blk)
-            except AttributeError:
-                raise TypeError(f"{request!r} is not a valid request") from None
 
         write_concern = self._write_concern_for(session)
-        bulk_api_result = blk.execute(write_concern, session, _Op.INSERT)
+        bulk_api_result = blk.execute(requests, write_concern, session, _Op.INSERT)
         if bulk_api_result is not None:
             return BulkWriteResult(bulk_api_result, True)
         return BulkWriteResult({}, False)
diff --git a/test/asynchronous/test_bulk.py b/test/asynchronous/test_bulk.py
@@ -299,6 +299,21 @@ async def test_numerous_inserts(self):
         self.assertEqual(n_docs, result.inserted_count)
         self.assertEqual(n_docs, await self.coll.count_documents({}))
 
+    async def test_numerous_inserts_generator(self):
+        # Ensure we don't exceed server's maxWriteBatchSize size limit.
+        n_docs = await async_client_context.max_write_batch_size + 100
+        requests = (InsertOne[dict]({}) for _ in range(n_docs))
+        result = await self.coll.bulk_write(requests, ordered=False)
+        self.assertEqual(n_docs, result.inserted_count)
+        self.assertEqual(n_docs, await self.coll.count_documents({}))
+
+        # Same with ordered bulk.
+        await self.coll.drop()
+        requests = (InsertOne[dict]({}) for _ in range(n_docs))
+        result = await self.coll.bulk_write(requests)
+        self.assertEqual(n_docs, result.inserted_count)
+        self.assertEqual(n_docs, await self.coll.count_documents({}))
+
     async def test_bulk_max_message_size(self):
         await self.coll.delete_many({})
         self.addAsyncCleanup(self.coll.delete_many, {})
diff --git a/test/test_bulk.py b/test/test_bulk.py
@@ -299,6 +299,21 @@ def test_numerous_inserts(self):
         self.assertEqual(n_docs, result.inserted_count)
         self.assertEqual(n_docs, self.coll.count_documents({}))
 
+    def test_numerous_inserts_generator(self):
+        # Ensure we don't exceed server's maxWriteBatchSize size limit.
+        n_docs = client_context.max_write_batch_size + 100
+        requests = (InsertOne[dict]({}) for _ in range(n_docs))
+        result = self.coll.bulk_write(requests, ordered=False)
+        self.assertEqual(n_docs, result.inserted_count)
+        self.assertEqual(n_docs, self.coll.count_documents({}))
+
+        # Same with ordered bulk.
+        self.coll.drop()
+        requests = (InsertOne[dict]({}) for _ in range(n_docs))
+        result = self.coll.bulk_write(requests)
+        self.assertEqual(n_docs, result.inserted_count)
+        self.assertEqual(n_docs, self.coll.count_documents({}))
+
     def test_bulk_max_message_size(self):
         self.coll.delete_many({})
         self.addCleanup(self.coll.delete_many, {})
