Skip to content

Commit

Permalink
feat: add FeatureGroup delete
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 633836995
  • Loading branch information
vertex-sdk-bot authored and copybara-github committed May 15, 2024
1 parent 710f33d commit f9011e0
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 0 deletions.
41 changes: 41 additions & 0 deletions tests/unit/vertexai/test_feature_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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."
),
]
)
30 changes: 30 additions & 0 deletions vertexai/resources/preview/feature_store/feature_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit f9011e0

Please sign in to comment.