Skip to content

Commit

Permalink
WIP: Add support for Vector Search Endpoints
Browse files Browse the repository at this point in the history
Unit & integration tests needs to be added...
  • Loading branch information
alexott committed Jan 31, 2024
1 parent 510fd60 commit 1d54f44
Show file tree
Hide file tree
Showing 4 changed files with 150 additions and 0 deletions.
46 changes: 46 additions & 0 deletions docs/resources/vector_search_endpoint.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
---
subcategory: "Vector Search"
---
# databricks_vector_search_endpoint Resource

This resource allows you to create [Vector Search Endpoint](https://docs.databricks.com/en/generative-ai/vector-search.html) in Databricks. Vector Search is a serverless similarity search engine that allows you to store a vector representation of your data, including metadata, in a vector database. The Vector Search Endpoint is used to create and access vector search indexes.

## Example Usage


```hcl
resource "databricks_vector_search_endpoint" "this" {
name = "vector-search-test"
endpoint_type = "STANDARD"
}
```

## Argument Reference

The following arguments are supported:

* `name` - (Required) Name of the Vector Search Endpoint to create. If name is changed, Vector Search Endpoint is recreated.
* `endpoint_type` (Required) type of Vector Search Endpoint. Currently only accepting single value: `STANDARD` (See [documentation](https://docs.databricks.com/api/workspace/vectorsearchendpoints/createendpoint) for the list of currently supported values). If it changed, Vector Search Endpoint is recreated.

## Attribute Reference

In addition to all arguments above, the following attributes are exported:

* `id` - The same as the name of the endpoint.
* `creator` - Creator of the endpoint.
* `creation_timestamp` - Timestamp of endpoint creation (milliseconds).
* `last_updated_user` - User who last updated the endpoint.
* `last_updated_timestamp` - Timestamp of last update to the endpoint (milliseconds).
* `endpoint_id` - Unique identifier of the endpoint.
* `num_indexes` - Current status of the endpoint.
* `endpoint_status` - Object describing the current status of the endpoint consisting of following fields:
* `state` - Current state of the endpoint. Currently following values are supported: `PROVISIONING`, `ONLINE`, `OFFLINE`.
* `message` - Additional status message.

## Import

The resource can be imported using the name of the Vector Search Endpoint

```bash
terraform import databricks_vector_search_endpoint.this <endpoint-name>
```
2 changes: 2 additions & 0 deletions provider/provider.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ import (
"github.com/databricks/terraform-provider-databricks/sql"
"github.com/databricks/terraform-provider-databricks/storage"
"github.com/databricks/terraform-provider-databricks/tokens"
"github.com/databricks/terraform-provider-databricks/vectorsearch"
"github.com/databricks/terraform-provider-databricks/workspace"
)

Expand Down Expand Up @@ -171,6 +172,7 @@ func DatabricksProvider() *schema.Provider {
"databricks_user": scim.ResourceUser(),
"databricks_user_instance_profile": aws.ResourceUserInstanceProfile(),
"databricks_user_role": aws.ResourceUserRole(),
"databricks_vector_search_endpoint": vectorsearch.ResourceVectorSearchEndpoint(),
"databricks_volume": catalog.ResourceVolume(),
"databricks_workspace_conf": workspace.ResourceWorkspaceConf(),
"databricks_workspace_file": workspace.ResourceWorkspaceFile(),
Expand Down
85 changes: 85 additions & 0 deletions vectorsearch/resource_vector_search_endpoint.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
package vectorsearch

import (
"context"
"time"

"github.com/databricks/databricks-sdk-go/retries"
"github.com/databricks/terraform-provider-databricks/common"
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"

"github.com/databricks/databricks-sdk-go/service/vectorsearch"
)

const DefaultProvisionTimeout = 45 * time.Minute

func ResourceVectorSearchEndpoint() *schema.Resource {
s := common.StructToSchema(
vectorsearch.EndpointInfo{},
func(s map[string]*schema.Schema) map[string]*schema.Schema {
common.CustomizeSchemaPath(s, "name").SetRequired()
s["name"].ForceNew = true
s["endpoint_type"].ForceNew = true
common.CustomizeSchemaPath(s, "endpoint_type").SetRequired()
delete(s, "id")
common.CustomizeSchemaPath(s, "creator").SetReadOnly()
common.CustomizeSchemaPath(s, "creation_timestamp").SetReadOnly()
common.CustomizeSchemaPath(s, "last_updated_timestamp").SetReadOnly()
common.CustomizeSchemaPath(s, "last_updated_user").SetReadOnly()
common.CustomizeSchemaPath(s, "endpoint_status").SetReadOnly()
common.CustomizeSchemaPath(s, "num_indexes").SetReadOnly()
common.CustomizeSchemaPath(s).AddNewField("endpoint_id", &schema.Schema{
Type: schema.TypeString,
Computed: true,
})

return s
})

return common.Resource{
Create: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
w, err := c.WorkspaceClient()
if err != nil {
return err
}
var req vectorsearch.CreateEndpoint
common.DataToStructPointer(d, s, &req)
endpoint, err := w.VectorSearchEndpoints.CreateEndpointAndWait(ctx, req,
retries.Timeout[vectorsearch.EndpointInfo](d.Timeout(schema.TimeoutCreate)))
if err != nil {
return err
}
d.SetId(endpoint.Name)
return nil
},
Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
w, err := c.WorkspaceClient()
if err != nil {
return err
}
endpoint, err := w.VectorSearchEndpoints.GetEndpointByEndpointName(ctx, d.Id())
if err != nil {
return err
}
err = common.StructToData(*endpoint, s, d)
if err != nil {
return err
}
d.Set("endpoint_id", endpoint.Id)
return nil
},
Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
w, err := c.WorkspaceClient()
if err != nil {
return err
}
return w.VectorSearchEndpoints.DeleteEndpointByEndpointName(ctx, d.Id())
},
StateUpgraders: []schema.StateUpgrader{},
Schema: s,
SchemaVersion: 0,
Timeouts: &schema.ResourceTimeout{
Create: schema.DefaultTimeout(DefaultProvisionTimeout),
},
}.ToResource()
}
17 changes: 17 additions & 0 deletions vectorsearch/resource_vector_search_endpoint_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
package vectorsearch

import (
"testing"

"github.com/databricks/terraform-provider-databricks/qa"

"github.com/stretchr/testify/assert"
)

func TestVectorSearchEndpointCornerCases(t *testing.T) {
qa.ResourceCornerCases(t, ResourceVectorSearchEndpoint())
}

func TestVectorSearchEndpointCreate(t *testing.T) {
assert.Equal(t, 1, 1)
}

0 comments on commit 1d54f44

Please sign in to comment.