feat(samples): add create-featurestore sample (#948)

* CreateFeaturestore Sample API * DeleteFeaturestore Sample API * Create, Update, Get featurestore API's * Revert Code * CreateFeaturestoreSampleTest update * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
GoogleCloudPlatform · Jun 21, 2022 · 3a4771e · 3a4771e
1 parent 68dda72
commit 3a4771e
Show file tree

Hide file tree

Showing 3 changed files with 329 additions and 0 deletions.
diff --git a/aiplatform/snippets/src/main/java/aiplatform/CreateFeaturestoreSample.java b/aiplatform/snippets/src/main/java/aiplatform/CreateFeaturestoreSample.java
@@ -0,0 +1,101 @@
+/*
+ * Copyright 2022 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.
+ *
+ *
+ * Create a featurestore resource to contain entity types and features. See
+ * https://cloud.google.com/vertex-ai/docs/featurestore/setup before running
+ * the code snippet
+ */
+
+package aiplatform;
+
+// [START aiplatform_create_featurestore_sample]
+
+import com.google.api.gax.longrunning.OperationFuture;
+import com.google.cloud.aiplatform.v1beta1.CreateFeaturestoreOperationMetadata;
+import com.google.cloud.aiplatform.v1beta1.CreateFeaturestoreRequest;
+import com.google.cloud.aiplatform.v1beta1.Featurestore;
+import com.google.cloud.aiplatform.v1beta1.Featurestore.OnlineServingConfig;
+import com.google.cloud.aiplatform.v1beta1.Featurestore.OnlineServingConfig.Scaling;
+import com.google.cloud.aiplatform.v1beta1.FeaturestoreServiceClient;
+import com.google.cloud.aiplatform.v1beta1.FeaturestoreServiceSettings;
+import com.google.cloud.aiplatform.v1beta1.LocationName;
+import java.io.IOException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+public class CreateFeaturestoreSample {
+
+  public static void main(String[] args)
+      throws IOException, InterruptedException, ExecutionException, TimeoutException {
+    // TODO(developer): Replace these variables before running the sample.
+    String project = "YOUR_PROJECT_ID";
+    String featurestoreId = "YOUR_FEATURESTORE_ID";
+    int minNodeCount = 1;
+    int maxNodeCount = 5;
+    String location = "us-central1";
+    String endpoint = "us-central1-aiplatform.googleapis.com:443";
+    int timeout = 900;
+    createFeaturestoreSample(
+        project, featurestoreId, minNodeCount, maxNodeCount, location, endpoint, timeout);
+  }
+
+  static void createFeaturestoreSample(
+      String project,
+      String featurestoreId,
+      int minNodeCount,
+      int maxNodeCount,
+      String location,
+      String endpoint,
+      int timeout)
+      throws IOException, InterruptedException, ExecutionException, TimeoutException {
+
+    FeaturestoreServiceSettings featurestoreServiceSettings =
+        FeaturestoreServiceSettings.newBuilder().setEndpoint(endpoint).build();
+
+    // Initialize client that will be used to send requests. This client only needs to be created
+    // once, and can be reused for multiple requests. After completing all of your requests, call
+    // the "close" method on the client to safely clean up any remaining background resources.
+    try (FeaturestoreServiceClient featurestoreServiceClient =
+        FeaturestoreServiceClient.create(featurestoreServiceSettings)) {
+
+      OnlineServingConfig.Builder builderValue =
+          OnlineServingConfig.newBuilder()
+              .setScaling(
+                  Scaling.newBuilder().setMinNodeCount(minNodeCount).setMaxNodeCount(maxNodeCount));
+      Featurestore featurestore =
+          Featurestore.newBuilder().setOnlineServingConfig(builderValue).build();
+      String parent = LocationName.of(project, location).toString();
+
+      CreateFeaturestoreRequest createFeaturestoreRequest =
+          CreateFeaturestoreRequest.newBuilder()
+              .setParent(parent)
+              .setFeaturestore(featurestore)
+              .setFeaturestoreId(featurestoreId)
+              .build();
+
+      OperationFuture<Featurestore, CreateFeaturestoreOperationMetadata> featurestoreFuture =
+          featurestoreServiceClient.createFeaturestoreAsync(createFeaturestoreRequest);
+      System.out.format(
+          "Operation name: %s%n", featurestoreFuture.getInitialFuture().get().getName());
+      System.out.println("Waiting for operation to finish...");
+      Featurestore featurestoreResponse = featurestoreFuture.get(timeout, TimeUnit.SECONDS);
+      System.out.println("Create Featurestore Response");
+      System.out.format("Name: %s%n", featurestoreResponse.getName());
+    }
+  }
+}
+// [END aiplatform_create_featurestore_sample]
diff --git a/aiplatform/snippets/src/main/java/aiplatform/DeleteFeaturestoreSample.java b/aiplatform/snippets/src/main/java/aiplatform/DeleteFeaturestoreSample.java
@@ -0,0 +1,86 @@
+/*
+ * Copyright 2022 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.
+ *
+ *
+ * Delete a featurestore. See
+ * https://cloud.google.com/vertex-ai/docs/featurestore/setup before running
+ * the code snippet
+ */
+
+package aiplatform;
+
+// [START aiplatform_delete_featurestore_sample]
+
+import com.google.api.gax.longrunning.OperationFuture;
+import com.google.cloud.aiplatform.v1.DeleteFeaturestoreRequest;
+import com.google.cloud.aiplatform.v1.DeleteOperationMetadata;
+import com.google.cloud.aiplatform.v1.FeaturestoreName;
+import com.google.cloud.aiplatform.v1.FeaturestoreServiceClient;
+import com.google.cloud.aiplatform.v1.FeaturestoreServiceSettings;
+import com.google.protobuf.Empty;
+import java.io.IOException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+public class DeleteFeaturestoreSample {
+
+  public static void main(String[] args)
+      throws IOException, InterruptedException, ExecutionException, TimeoutException {
+    // TODO(developer): Replace these variables before running the sample.
+    String project = "YOUR_PROJECT_ID";
+    String featurestoreId = "YOUR_FEATURESTORE_ID";
+    boolean useForce = true;
+    String location = "us-central1";
+    String endpoint = "us-central1-aiplatform.googleapis.com:443";
+    int timeout = 60;
+    deleteFeaturestoreSample(project, featurestoreId, useForce, location, endpoint, timeout);
+  }
+
+  static void deleteFeaturestoreSample(
+      String project,
+      String featurestoreId,
+      boolean useForce,
+      String location,
+      String endpoint,
+      int timeout)
+      throws IOException, InterruptedException, ExecutionException, TimeoutException {
+
+    FeaturestoreServiceSettings featurestoreServiceSettings =
+        FeaturestoreServiceSettings.newBuilder().setEndpoint(endpoint).build();
+
+    // Initialize client that will be used to send requests. This client only needs to be created
+    // once, and can be reused for multiple requests. After completing all of your requests, call
+    // the "close" method on the client to safely clean up any remaining background resources.
+    try (FeaturestoreServiceClient featurestoreServiceClient =
+        FeaturestoreServiceClient.create(featurestoreServiceSettings)) {
+
+      DeleteFeaturestoreRequest deleteFeaturestoreRequest =
+          DeleteFeaturestoreRequest.newBuilder()
+              .setName(FeaturestoreName.of(project, location, featurestoreId).toString())
+              .setForce(useForce)
+              .build();
+
+      OperationFuture<Empty, DeleteOperationMetadata> operationFuture =
+          featurestoreServiceClient.deleteFeaturestoreAsync(deleteFeaturestoreRequest);
+      System.out.format("Operation name: %s%n", operationFuture.getInitialFuture().get().getName());
+      System.out.println("Waiting for operation to finish...");
+      operationFuture.get(timeout, TimeUnit.SECONDS);
+
+      System.out.format("Deleted Featurestore.");
+    }
+  }
+}
+// [END aiplatform_delete_featurestore_sample]
diff --git a/aiplatform/snippets/src/test/java/aiplatform/CreateFeaturestoreSampleTest.java b/aiplatform/snippets/src/test/java/aiplatform/CreateFeaturestoreSampleTest.java
@@ -0,0 +1,142 @@
+/*
+ * Copyright 2022 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.
+ */
+
+package aiplatform;
+
+import static com.google.common.truth.Truth.assertThat;
+import static junit.framework.TestCase.assertNotNull;
+
+import com.google.api.gax.longrunning.OperationFuture;
+import com.google.cloud.aiplatform.v1.DeleteFeaturestoreRequest;
+import com.google.cloud.aiplatform.v1.DeleteOperationMetadata;
+import com.google.cloud.aiplatform.v1.FeaturestoreName;
+import com.google.cloud.aiplatform.v1.FeaturestoreServiceClient;
+import com.google.cloud.aiplatform.v1.FeaturestoreServiceSettings;
+import com.google.protobuf.Empty;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class CreateFeaturestoreSampleTest {
+
+  private static final String PROJECT_ID = System.getenv("UCAIP_PROJECT_ID");
+  private static final int MIN_NODE_COUNT = 1;
+  private static final int MAX_NODE_COUNT = 5;
+  private static final boolean USE_FORCE = true;
+  private static final String LOCATION = "us-central1";
+  private static final String ENDPOINT = "us-central1-aiplatform.googleapis.com:443";
+  private static final int TIMEOUT = 900;
+  private ByteArrayOutputStream bout;
+  private PrintStream out;
+  private PrintStream originalPrintStream;
+  private String featurestoreId;
+
+  private static void requireEnvVar(String varName) {
+    String errorMessage =
+        String.format("Environment variable '%s' is required to perform these tests.", varName);
+    assertNotNull(errorMessage, System.getenv(varName));
+  }
+
+  @BeforeClass
+  public static void checkRequirements() {
+    requireEnvVar("GOOGLE_APPLICATION_CREDENTIALS");
+    requireEnvVar("UCAIP_PROJECT_ID");
+  }
+
+  @Before
+  public void setUp() {
+    bout = new ByteArrayOutputStream();
+    out = new PrintStream(bout);
+    originalPrintStream = System.out;
+    System.setOut(out);
+  }
+
+  static void deleteFeaturestoreSample(
+      String project,
+      String featurestoreId,
+      boolean useForce,
+      String location,
+      String endpoint,
+      int timeout)
+      throws IOException, InterruptedException, ExecutionException, TimeoutException {
+
+    FeaturestoreServiceSettings featurestoreServiceSettings =
+        FeaturestoreServiceSettings.newBuilder().setEndpoint(endpoint).build();
+
+    // Initialize client that will be used to send requests. This client only needs to be created
+    // once, and can be reused for multiple requests. After completing all of your requests, call
+    // the "close" method on the client to safely clean up any remaining background resources.
+    try (FeaturestoreServiceClient featurestoreServiceClient =
+        FeaturestoreServiceClient.create(featurestoreServiceSettings)) {
+
+      DeleteFeaturestoreRequest deleteFeaturestoreRequest =
+          DeleteFeaturestoreRequest.newBuilder()
+              .setName(FeaturestoreName.of(project, location, featurestoreId).toString())
+              .setForce(useForce)
+              .build();
+
+      OperationFuture<Empty, DeleteOperationMetadata> operationFuture =
+          featurestoreServiceClient.deleteFeaturestoreAsync(deleteFeaturestoreRequest);
+      System.out.format("Operation name: %s%n", operationFuture.getInitialFuture().get().getName());
+      System.out.println("Waiting for operation to finish...");
+      operationFuture.get(timeout, TimeUnit.SECONDS);
+
+      System.out.format("Deleted Featurestore.");
+    }
+  }
+
+  @After
+  public void tearDown()
+      throws InterruptedException, ExecutionException, IOException, TimeoutException {
+
+    // Delete the featurestore
+    deleteFeaturestoreSample(PROJECT_ID, featurestoreId, USE_FORCE, LOCATION, ENDPOINT, 60);
+
+    // Assert
+    String deleteFeaturestoreResponse = bout.toString();
+    assertThat(deleteFeaturestoreResponse).contains("Deleted Featurestore");
+    System.out.flush();
+    System.setOut(originalPrintStream);
+  }
+
+  @Test
+  public void testCreateFeaturestoreSample()
+      throws IOException, InterruptedException, ExecutionException, TimeoutException {
+    // Create the featurestore
+    String tempUuid = UUID.randomUUID().toString().replaceAll("-", "_").substring(0, 26);
+    String id = String.format("temp_create_featurestore_test_%s", tempUuid);
+    CreateFeaturestoreSample.createFeaturestoreSample(
+        PROJECT_ID, id, MIN_NODE_COUNT, MAX_NODE_COUNT, LOCATION, ENDPOINT, TIMEOUT);
+
+    // Assert
+    String createFeaturestoreResponse = bout.toString();
+    assertThat(createFeaturestoreResponse).contains("Create Featurestore Response");
+    featurestoreId =
+        createFeaturestoreResponse.split("Name: ")[1].split("featurestores/")[1].split("\n")[0]
+            .trim();
+  }
+}