feat: code snippets for feature store control plane

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
 
@@ -83,6 +83,14 @@ 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()
 def teardown_batch_prediction_job(shared_state, job_client):
@@ -201,6 +209,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,45 @@
+# 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 = aiplatform.CreateEntityTypeRequest(
+        parent=parent,
+        entity_type_id=entity_type_id,
+        entity_type=aiplatform.EntityType(description=description),
+    )
+    lro_response = client.create_entity_type(request=create_entity_type_request)
+    print("Long running operation:", lro_response.operation.name)
+    create_entity_type_response = lro_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: aiplatform.Feature.ValueType,
+    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 = aiplatform.CreateFeatureRequest(
+        parent=parent,
+        feature=aiplatform.Feature(value_type=value_type, description=description),
+        feature_id=feature_id,
+    )
+    lro_response = client.create_feature(request=create_feature_request)
+    print("Long running operation:", lro_response.operation.name)
+    create_feature_response = lro_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
+from google.cloud import aiplatform_v1beta1 as aiplatform
+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]
+    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=aiplatform.Feature.ValueType.INT64,
+    )
+    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,48 @@
+# 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 = aiplatform.CreateFeaturestoreRequest(
+        parent=parent,
+        featurestore_id=featurestore_id,
+        featurestore=aiplatform.Featurestore(
+            online_serving_config=aiplatform.Featurestore.OnlineServingConfig(
+                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]