From cbfef72fbdaff59986e062e10240fbc6ad8ad002 Mon Sep 17 00:00:00 2001 From: Stefano Lottini Date: Wed, 6 Mar 2024 02:39:11 +0100 Subject: [PATCH] update_many --- astrapy/db.py | 24 ++++++++-- astrapy/idiomatic/collection.py | 44 +++++++++++++++++++ tests/idiomatic/integration/test_dml_async.py | 33 ++++++++++++++ tests/idiomatic/integration/test_dml_sync.py | 33 ++++++++++++++ 4 files changed, 130 insertions(+), 4 deletions(-) diff --git a/astrapy/db.py b/astrapy/db.py index fd0f91aa..38622cce 100644 --- a/astrapy/db.py +++ b/astrapy/db.py @@ -879,7 +879,10 @@ def update_one( return response def update_many( - self, filter: Dict[str, Any], update: Dict[str, Any] + self, + filter: Dict[str, Any], + update: Dict[str, Any], + options: Optional[Dict[str, Any]] = None, ) -> API_RESPONSE: """ Updates multiple documents in the collection. @@ -889,7 +892,12 @@ def update_many( Returns: dict: The response from the database after the update operation. """ - json_query = make_payload(top_level="updateMany", filter=filter, update=update) + json_query = make_payload( + top_level="updateMany", + filter=filter, + update=update, + options=options, + ) response = self._request( method=http_methods.POST, @@ -1910,7 +1918,10 @@ async def update_one( return response async def update_many( - self, filter: Dict[str, Any], update: Dict[str, Any] + self, + filter: Dict[str, Any], + update: Dict[str, Any], + options: Optional[Dict[str, Any]] = None, ) -> API_RESPONSE: """ Updates multiple documents in the collection. @@ -1920,7 +1931,12 @@ async def update_many( Returns: dict: The response from the database after the update operation. """ - json_query = make_payload(top_level="updateMany", filter=filter, update=update) + json_query = make_payload( + top_level="updateMany", + filter=filter, + update=update, + options=options, + ) response = await self._request( method=http_methods.POST, diff --git a/astrapy/idiomatic/collection.py b/astrapy/idiomatic/collection.py index f98dbe7f..8ec33209 100644 --- a/astrapy/idiomatic/collection.py +++ b/astrapy/idiomatic/collection.py @@ -385,6 +385,28 @@ def update_one( f"(gotten '${json.dumps(fo_response)}')" ) + def update_many( + self, + filter: Dict[str, Any], + update: Dict[str, Any], + *, + upsert: bool = False, + ) -> UpdateResult: + options = { + "upsert": upsert, + } + um_response = self._astra_db_collection.update_many( + update=update, + filter=filter, + options=options, + ) + um_status = um_response.get("status") or {} + _update_info = _prepare_update_info(um_status) + return UpdateResult( + raw_result=um_status, + update_info=_update_info, + ) + def find_one_and_delete( self, filter: Dict[str, Any], @@ -800,6 +822,28 @@ async def update_one( f"(gotten '${json.dumps(fo_response)}')" ) + async def update_many( + self, + filter: Dict[str, Any], + update: Dict[str, Any], + *, + upsert: bool = False, + ) -> UpdateResult: + options = { + "upsert": upsert, + } + um_response = await self._astra_db_collection.update_many( + update=update, + filter=filter, + options=options, + ) + um_status = um_response.get("status") or {} + _update_info = _prepare_update_info(um_status) + return UpdateResult( + raw_result=um_status, + update_info=_update_info, + ) + async def find_one_and_delete( self, filter: Dict[str, Any], diff --git a/tests/idiomatic/integration/test_dml_async.py b/tests/idiomatic/integration/test_dml_async.py index 260d32c2..72556896 100644 --- a/tests/idiomatic/integration/test_dml_async.py +++ b/tests/idiomatic/integration/test_dml_async.py @@ -717,6 +717,39 @@ async def test_collection_update_one_async( assert result4.update_info["nModified"] == 1 assert "upserted" not in result4.update_info + @pytest.mark.describe("test of update_many, async") + async def test_collection_update_many_async( + self, + async_empty_collection: AsyncCollection, + ) -> None: + acol = async_empty_collection + await acol.insert_many([{"a": 1, "seq": i} for i in range(4)]) + await acol.insert_many([{"a": 2, "seq": i} for i in range(2)]) + + resp1 = await acol.update_many({"a": 1}, {"$set": {"n": 1}}) + assert resp1.update_info["n"] == 4 + assert resp1.update_info["updatedExisting"] is True + assert resp1.update_info["nModified"] == 4 + assert "upserted" not in resp1.update_info + + resp2 = await acol.update_many({"a": 1}, {"$set": {"n": 2}}, upsert=True) + assert resp2.update_info["n"] == 4 + assert resp2.update_info["updatedExisting"] is True + assert resp2.update_info["nModified"] == 4 + assert "upserted" not in resp2.update_info + + resp3 = await acol.update_many({"a": 3}, {"$set": {"n": 3}}) + assert resp3.update_info["n"] == 0 + assert resp3.update_info["updatedExisting"] is False + assert resp3.update_info["nModified"] == 0 + assert "upserted" not in resp3.update_info + + resp4 = await acol.update_many({"a": 3}, {"$set": {"n": 4}}, upsert=True) + assert resp4.update_info["n"] == 1 + assert resp4.update_info["updatedExisting"] is False + assert resp4.update_info["nModified"] == 0 + assert "upserted" in resp4.update_info + @pytest.mark.describe("test of collection find_one_and_delete, async") async def test_collection_find_one_and_delete_async( self, diff --git a/tests/idiomatic/integration/test_dml_sync.py b/tests/idiomatic/integration/test_dml_sync.py index b2d8d04d..fb7359e9 100644 --- a/tests/idiomatic/integration/test_dml_sync.py +++ b/tests/idiomatic/integration/test_dml_sync.py @@ -699,6 +699,39 @@ def test_collection_update_one_sync( assert result4.update_info["nModified"] == 1 assert "upserted" not in result4.update_info + @pytest.mark.describe("test of update_many, sync") + def test_collection_update_many_sync( + self, + sync_empty_collection: Collection, + ) -> None: + col = sync_empty_collection + col.insert_many([{"a": 1, "seq": i} for i in range(4)]) + col.insert_many([{"a": 2, "seq": i} for i in range(2)]) + + resp1 = col.update_many({"a": 1}, {"$set": {"n": 1}}) + assert resp1.update_info["n"] == 4 + assert resp1.update_info["updatedExisting"] is True + assert resp1.update_info["nModified"] == 4 + assert "upserted" not in resp1.update_info + + resp2 = col.update_many({"a": 1}, {"$set": {"n": 2}}, upsert=True) + assert resp2.update_info["n"] == 4 + assert resp2.update_info["updatedExisting"] is True + assert resp2.update_info["nModified"] == 4 + assert "upserted" not in resp2.update_info + + resp3 = col.update_many({"a": 3}, {"$set": {"n": 3}}) + assert resp3.update_info["n"] == 0 + assert resp3.update_info["updatedExisting"] is False + assert resp3.update_info["nModified"] == 0 + assert "upserted" not in resp3.update_info + + resp4 = col.update_many({"a": 3}, {"$set": {"n": 4}}, upsert=True) + assert resp4.update_info["n"] == 1 + assert resp4.update_info["updatedExisting"] is False + assert resp4.update_info["nModified"] == 0 + assert "upserted" in resp4.update_info + @pytest.mark.describe("test of collection find_one_and_delete, sync") def test_collection_find_one_and_delete_sync( self,