fix: pass pymongo kwargs to the bulk writer (#406)

* fix: pass pymongo kwargs to the bulk writer
BeanieODM · Nov 7, 2022 · 3ace946 · 3ace946
1 parent e96ed44
commit 3ace946
Show file tree

Hide file tree

Showing 9 changed files with 72 additions and 96 deletions.
diff --git a/beanie/__init__.py b/beanie/__init__.py
@@ -26,7 +26,7 @@
 from beanie.odm.views import View
 from beanie.odm.union_doc import UnionDoc
 
-__version__ = "1.15.0"
+__version__ = "1.15.1"
 __all__ = [
     # ODM
     "Document",

diff --git a/beanie/odm/bulk.py b/beanie/odm/bulk.py
@@ -1,6 +1,6 @@
 from typing import Dict, Any, List, Optional, Union, Type, Mapping
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from pymongo import (
     InsertOne,
     DeleteOne,
@@ -22,6 +22,7 @@ class Operation(BaseModel):
     ]
     first_query: Mapping[str, Any]
     second_query: Optional[Dict[str, Any]] = None
+    pymongo_kwargs: Dict[str, Any] = Field(default_factory=dict)
     object_class: Type
 
     class Config:
@@ -51,9 +52,11 @@ async def commit(self):
                         "All the operations should be for a single document model"
                     )
                 if op.operation in [InsertOne, DeleteOne]:
-                    query = op.operation(op.first_query)
+                    query = op.operation(op.first_query, **op.pymongo_kwargs)
                 else:
-                    query = op.operation(op.first_query, op.second_query)
+                    query = op.operation(
+                        op.first_query, op.second_query, **op.pymongo_kwargs
+                    )
                 requests.append(query)
 
             await obj_class.get_motor_collection().bulk_write(requests)  # type: ignore

diff --git a/beanie/odm/queries/delete.py b/beanie/odm/queries/delete.py
@@ -55,6 +55,7 @@ def __await__(
                     operation=DeleteManyPyMongo,
                     first_query=self.find_query,
                     object_class=self.document_model,
+                    pymongo_kwargs=self.pymongo_kwargs,
                 )
             )
             return None
@@ -84,6 +85,7 @@ def __await__(
                     operation=DeleteOnePyMongo,
                     first_query=self.find_query,
                     object_class=self.document_model,
+                    pymongo_kwargs=self.pymongo_kwargs,
                 )
             )
             return None
diff --git a/beanie/odm/queries/find.py b/beanie/odm/queries/find.py
@@ -141,7 +141,6 @@ def upsert(
         *args: Mapping[str, Any],
         on_insert: "DocType",
         session: Optional[ClientSession] = None,
-        bulk_writer: Optional[BulkWriter] = None,
         **pymongo_kwargs,
     ):
         """
@@ -163,7 +162,6 @@ def upsert(
             .upsert(
                 *args,
                 on_insert=on_insert,
-                bulk_writer=bulk_writer,
                 **pymongo_kwargs,
             )
             .set_session(session=self.session)
@@ -827,6 +825,7 @@ async def replace_one(
                         by_alias=True, exclude={"_id"}
                     ).encode(document),
                     object_class=self.document_model,
+                    pymongo_kwargs=self.pymongo_kwargs,
                 )
             )
             return None

diff --git a/beanie/odm/queries/update.py b/beanie/odm/queries/update.py
@@ -98,7 +98,6 @@ def upsert(
         *args: Mapping[str, Any],
         on_insert: "DocType",
         session: Optional[ClientSession] = None,
-        bulk_writer: Optional[BulkWriter] = None,
         **pymongo_kwargs,
     ) -> "UpdateQuery":
         """
