Skip to content

Commit

Permalink
Allow unordered parameter on BulkWriter
Browse files Browse the repository at this point in the history
  • Loading branch information
thiagosalvatore committed Jun 10, 2024
1 parent 51e73eb commit 0c8f827
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 2 deletions.
7 changes: 5 additions & 2 deletions beanie/odm/bulk.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,12 @@ class Config:


class BulkWriter:
def __init__(self, session: Optional[ClientSession] = None):
def __init__(
self, session: Optional[ClientSession] = None, ordered: bool = True
):
self.operations: List[Operation] = []
self.session = session
self.ordered = ordered

async def __aenter__(self):
return self
Expand Down Expand Up @@ -78,7 +81,7 @@ async def commit(self) -> Optional[BulkWriteResult]:
requests.append(query)

return await obj_class.get_motor_collection().bulk_write( # type: ignore
requests, session=self.session
requests, session=self.session, ordered=self.ordered
)
return None

Expand Down
22 changes: 22 additions & 0 deletions tests/odm/documents/test_bulk_write.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,28 @@ async def test_update(documents, document_not_inserted):
)


async def test_unordered_update(documents, document):
await documents(5)
doc = await DocumentTestModel.find_one(DocumentTestModel.test_int == 0)
doc.test_int = 100
with pytest.raises(BulkWriteError):
async with BulkWriter(ordered=False) as bulk_writer:
await DocumentTestModel.insert_one(
document, bulk_writer=bulk_writer
)
await doc.save_changes(bulk_writer=bulk_writer)

assert len(await DocumentTestModel.find_all().to_list()) == 6
assert (
len(
await DocumentTestModel.find(
DocumentTestModel.test_int == 100
).to_list()
)
== 1
)


async def test_delete(documents, document_not_inserted):
await documents(5)
doc = await DocumentTestModel.find_one(DocumentTestModel.test_int == 0)
Expand Down

0 comments on commit 0c8f827

Please sign in to comment.