Skip to content

Commit

Permalink
feat(vertexai): add the data source for Vertex AI Index (GoogleCloudP…
Browse files Browse the repository at this point in the history
  • Loading branch information
shotarok authored and ericayyliu committed Jul 26, 2023
1 parent 92f3cd5 commit 910b56d
Show file tree
Hide file tree
Showing 4 changed files with 151 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
package google

import (
"fmt"

"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
"github.com/hashicorp/terraform-provider-google/google/tpgresource"
transport_tpg "github.com/hashicorp/terraform-provider-google/google/transport"
)

func dataSourceVertexAIIndex() *schema.Resource {

dsSchema := datasourceSchemaFromResourceSchema(ResourceVertexAIIndex().Schema)

addRequiredFieldsToSchema(dsSchema, "name", "region")
addOptionalFieldsToSchema(dsSchema, "project")

return &schema.Resource{
Read: dataSourceVertexAIIndexRead,
Schema: dsSchema,
}
}

func dataSourceVertexAIIndexRead(d *schema.ResourceData, meta interface{}) error {
config := meta.(*transport_tpg.Config)

id, err := tpgresource.ReplaceVars(d, config, "projects/{{project}}/locations/{{region}}/indexes/{{name}}")
if err != nil {
return fmt.Errorf("Error constructing id: %s", err)
}
d.SetId(id)
return resourceVertexAIIndexRead(d, meta)
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
package google

import (
"testing"

"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
"github.com/hashicorp/terraform-provider-google/google/acctest"
)

func TestAccDataSourceVertexAIIndex_basic(t *testing.T) {
t.Parallel()

context := map[string]interface{}{
"project": acctest.GetTestProjectFromEnv(),
"random_suffix": RandString(t, 10),
}

VcrTest(t, resource.TestCase{
PreCheck: func() { acctest.AccTestPreCheck(t) },
ProtoV5ProviderFactories: ProtoV5ProviderFactories(t),
CheckDestroy: testAccCheckVertexAIIndexDestroyProducer_basic(t),
Steps: []resource.TestStep{
{
Config: testAccDataSourceVertexAIIndex_basic(context),
Check: resource.ComposeTestCheckFunc(
acctest.CheckDataSourceStateMatchesResourceStateWithIgnores(
"data.google_vertex_ai_index.foo",
"google_vertex_ai_index.index",
// The projects.locations.indexes.get doesn't return the following fields
map[string]struct{}{
"metadata.0.contents_delta_uri": {},
"metadata.0.is_complete_overwrite": {},
},
),
),
},
},
})
}

func testAccDataSourceVertexAIIndex_basic(context map[string]interface{}) string {
return Nprintf(`
resource "google_storage_bucket" "bucket" {
name = "%{project}-tf-test-vertex-ai-index-test%{random_suffix}" # Every bucket name must be globally unique
location = "us-central1"
uniform_bucket_level_access = true
}
# The sample data comes from the following link:
# https://cloud.google.com/vertex-ai/docs/matching-engine/filtering#specify-namespaces-tokens
resource "google_storage_bucket_object" "data" {
name = "contents/data.json"
bucket = google_storage_bucket.bucket.name
content = <<EOF
{"id": "42", "embedding": [0.5, 1.0], "restricts": [{"namespace": "class", "allow": ["cat", "pet"]},{"namespace": "category", "allow": ["feline"]}]}
{"id": "43", "embedding": [0.6, 1.0], "restricts": [{"namespace": "class", "allow": ["dog", "pet"]},{"namespace": "category", "allow": ["canine"]}]}
EOF
}
resource "google_vertex_ai_index" "index" {
labels = {
foo = "bar"
}
region = "us-central1"
display_name = "tf-test-test-index%{random_suffix}"
description = "index for test"
metadata {
contents_delta_uri = "gs://${google_storage_bucket.bucket.name}/contents"
config {
dimensions = 2
approximate_neighbors_count = 150
distance_measure_type = "DOT_PRODUCT_DISTANCE"
algorithm_config {
tree_ah_config {
leaf_node_embedding_count = 500
leaf_nodes_to_search_percent = 7
}
}
}
}
index_update_method = "BATCH_UPDATE"
}
data "google_vertex_ai_index" "foo" {
name = google_vertex_ai_index.index.name
region = google_vertex_ai_index.index.region
project = google_vertex_ai_index.index.project
}
`, context)
}
1 change: 1 addition & 0 deletions mmv1/third_party/terraform/utils/provider.go.erb
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,7 @@ func Provider() *schema.Provider {
"google_tpu_tensorflow_versions": DataSourceTpuTensorflowVersions(),
"google_vpc_access_connector": DataSourceVPCAccessConnector(),
"google_redis_instance": DataSourceGoogleRedisInstance(),
"google_vertex_ai_index": dataSourceVertexAIIndex(),
// ####### END datasources ###########
},
ResourcesMap: ResourceMap(),
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
---
subcategory: "VertexAI"
description: |-
A representation of a collection of database items organized in a way that allows for approximate nearest neighbor (a.k.a ANN) algorithms search.
---

# google\_vertex\_ai\_index

A representation of a collection of database items organized in a way that allows for approximate nearest neighbor (a.k.a ANN) algorithms search.


## Argument Reference

The following arguments are supported:

* `name` - (Required) The name of the index.

* `region` - (Required) The region of the index.

- - -

* `project` - (Optional) The ID of the project in which the resource belongs.

## Attributes Reference

See [google_vertex_ai_index](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/vertex_ai_index) resource for details of the available attributes.

0 comments on commit 910b56d

Please sign in to comment.