@@ -112,9 +111,7 @@ def upsert(
         :return: UpdateMany query
         """
         self.upsert_insert_doc = on_insert  # type: ignore
-        self.update(
-            *args, session=session, bulk_writer=bulk_writer, **pymongo_kwargs
-        )
+        self.update(*args, session=session, **pymongo_kwargs)
         return self
 
     @abstractmethod
@@ -191,6 +188,7 @@ async def _update(self):
                     first_query=self.find_query,
                     second_query=self.update_query,
                     object_class=self.document_model,
+                    pymongo_kwargs=self.pymongo_kwargs,
                 )
             )
 
@@ -239,5 +237,6 @@ async def _update(self):
                     first_query=self.find_query,
                     second_query=self.update_query,
                     object_class=self.document_model,
+                    pymongo_kwargs=self.pymongo_kwargs,
                 )
             )
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -2,6 +2,16 @@
 
 Beanie project
 
+## [1.15.1] - 2022-11-07
+
+### Fix
+
+- Pass pymongo kwargs to the bulk writer
+
+### Implementation
+
+- PR <https://github.com/roman-right/beanie/pull/406>
+
 ## [1.15.0] - 2022-11-05
 
 ### Feature
@@ -1055,4 +1065,6 @@ how specific type should be presented in the database
 
 [1.14.0]: https://pypi.org/project/beanie/1.14.0
 
-[1.15.0]: https://pypi.org/project/beanie/1.15.0
+[1.15.0]: https://pypi.org/project/beanie/1.15.0
+
+[1.15.1]: https://pypi.org/project/beanie/1.15.1
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "beanie"
-version = "1.15.0"
+version = "1.15.1"
 description = "Asynchronous Python ODM for MongoDB"
 authors = ["Roman <roman-right@protonmail.com>"]
 license = "Apache-2.0"

diff --git a/tests/odm/documents/test_bulk_write.py b/tests/odm/documents/test_bulk_write.py
@@ -3,7 +3,7 @@
 
 from beanie.odm.bulk import BulkWriter
 from beanie.odm.operators.update.general import Set
-from tests.odm.models import DocumentTestModel
+from tests.odm.models import DocumentTestModel, SubDocument
 
 
 async def test_insert(documents_not_inserted):
@@ -101,101 +101,62 @@ async def test_replace(documents, document_not_inserted):
     )
 
 
-async def test_upsert_find_many_not_found(documents, document_not_inserted):
-    await documents(5)
-    document_not_inserted.test_int = -10000
-    async with BulkWriter() as bulk_writer:
-        await DocumentTestModel.find(
-            DocumentTestModel.test_int < -1000
-        ).upsert(
-            {"$set": {DocumentTestModel.test_int: 0}},
-            on_insert=document_not_inserted,
-        )
-
-        await bulk_writer.commit()
-
-    assert len(await DocumentTestModel.find_all().to_list()) == 6
-    assert (
-        len(
-            await DocumentTestModel.find(
-                DocumentTestModel.test_int == -10000
-            ).to_list()
-        )
-        == 1
-    )
+async def test_internal_error(document):
+    with pytest.raises(BulkWriteError):
+        async with BulkWriter() as bulk_writer:
+            await DocumentTestModel.insert_one(
+                document, bulk_writer=bulk_writer
+            )
 
 
-async def test_upsert_find_one_not_found(documents, document_not_inserted):
+async def test_native_upsert_found(documents, document_not_inserted):
     await documents(5)
-    document_not_inserted.test_int = -10000
+    document_not_inserted.test_int = -1000
     async with BulkWriter() as bulk_writer:
         await DocumentTestModel.find_one(
-            DocumentTestModel.test_int < -1000
-        ).upsert(
-            {"$set": {DocumentTestModel.test_int: 0}},
-            on_insert=document_not_inserted,
-        )
-
-        await bulk_writer.commit()
-
-    assert len(await DocumentTestModel.find_all().to_list()) == 6
-    assert (
-        len(
-            await DocumentTestModel.find(
-                DocumentTestModel.test_int == -10000
-            ).to_list()
-        )
-        == 1
-    )
-
-
-async def test_upsert_find_many_found(documents, document_not_inserted):
-    await documents(5)
-    async with BulkWriter() as bulk_writer:
-        await DocumentTestModel.find(DocumentTestModel.test_int == 1).upsert(
-            {"$set": {DocumentTestModel.test_int: -10000}},
-            on_insert=document_not_inserted,
+            DocumentTestModel.test_int == 1
+        ).update_one(
+            {
+                "$addToSet": {
+                    "test_list": {
+                        "$each": [
+                            SubDocument(test_str="TEST_ONE"),
+                            SubDocument(test_str="TEST_TWO"),
+                        ]
+                    }
+                },
+                "$setOnInsert": {},
+            },
+            bulk_writer=bulk_writer,
+            upsert=True,
         )
-
         await bulk_writer.commit()
 
-    assert len(await DocumentTestModel.find_all().to_list()) == 5
-    assert (
-        len(
-            await DocumentTestModel.find(
-                DocumentTestModel.test_int == -10000
-            ).to_list()
-        )
-        == 1
-    )
+    doc = await DocumentTestModel.find_one(DocumentTestModel.test_int == 1)
+    assert len(doc.test_list) == 4
 
 
-async def test_upsert_find_one_found(documents, document_not_inserted):
+async def test_native_upsert_not_found(documents, document_not_inserted):
     await documents(5)
+    document_not_inserted.test_int = -1000
     async with BulkWriter() as bulk_writer:
         await DocumentTestModel.find_one(
-            DocumentTestModel.test_int == 1
-        ).upsert(
-            {"$set": {DocumentTestModel.test_int: -10000}},
-            on_insert=document_not_inserted,
+            DocumentTestModel.test_int == -1000
+        ).update_one(
+            {
+                "$addToSet": {
+                    "test_list": {
+                        "$each": [
+                            SubDocument(test_str="TEST_ONE"),
+                            SubDocument(test_str="TEST_TWO"),
+                        ]
+                    }
+                },
+                "$setOnInsert": {"TEST": "VALUE"},
+            },
+            bulk_writer=bulk_writer,
+            upsert=True,
         )
-
         await bulk_writer.commit()
 
-    assert len(await DocumentTestModel.find_all().to_list()) == 5
-    assert (
-        len(
-            await DocumentTestModel.find(
-                DocumentTestModel.test_int == -10000
-            ).to_list()
-        )
-        == 1
-    )
-
-
-async def test_internal_error(document):
-    with pytest.raises(BulkWriteError):
-        async with BulkWriter() as bulk_writer:
-            await DocumentTestModel.insert_one(
-                document, bulk_writer=bulk_writer
-            )
+    assert await DocumentTestModel.count() == 6
diff --git a/tests/test_beanie.py b/tests/test_beanie.py
@@ -2,4 +2,4 @@
 
 
 def test_version():
-    assert __version__ == "1.15.0"
+    assert __version__ == "1.15.1"