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

* fix: pass pymongo kwargs to the bulk writer * version: 0.1.1
BeanieODM · Nov 7, 2022 · 7bf74eb · 7bf74eb
1 parent 9651228
commit 7bf74eb
Show file tree

Hide file tree

Showing 10 changed files with 70 additions and 92 deletions.
diff --git a/README.md b/README.md
@@ -54,14 +54,14 @@ class Product(Document):
 
 
 
-# Beanie uses Pymongo client under the hood 
+# Bunnet uses Pymongo client under the hood 
 client = MongoClient("mongodb://user:pass@host:27017")
 
 # Initialize bunnet with the Product document class
 init_bunnet(database=client.db_name, document_models=[Product])
 
 chocolate = Category(name="Chocolate", description="A preparation of roasted and ground cacao seeds.")
-# Beanie documents work just like pydantic models
+# Bunnet documents work just like pydantic models
 tonybar = Product(name="Tony's", price=5.95, category=chocolate)
 # And can be inserted into the database
 tonybar.insert() 

diff --git a/bunnet/__init__.py b/bunnet/__init__.py
@@ -24,7 +24,7 @@
 from bunnet.odm.views import View
 from bunnet.odm.union_doc import UnionDoc
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"
 __all__ = [
     # ODM
     "Document",

diff --git a/bunnet/odm/bulk.py b/bunnet/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 @@ 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)
 
             obj_class.get_motor_collection().bulk_write(requests)  # type: ignore

diff --git a/bunnet/odm/queries/delete.py b/bunnet/odm/queries/delete.py
@@ -55,6 +55,7 @@ def run(self) -> Union[DeleteResult, None, Optional[DeleteResult]]:
                     operation=DeleteManyPyMongo,
                     first_query=self.find_query,
                     object_class=self.document_model,
+                    pymongo_kwargs=self.pymongo_kwargs,
                 )
             )
             return None
@@ -76,6 +77,7 @@ def run(self) -> Union[DeleteResult, None, Optional[DeleteResult]]:
                     operation=DeleteOnePyMongo,
                     first_query=self.find_query,
                     object_class=self.document_model,
+                    pymongo_kwargs=self.pymongo_kwargs,
                 )
             )
             return None
diff --git a/bunnet/odm/queries/find.py b/bunnet/odm/queries/find.py
@@ -137,7 +137,6 @@ def upsert(
         *args: Mapping[str, Any],
         on_insert: "DocType",
         session: Optional[ClientSession] = None,
-        bulk_writer: Optional[BulkWriter] = None,
         **pymongo_kwargs,
     ):
         """
@@ -159,7 +158,6 @@ def upsert(
             .upsert(
                 *args,
                 on_insert=on_insert,
-                bulk_writer=bulk_writer,
                 **pymongo_kwargs,
             )
             .set_session(session=self.session)
@@ -821,6 +819,7 @@ 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/bunnet/odm/queries/update.py b/bunnet/odm/queries/update.py
@@ -97,7 +97,6 @@ def upsert(
         *args: Mapping[str, Any],
         on_insert: "DocType",
         session: Optional[ClientSession] = None,
-        bulk_writer: Optional[BulkWriter] = None,
         **pymongo_kwargs,
     ) -> "UpdateQuery":
         """
@@ -111,9 +110,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
@@ -186,6 +183,7 @@ def _update(self):
                     first_query=self.find_query,
                     second_query=self.update_query,
                     object_class=self.document_model,
+                    pymongo_kwargs=self.pymongo_kwargs,
                 )
             )
 
@@ -234,5 +232,6 @@ 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 @@
 
 Bunnet project
 
+## [1.1.1] - 2022-11-07
+
+### Fix
+
+- Pass pymongo kwargs to the bulk writer
+
+### Implementation
+
+- PR <https://github.com/roman-right/bunnet/pull/3>
+
 ## [0.1.0] - 2022-11-05
 
 ### Feature

diff --git a/docs/index.md b/docs/index.md
@@ -54,14 +54,14 @@ class Product(Document):
 
 
 
-# Beanie uses Pymongo client under the hood 
+# Bunnet uses Pymongo client under the hood 
 client = MongoClient("mongodb://user:pass@host:27017")
 
 # Initialize bunnet with the Product document class
 init_bunnet(database=client.db_name, document_models=[Product])
 
 chocolate = Category(name="Chocolate", description="A preparation of roasted and ground cacao seeds.")
