From 3ace9467dae5a8c120819cb08ff166227dcc3fe4 Mon Sep 17 00:00:00 2001 From: Roman Right Date: Mon, 7 Nov 2022 12:47:47 +0100 Subject: [PATCH] fix: pass pymongo kwargs to the bulk writer (#406) * fix: pass pymongo kwargs to the bulk writer --- beanie/__init__.py | 2 +- beanie/odm/bulk.py | 9 +- beanie/odm/queries/delete.py | 2 + beanie/odm/queries/find.py | 3 +- beanie/odm/queries/update.py | 7 +- docs/changelog.md | 14 ++- pyproject.toml | 2 +- tests/odm/documents/test_bulk_write.py | 127 +++++++++---------------- tests/test_beanie.py | 2 +- 9 files changed, 72 insertions(+), 96 deletions(-) diff --git a/beanie/__init__.py b/beanie/__init__.py index bbfb17b6..95111831 100644 --- 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 index e7440adc..04e6e599 100644 --- 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 index da483c48..e3c96053 100644 --- 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 index 3a2990e5..f32a494d 100644 --- 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 index 8be48e43..e1a42bb1 100644 --- 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 index dad2014c..ab27f5a6 100644 --- 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 + ## [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 \ No newline at end of file +[1.15.0]: https://pypi.org/project/beanie/1.15.0 + +[1.15.1]: https://pypi.org/project/beanie/1.15.1 \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 7e5bfe58..81befe2e 100644 --- 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 "] license = "Apache-2.0" diff --git a/tests/odm/documents/test_bulk_write.py b/tests/odm/documents/test_bulk_write.py index e3e20d66..050365e8 100644 --- 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 index 941710d2..de232786 100644 --- 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"