update_many

datastax · Mar 6, 2024 · cbfef72 · cbfef72
1 parent b268c3f
commit cbfef72
Show file tree

Hide file tree

Showing 4 changed files with 130 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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,