feat(firestore): add DML support to pipelines

packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py

@@ -610,3 +610,55 @@ def distinct(self, *fields: str | Selectable) -> "_BasePipeline":
             A new Pipeline object with this stage appended to the stage list
         """
         return self._append(stages.Distinct(*fields))
+
+    def delete(self) -> "_BasePipeline":
+        """
+        Deletes the documents from the current pipeline stage.
+
+        Example:
+            >>> from google.cloud.firestore_v1.pipeline_expressions import Field
+            >>> pipeline = client.pipeline().collection("logs")
+            >>> # Delete all documents in the "logs" collection where "status" is "archived"
+            >>> pipeline = pipeline.where(Field.of("status").equal("archived")).delete()
+            >>> pipeline.execute()
+
+        Returns:
+            A new Pipeline object with this stage appended to the stage list
+        """
+        return self._append(stages.Delete())
+
+    def update(self, *transformed_fields: "Selectable") -> "_BasePipeline":
+        """
+        Updates the documents with the given transformations.
+
+        This method updates the documents in place based on the data flowing through the pipeline.
+        To specify transformations, use `*transformed_fields`.
+
+        Example 1: Update a collection's schema by adding a new field and removing an old one.
+            >>> from google.cloud.firestore_v1.pipeline_expressions import Constant
+            >>> pipeline = client.pipeline().collection("books")
+            >>> pipeline = pipeline.add_fields(Constant.of("Fiction").as_("genre"))
+            >>> pipeline = pipeline.remove_fields("old_genre").update()
+            >>> pipeline.execute()
+
+        Example 2: Update documents in place with data from literals.
+            >>> pipeline = client.pipeline().literals(
+            ...     {"__name__": client.collection("books").document("book1"), "status": "Updated"}
+            ... ).update()
+            >>> pipeline.execute()
+
+        Example 3: Update documents from previous stages with specified transformations.
+            >>> from google.cloud.firestore_v1.pipeline_expressions import Field, Constant
+            >>> pipeline = client.pipeline().collection("books")
+            >>> # Update the "status" field to "Discounted" for all books where price > 50
+            >>> pipeline = pipeline.where(Field.of("price").greater_than(50))
+            >>> pipeline = pipeline.update(Constant.of("Discounted").as_("status"))
+            >>> pipeline.execute()
+
+        Args:
+            *transformed_fields: The transformations to apply.
+
+        Returns:
+            A new Pipeline object with this stage appended to the stage list
+        """
+        return self._append(stages.Update(*transformed_fields))

packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_stages.py

@@ -494,3 +494,30 @@ def __init__(self, condition: BooleanExpression):
 
     def _pb_args(self):
         return [self.condition._to_pb()]
+
+
+class Delete(Stage):
+    """Deletes documents matching the pipeline criteria."""
+
+    def __init__(self):
+        super().__init__("delete")
+
+    def _pb_args(self) -> list[Value]:
+        return []
+
+    def _pb_options(self) -> dict[str, Value]:
+        return {}
+
+
+class Update(Stage):
+    """Updates documents with transformed fields."""
+
+    def __init__(self, *transformed_fields: Selectable):
+        super().__init__("update")
+        self.transformed_fields = list(transformed_fields)
+
+    def _pb_args(self) -> list[Value]:
+        return [Selectable._to_value(self.transformed_fields)]
+
+    def _pb_options(self) -> dict[str, Value]:
+        return {}

packages/google-cloud-firestore/tests/system/pipeline_e2e/dml.yaml

@@ -0,0 +1,23 @@
+data:
+  dml_delete_coll:
+    doc1: { score: 10 }
+    doc2: { score: 60 }
+  dml_update_coll:
+    doc1: { status: "pending", score: 50 }
+
+tests:
+  - description: "Basic DML delete"
+    pipeline:
+      - Collection: dml_delete_coll
+      - Where: Field.of("score").less_than(50)
+      - Delete:
+    assert_end_state:
+      dml_delete_coll/doc1: null
+      dml_delete_coll/doc2: { score: 60 }
+
+  - description: "Basic DML update"
+    pipeline:
+      - Collection: dml_update_coll
+      - Update: Field.of("status").set("active")
+    assert_end_state:
+      dml_update_coll/doc1: { status: "active", score: 50 }

packages/google-cloud-firestore/tests/system/test_pipeline_acceptance.py

@@ -206,6 +206,59 @@ async def test_pipeline_results_async(test_dict, async_client):
         assert len(got_results) == expected_count
 
 
+@pytest.mark.parametrize(
+    "test_dict",
+    [t for t in yaml_loader() if "assert_end_state" in t],
+    ids=id_format,
+)
+def test_pipeline_end_state(test_dict, client):
+    """
+    Ensure pipeline leaves database in expected state
+    """
+    expected_state = _parse_yaml_types(test_dict["assert_end_state"])
+    pipeline = parse_pipeline(client, test_dict["pipeline"])
+
+    # Execute the DML pipeline
+    pipeline.execute()
+
+    # Assert end state
+    for doc_path, expected_content in expected_state.items():
+        doc_ref = client.document(doc_path)
+        snapshot = doc_ref.get()
+        if expected_content is None or expected_content == {}:
+            assert not snapshot.exists, f"Expected {doc_path} to be deleted, but it exists"
+        else:
+            assert snapshot.exists, f"Expected {doc_path} to exist, but it was deleted"
+            assert snapshot.to_dict() == expected_content
+
+
+@pytest.mark.parametrize(
+    "test_dict",
+    [t for t in yaml_loader() if "assert_end_state" in t],
+    ids=id_format,
+)
+@pytest.mark.asyncio
+async def test_pipeline_end_state_async(test_dict, async_client):
+    """
+    Ensure pipeline leaves database in expected state
+    """
+    expected_state = _parse_yaml_types(test_dict["assert_end_state"])
+    pipeline = parse_pipeline(async_client, test_dict["pipeline"])
+
+    # Execute the DML pipeline
+    await pipeline.execute()
+
+    # Assert end state
+    for doc_path, expected_content in expected_state.items():
+        doc_ref = async_client.document(doc_path)
+        snapshot = await doc_ref.get()
+        if expected_content is None or expected_content == {}:
+            assert not snapshot.exists, f"Expected {doc_path} to be deleted, but it exists"
+        else:
+            assert snapshot.exists, f"Expected {doc_path} to exist, but it was deleted"
+            assert snapshot.to_dict() == expected_content
+
+
 #################################################################################
 # Helpers & Fixtures
 #################################################################################

packages/google-cloud-firestore/tests/unit/v1/test_pipeline_stages.py

@@ -960,3 +960,40 @@ def test_to_pb(self):
         assert got_fn.args[0].field_reference_value == "city"
         assert got_fn.args[1].string_value == "SF"
         assert len(result.options) == 0
+
+
+class TestDelete:
+    def _make_one(self):
+        return stages.Delete()
+
+    def test_to_pb(self):
+        instance = self._make_one()
+        result = instance._to_pb()
+        assert result.name == "delete"
+        assert len(result.args) == 0
+        assert len(result.options) == 0
+
+
+class TestUpdate:
+    def _make_one(self, *args):
+        return stages.Update(*args)
+
+    def test_to_pb_empty(self):
+        instance = self._make_one()
+        result = instance._to_pb()
+        assert result.name == "update"
+        assert len(result.args) == 1
+        assert result.args[0].map_value.fields == {}
+        assert len(result.options) == 0
+
+    def test_to_pb_with_fields(self):
+        instance = self._make_one(
+            Field.of("score").add(10).as_("score"),
+            Field.of("status").set("active").as_("status")
+        )
+        result = instance._to_pb()
+        assert result.name == "update"
+        assert len(result.args) == 1
+        assert "score" in result.args[0].map_value.fields
+        assert "status" in result.args[0].map_value.fields
+        assert len(result.options) == 0