From 7bf74eb39656083ac4ef8b8ef4ed14b0e354fe61 Mon Sep 17 00:00:00 2001 From: Roman Right Date: Mon, 7 Nov 2022 12:47:55 +0100 Subject: [PATCH] fix: pass pymongo kwargs to the bulk writer (#3) * fix: pass pymongo kwargs to the bulk writer * version: 0.1.1 --- README.md | 4 +- bunnet/__init__.py | 2 +- bunnet/odm/bulk.py | 9 +- bunnet/odm/queries/delete.py | 2 + bunnet/odm/queries/find.py | 3 +- bunnet/odm/queries/update.py | 7 +- docs/changelog.md | 10 +++ docs/index.md | 4 +- pyproject.toml | 2 +- tests/odm/documents/test_bulk_write.py | 119 +++++++++---------------- 10 files changed, 70 insertions(+), 92 deletions(-) diff --git a/README.md b/README.md index f1193a5..850fb1d 100644 --- 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 index de04ca8..4779269 100644 --- 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 index 9b67444..1bb737f 100644 --- 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 index f063284..67d457c 100644 --- 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 index 529dd1b..18cfc0a 100644 --- 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 index 698377c..9cb4b9a 100644 --- 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 index 4afc5dc..39ee8ba 100644 --- 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 + ## [0.1.0] - 2022-11-05 ### Feature diff --git a/docs/index.md b/docs/index.md index f1193a5..850fb1d 100644 --- 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 index a4147d1..f9df7e3 100644 --- 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 "] license = "Apache-2.0" diff --git a/tests/odm/documents/test_bulk_write.py b/tests/odm/documents/test_bulk_write.py index 65298cc..4c4ba8b 100644 --- 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