From 769fec1f90dc61896af21f788429de6489e7d08d Mon Sep 17 00:00:00 2001
From: Eric Dong <itseric@google.com>
Date: Wed, 20 Dec 2023 10:40:28 -0500
Subject: [PATCH] feat: update/add new samples for metadata store (#559)

* feat: update/add new samples for metadata store

* fix: address review comments

* address review comments 2
---
 vertex_ai/metadata_store/main.tf      | 12 +++---
 vertex_ai/metadata_store_cmek/main.tf | 54 +++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 6 deletions(-)
 create mode 100644 vertex_ai/metadata_store_cmek/main.tf

diff --git a/vertex_ai/metadata_store/main.tf b/vertex_ai/metadata_store/main.tf
index 1699ed15e..7b34087d2 100644
--- a/vertex_ai/metadata_store/main.tf
+++ b/vertex_ai/metadata_store/main.tf
@@ -16,14 +16,14 @@
 
 
 # [START aiplatform_create_metadata_store_sample]
-resource "random_id" "store_prefix" {
-  byte_length = 8
-}
-
-resource "google_vertex_ai_metadata_store" "main" {
-  name        = "${random_id.store_prefix.hex}-test-store"
+resource "google_vertex_ai_metadata_store" "default" {
+  name        = "${random_id.default.hex}-example-store"
   description = "Example metadata store"
   provider    = google-beta
   region      = "us-central1"
 }
+
+resource "random_id" "default" {
+  byte_length = 8
+}
 # [END aiplatform_create_metadata_store_sample]
diff --git a/vertex_ai/metadata_store_cmek/main.tf b/vertex_ai/metadata_store_cmek/main.tf
new file mode 100644
index 000000000..4f221d1f8
--- /dev/null
+++ b/vertex_ai/metadata_store_cmek/main.tf
@@ -0,0 +1,54 @@
+/**
+ * Copyright 2023 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.
+ */
+
+
+# [START aiplatform_create_metadata_store_cmek]
+resource "google_vertex_ai_metadata_store" "default" {
+  name        = "${random_id.default.hex}-example-store"
+  description = "Example metadata store"
+  provider    = google-beta
+  region      = "us-central1"
+  encryption_spec {
+    kms_key_name = google_kms_crypto_key.default.id
+  }
+
+  depends_on = [google_project_iam_member.default]
+}
+
+resource "random_id" "default" {
+  byte_length = 8
+}
+
+resource "google_kms_key_ring" "default" {
+  name     = "${random_id.default.hex}-example-keyring"
+  location = "us-central1"
+}
+
+resource "google_kms_crypto_key" "default" {
+  name     = "example-key"
+  key_ring = google_kms_key_ring.default.id
+}
+
+data "google_project" "default" {
+}
+
+# Enable the service account to encrypt/decrypt Cloud KMS keys
+resource "google_project_iam_member" "default" {
+  project = data.google_project.default.project_id
+  role    = "roles/cloudkms.cryptoKeyEncrypterDecrypter"
+  member  = "serviceAccount:service-${data.google_project.default.number}@gcp-sa-aiplatform.iam.gserviceaccount.com"
+}
+# [END aiplatform_create_metadata_store_cmek]