feat: Add code snippets for feature store service

Add resource creation and deletion including featurestore, entity
type and feature.

samples/snippets/conftest.py

@@ -15,7 +15,7 @@
 import os
 from uuid import uuid4
 
-from google.cloud import aiplatform
+from google.cloud import aiplatform, aiplatform_v1beta1
 from google.cloud import storage
 import pytest
 
@@ -82,6 +82,12 @@ def dataset_client():
     )
     yield dataset_client
 
+@pytest.fixture
+def featurestore_client():
+    featurestore_client = aiplatform_v1beta1.FeaturestoreServiceClient(
+        client_options={"api_endpoint": "us-central1-aiplatform.googleapis.com"}
+    )
+    yield featurestore_client
 
 # Shared setup/teardown.
 @pytest.fixture()
@@ -201,6 +207,30 @@ def teardown_dataset(shared_state, dataset_client):
     dataset_client.delete_dataset(name=shared_state["dataset_name"])
 
 
+@pytest.fixture()
+def teardown_featurestore(shared_state, featurestore_client):
+    yield
+
+    # Delete the created featurestore
+    featurestore_client.delete_featurestore(name=shared_state["featurestore_name"])
+
+
+@pytest.fixture()
+def teardown_entity_type(shared_state, featurestore_client):
+    yield
+
+    # Delete the created entity type
+    featurestore_client.delete_entity_type(name=shared_state["entity_type_name"])
+
+
+@pytest.fixture()
+def teardown_feature(shared_state, featurestore_client):
+    yield
+
+    # Delete the created feature
+    featurestore_client.delete_feature(name=shared_state["feature_name"])
+
+
 @pytest.fixture()
 def create_endpoint(shared_state, endpoint_client):
     def create(project, location, test_name="temp_deploy_model_test"):

samples/snippets/feature_store_service/create_entity_type_sample.py

@@ -0,0 +1,46 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# [START aiplatform_create_entity_type_sample]
+from google.cloud import aiplatform_v1beta1 as aiplatform
+
+
+def create_entity_type_sample(
+    project: str,
+    featurestore_id: str,
+    entity_type_id: str,
+    description: str = "sample entity type",
+    location: str = "us-central1",
+    api_endpoint: str = "us-central1-aiplatform.googleapis.com",
+    timeout: int = 300,
+):
+    # The AI Platform services require regional API endpoints.
+    client_options = {"api_endpoint": api_endpoint}
+    # Initialize client that will be used to create and send requests.
+    # This client only needs to be created once, and can be reused for multiple requests.
+    client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
+    parent = f"projects/{project}/locations/{location}/featurestores/{featurestore_id}"
+    create_entity_type_request = {
+        "parent": parent,
+        "entity_type_id": entity_type_id,
+        "entity_type": {"description": description},
+    }
+
+    response = client.create_entity_type(request=create_entity_type_request)
+    print("Long running operation:", response.operation.name)
+    create_entity_type_response = response.result(timeout=timeout)
+    print("create_entity_type_response:", create_entity_type_response)
+
+
+# [END aiplatform_create_entity_type_sample]

samples/snippets/feature_store_service/create_entity_type_sample_test.py

@@ -0,0 +1,42 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from uuid import uuid4
+
+import create_entity_type_sample
+import pytest
+
+import helpers
+
+PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT")
+
+
+@pytest.fixture(scope="function", autouse=True)
+def teardown(teardown_entity_type):
+    yield
+
+
+def test_ucaip_generated_create_entity_type_sample_vision(capsys, shared_state):
+    featurestore_id = "perm_sample_featurestore"
+    entity_type_id = f"temp_create_entity_type_test_{uuid4()}".replace("-", "_")[:60]
+    create_entity_type_sample.create_entity_type_sample(
+        project=PROJECT_ID,
+        featurestore_id=featurestore_id,
+        entity_type_id=entity_type_id,
+    )
+    out, _ = capsys.readouterr()
+    assert "create_entity_type_response" in out
+
+    shared_state["entity_type_name"] = helpers.get_featurestore_resource_name(out)

samples/snippets/feature_store_service/create_feature_sample.py

@@ -0,0 +1,47 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# [START aiplatform_create_feature_sample]
+from google.cloud import aiplatform_v1beta1 as aiplatform
+
+
+def create_feature_sample(
+    project: str,
+    featurestore_id: str,
+    entity_type_id: str,
+    feature_id: str,
+    value_type: str,
+    description: str = "sample feature",
+    location: str = "us-central1",
+    api_endpoint: str = "us-central1-aiplatform.googleapis.com",
+    timeout: int = 300,
+):
+    # The AI Platform services require regional API endpoints.
+    client_options = {"api_endpoint": api_endpoint}
+    # Initialize client that will be used to create and send requests.
+    # This client only needs to be created once, and can be reused for multiple requests.
+    client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
+    parent = f"projects/{project}/locations/{location}/featurestores/{featurestore_id}/entityTypes/{entity_type_id}"
+    create_feature_request = {
+        "parent": parent,
+        "feature": {"value_type": value_type, "description": description,},
+        "feature_id": feature_id,
+    }
+    response = client.create_feature(request=create_feature_request)
+    print("Long running operation:", response.operation.name)
+    create_feature_response = response.result(timeout=timeout)
+    print("create_feature_response:", create_feature_response)
+
+
+# [END aiplatform_create_feature_sample]