-# Beanie documents work just like pydantic models
+# Bunnet documents work just like pydantic models
 tonybar = Product(name="Tony's", price=5.95, category=chocolate)
 # And can be inserted into the database
 tonybar.insert() 

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "bunnet"
-version = "0.1.0"
+version = "0.1.1"
 description = "Synchronous Python ODM for MongoDB"
 authors = ["Roman <[email protected]>"]
 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 bunnet.odm.bulk import BulkWriter
 from bunnet.odm.operators.update.general import Set
-from tests.models import DocumentTestModel
+from tests.models import DocumentTestModel, SubDocument
 
 
 def test_insert(documents_not_inserted):
@@ -93,93 +93,58 @@ def test_replace(documents, document_not_inserted):
     )
 
 
-def test_upsert_find_many_not_found(documents, document_not_inserted):
-    documents(5)
-    document_not_inserted.test_int = -10000
-    with BulkWriter() as bulk_writer:
-        DocumentTestModel.find(DocumentTestModel.test_int < -1000).upsert(
-            {"$set": {DocumentTestModel.test_int: 0}},
-            on_insert=document_not_inserted,
-        ).run()
-
-        bulk_writer.commit()
-
-    assert len(DocumentTestModel.find_all().to_list()) == 6
-    assert (
-        len(
-            DocumentTestModel.find(
-                DocumentTestModel.test_int == -10000
-            ).to_list()
-        )
-        == 1
-    )
-
-
-def test_upsert_find_one_not_found(documents, document_not_inserted):
-    documents(5)
-    document_not_inserted.test_int = -10000
-    with BulkWriter() as bulk_writer:
-        DocumentTestModel.find_one(DocumentTestModel.test_int < -1000).upsert(
-            {"$set": {DocumentTestModel.test_int: 0}},
-            on_insert=document_not_inserted,
-        ).run()
-
-        bulk_writer.commit()
-
-    assert len(DocumentTestModel.find_all().to_list()) == 6
-    assert (
-        len(
-            DocumentTestModel.find(
-                DocumentTestModel.test_int == -10000
-            ).to_list()
-        )
-        == 1
-    )
+def test_internal_error(document):
+    with pytest.raises(BulkWriteError):
+        with BulkWriter() as bulk_writer:
+            DocumentTestModel.insert_one(document, bulk_writer=bulk_writer)
 
 
-def test_upsert_find_many_found(documents, document_not_inserted):
+def test_native_upsert_found(documents, document_not_inserted):
     documents(5)
+    document_not_inserted.test_int = -1000
     with BulkWriter() as bulk_writer:
-        DocumentTestModel.find(DocumentTestModel.test_int == 1).upsert(
-            {"$set": {DocumentTestModel.test_int: -10000}},
-            on_insert=document_not_inserted,
+        DocumentTestModel.find_one(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,
         ).run()
-
         bulk_writer.commit()
 
-    assert len(DocumentTestModel.find_all().to_list()) == 5
-    assert (
-        len(
-            DocumentTestModel.find(
-                DocumentTestModel.test_int == -10000
-            ).to_list()
-        )
-        == 1
-    )
+    doc = DocumentTestModel.find_one(DocumentTestModel.test_int == 1).run()
+    assert len(doc.test_list) == 4
 
 
-def test_upsert_find_one_found(documents, document_not_inserted):
+def test_native_upsert_not_found(documents, document_not_inserted):
     documents(5)
+    document_not_inserted.test_int = -1000
     with BulkWriter() as bulk_writer:
-        DocumentTestModel.find_one(DocumentTestModel.test_int == 1).upsert(
-            {"$set": {DocumentTestModel.test_int: -10000}},
-            on_insert=document_not_inserted,
+        DocumentTestModel.find_one(
+            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,
         ).run()
-
         bulk_writer.commit()
 
-    assert len(DocumentTestModel.find_all().to_list()) == 5
-    assert (
-        len(
-            DocumentTestModel.find(
-                DocumentTestModel.test_int == -10000
-            ).to_list()
-        )
-        == 1
-    )
-
-
-def test_internal_error(document):
-    with pytest.raises(BulkWriteError):
-        with BulkWriter() as bulk_writer:
-            DocumentTestModel.insert_one(document, bulk_writer=bulk_writer)
+    assert DocumentTestModel.count() == 6