From f9011e0b1b4ea8470849ecdd5ba9e086c73b778b Mon Sep 17 00:00:00 2001
From: A Vertex SDK engineer <vertex-sdk-bot@google.com>
Date: Wed, 15 May 2024 00:19:15 -0700
Subject: [PATCH] feat: add FeatureGroup delete

PiperOrigin-RevId: 633836995
---
 tests/unit/vertexai/test_feature_group.py     | 41 +++++++++++++++++++
 .../preview/feature_store/feature_group.py    | 30 ++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/tests/unit/vertexai/test_feature_group.py b/tests/unit/vertexai/test_feature_group.py
index 477d41e0d3..a5dbf06f7f 100644
--- a/tests/unit/vertexai/test_feature_group.py
+++ b/tests/unit/vertexai/test_feature_group.py
@@ -108,6 +108,17 @@ def list_fg_mock():
         yield list_fg_mock
 
 
+@pytest.fixture
+def delete_fg_mock():
+    with patch.object(
+        feature_registry_service_client.FeatureRegistryServiceClient,
+        "delete_feature_group",
+    ) as delete_fg_mock:
+        delete_fg_lro_mock = mock.Mock(ga_operation.Operation)
+        delete_fg_mock.return_value = delete_fg_lro_mock
+        yield delete_fg_mock
+
+
 def fg_eq(
     fg_to_check: FeatureGroup,
     name: str,
@@ -293,3 +304,33 @@ def test_list(list_fg_mock):
         location=_TEST_LOCATION,
         labels=_TEST_FG3_LABELS,
     )
+
+
+@pytest.mark.parametrize("force", [True, False])
+def test_delete(force, delete_fg_mock, get_fg_mock, fg_logger_mock, sync=True):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    fg = FeatureGroup(_TEST_FG1_ID)
+    fg.delete(force=force, sync=sync)
+
+    if not sync:
+        fg.wait()
+
+    delete_fg_mock.assert_called_once_with(
+        name=_TEST_FG1_PATH,
+        force=force,
+    )
+
+    fg_logger_mock.assert_has_calls(
+        [
+            call(
+                "Deleting FeatureGroup resource: projects/test-project/locations/us-central1/featureGroups/my_fg1"
+            ),
+            call(
+                f"Delete FeatureGroup backing LRO: {delete_fg_mock.return_value.operation.name}"
+            ),
+            call(
+                "FeatureGroup resource projects/test-project/locations/us-central1/featureGroups/my_fg1 deleted."
+            ),
+        ]
+    )
diff --git a/vertexai/resources/preview/feature_store/feature_group.py b/vertexai/resources/preview/feature_store/feature_group.py
index 90198a4e01..857e73a9a3 100644
--- a/vertexai/resources/preview/feature_store/feature_group.py
+++ b/vertexai/resources/preview/feature_store/feature_group.py
@@ -208,6 +208,36 @@ def create(
 
         return feature_group_obj
 
+    @base.optional_sync()
+    def delete(self, force: bool = False, sync: bool = True) -> None:
+        """Deletes this feature group.
+
+        WARNING: This deletion is permanent.
+
+        Args:
+            force:
+                If set to True, all features under this online store will be
+                deleted prior to online store deletion. Otherwise, deletion
+                will only succeed if the online store has no FeatureViews.
+
+                If set to true, any Features under this FeatureGroup will also
+                be deleted. (Otherwise, the request will only work if the
+                FeatureGroup has no Features.)
+            sync:
+                Whether to execute this deletion synchronously. If False, this
+                method will be executed in concurrent Future and any downstream
+                object will be immediately returned and synced when the Future
+                has completed.
+        """
+
+        lro = getattr(self.api_client, self._delete_method)(
+            name=self.resource_name,
+            force=force,
+        )
+        _LOGGER.log_delete_with_lro(self, lro)
+        lro.result()
+        _LOGGER.log_delete_complete(self)
+
     @property
     def source(self) -> FeatureGroupBigQuerySource:
         return FeatureGroupBigQuerySource(