samples/snippets/feature_store_service/create_feature_sample_test.py

@@ -0,0 +1,46 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from uuid import uuid4
+
+import create_feature_sample
+import pytest
+
+import helpers
+
+PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT")
+
+
+@pytest.fixture(scope="function", autouse=True)
+def teardown(teardown_feature):
+    yield
+
+
+def test_ucaip_generated_create_feature_sample_vision(capsys, shared_state):
+    featurestore_id = "perm_sample_featurestore"
+    entity_type_id = "perm_sample_entity_type"
+    feature_id = f"temp_create_feature_test_{uuid4()}".replace("-", "_")[:60]
+    int_64_type = "INT64"
+    create_feature_sample.create_feature_sample(
+        project=PROJECT_ID,
+        featurestore_id=featurestore_id,
+        entity_type_id=entity_type_id,
+        feature_id=feature_id,
+        value_type=int_64_type,
+    )
+    out, _ = capsys.readouterr()
+    assert "create_feature_response" in out
+
+    shared_state["feature_name"] = helpers.get_featurestore_resource_name(out)

samples/snippets/feature_store_service/create_featurestore_sample.py

@@ -0,0 +1,46 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# [START aiplatform_create_featurestore_sample]
+from google.cloud import aiplatform_v1beta1 as aiplatform
+
+
+def create_featurestore_sample(
+    project: str,
+    featurestore_id: str,
+    fixed_node_count: int = 1,
+    location: str = "us-central1",
+    api_endpoint: str = "us-central1-aiplatform.googleapis.com",
+    timeout: int = 300,
+):
+    # The AI Platform services require regional API endpoints.
+    client_options = {"api_endpoint": api_endpoint}
+    # Initialize client that will be used to create and send requests.
+    # This client only needs to be created once, and can be reused for multiple requests.
+    client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
+    parent = f"projects/{project}/locations/{location}"
+    create_featurestore_request = {
+        "parent": parent,
+        "featurestore_id": featurestore_id,
+        "featurestore": {
+            "online_serving_config": {"fixed_node_count": fixed_node_count}
+        },
+    }
+    lro_response = client.create_featurestore(request=create_featurestore_request)
+    print("Long running operation:", lro_response.operation.name)
+    create_featurestore_response = lro_response.result(timeout=timeout)
+    print("create_featurestore_response:", create_featurestore_response)
+
+
+# [END aiplatform_create_featurestore_sample]

samples/snippets/feature_store_service/create_featurestore_sample_test.py

@@ -0,0 +1,39 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from uuid import uuid4
+
+import create_featurestore_sample
+import pytest
+
+import helpers
+
+PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT")
+
+
+@pytest.fixture(scope="function", autouse=True)
+def teardown(teardown_featurestore):
+    yield
+
+
+def test_ucaip_generated_create_featurestore_sample_vision(capsys, shared_state):
+    featurestore_id = f"temp_create_featurestore_test_{uuid4()}".replace("-", "_")[:60]
+    create_featurestore_sample.create_featurestore_sample(
+        project=PROJECT_ID, featurestore_id=featurestore_id, fixed_node_count=1
+    )
+    out, _ = capsys.readouterr()
+    assert "create_featurestore_response" in out
+
+    shared_state["featurestore_name"] = helpers.get_featurestore_resource_name(out)

samples/snippets/feature_store_service/delete_featurestore_sample.py

@@ -0,0 +1,40 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# [START aiplatform_delete_featurestore_sample]
+from google.cloud import aiplatform_v1beta1 as aiplatform
+
+
+def delete_featurestore_sample(
+    project: str,
+    featurestore_id: str,
+    location: str = "us-central1",
+    api_endpoint: str = "us-central1-aiplatform.googleapis.com",
+    timeout: int = 300,
+):
+    # The AI Platform services require regional API endpoints.
+    client_options = {"api_endpoint": api_endpoint}
+    # Initialize client that will be used to create and send requests.
+    # This client only needs to be created once, and can be reused for multiple requests.
+    client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
+    name = client.featurestore_path(
+        project=project, location=location, featurestore=featurestore_id
+    )
+    response = client.delete_featurestore(name=name)
+    print("Long running operation:", response.operation.name)
+    delete_featurestore_response = response.result(timeout=timeout)
+    print("delete_featurestore_response:", delete_featurestore_response)
+
+
+# [END aiplatform_delete_featurestore_sample]

samples/snippets/helpers.py

@@ -18,6 +18,13 @@ def get_state(out):
     return state
 
 
+def get_featurestore_resource_name(out, key="name"):
+    pattern = re.compile(fr'{key}:\s*"([\_\-a-zA-Z0-9/]+)"')
+    name = re.search(pattern, out).group(1)
+
+    return name
+
+
 def wait_for_job_state(
     get_job_method: Callable[[str], "proto.Message"],  # noqa: F821
     name: str,