[Feature] Allow to specify budget policy for databricks_vector_search_endpoint (#4707)

## Changes
<!-- Summary of your changes that are easy to understand -->

For some reason, `custom_tags` cause configuration drift even if they
are marked as read-only, so temporary removed it (and we need to wait
for API anyway).


## Tests
<!-- 
How is this tested? Please see the checklist below and also describe any
other relevant tests
-->

- [x] `make test` run locally
- [x] relevant change in `docs/` folder
- [ ] covered with integration tests in `internal/acceptance`
- [ ] using Go SDK
- [ ] using TF Plugin Framework

Co-authored-by: vuong-nguyen <[email protected]>

Author: alexott

NEXT_CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### New Features and Improvements
 
+* Allow to specify budget policy for `databricks_vector_search_endpoint` [#4707](https://github.com/databricks/terraform-provider-databricks/pull/4707)
+
 ### Bug Fixes
 
  * Don't fail delete when `databricks_system_schema` can be disabled only by Databricks [#4727](https://github.com/databricks/terraform-provider-databricks/pull/4727)

docs/resources/vector_search_endpoint.md

@@ -18,10 +18,11 @@ resource "databricks_vector_search_endpoint" "this" {
 
 ## Argument Reference
 
-The following arguments are supported (change of any parameter leads to recreation of the resource):
+The following arguments are supported:
 
-* `name` - (Required) Name of the Mosaic AI Vector Search Endpoint to create.
-* `endpoint_type` (Required) Type of Mosaic AI 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).
+* `name` - (Required) Name of the Mosaic AI Vector Search Endpoint to create. (Change leads to recreation of the resource).
+* `endpoint_type` (Required) Type of Mosaic AI 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). (Change leads to recreation of the resource).
+* `budget_policy_id` - (Optional) The Budget Policy ID set for this resource.
 
 ## Attribute Reference
 
@@ -32,6 +33,7 @@ In addition to all the arguments above, the following attributes are exported:
 * `creation_timestamp` - Timestamp of endpoint creation (milliseconds).
 * `last_updated_user` - User who last updated the endpoint.
 * `last_updated_timestamp` - Timestamp of the last update to the endpoint (milliseconds).
+* `effective_budget_policy_id` - The effective budget policy ID.
 * `endpoint_id` - Unique internal identifier of the endpoint (UUID).
 * `num_indexes` - Number of indexes on the endpoint.
 * `endpoint_status` - Object describing the current status of the endpoint consisting of the following fields:

vectorsearch/resource_vector_search_endpoint.go

@@ -21,16 +21,20 @@ func ResourceVectorSearchEndpoint() common.Resource {
 			common.CustomizeSchemaPath(s, "name").SetRequired().SetForceNew()
 			common.CustomizeSchemaPath(s, "endpoint_type").SetRequired().SetForceNew()
 			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()
+			delete(s, "custom_tags")
+			for _, field := range []string{"creator", "creation_timestamp", "last_updated_timestamp",
+				"last_updated_user", "endpoint_status", "num_indexes", "effective_budget_policy_id"} {
+				common.CustomizeSchemaPath(s, field).SetReadOnly()
+			}
 			common.CustomizeSchemaPath(s).AddNewField("endpoint_id", &schema.Schema{
 				Type:     schema.TypeString,
 				Computed: true,
 			})
+			common.CustomizeSchemaPath(s).AddNewField("budget_policy_id", &schema.Schema{
+				Type:     schema.TypeString,
+				Optional: true,
+				Computed: true,
+			})
 
 			return s
 		})
@@ -72,9 +76,27 @@ func ResourceVectorSearchEndpoint() common.Resource {
 			if err != nil {
 				return err
 			}
+			d.Set("budget_policy_id", endpoint.EffectiveBudgetPolicyId)
 			d.Set("endpoint_id", endpoint.Id)
 			return nil
 		},
+		Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
+			w, err := c.WorkspaceClient()
+			if err != nil {
+				return err
+			}
+			if d.HasChange("budget_policy_id") {
+				_, err := w.VectorSearchEndpoints.UpdateEndpointBudgetPolicy(ctx, vectorsearch.PatchEndpointBudgetPolicyRequest{
+					EndpointName:   d.Id(),
+					BudgetPolicyId: d.Get("budget_policy_id").(string),
+				})
+				if err != nil {
+					return err
+				}
+			}
+			return nil
+		},
+
 		Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
 			w, err := c.WorkspaceClient()
 			if err != nil {

vectorsearch/resource_vector_search_endpoint_test.go

@@ -105,3 +105,41 @@ func TestVectorSearchEndpointCreateTimeoutError(t *testing.T) {
 	}.ExpectError(t, "timeout")
 
 }
+
+func TestVectorSearchEndpointUpdateBudgetPolicy(t *testing.T) {
+	qa.ResourceFixture{
+		MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
+			e := w.GetMockVectorSearchEndpointsAPI().EXPECT()
+			e.UpdateEndpointBudgetPolicy(mock.Anything, vectorsearch.PatchEndpointBudgetPolicyRequest{
+				EndpointName:   "abc",
+				BudgetPolicyId: "budget-123",
+			}).Return(&vectorsearch.PatchEndpointBudgetPolicyResponse{
+				EffectiveBudgetPolicyId: "budget-123",
+			}, nil)
+			e.GetEndpointByEndpointName(mock.Anything, "abc").Return(&vectorsearch.EndpointInfo{
+				Id:                      "1234-5678",
+				EffectiveBudgetPolicyId: "budget-123",
+				EndpointStatus:          &vectorsearch.EndpointStatus{State: "ONLINE"},
+				EndpointType:            "STANDARD",
+				Name:                    "abc",
+			}, nil)
+		},
+		Resource: ResourceVectorSearchEndpoint(),
+		Update:   true,
+		ID:       "abc",
+		InstanceState: map[string]string{
+			"endpoint_id":      "1234-5678",
+			"name":             "abc",
+			"endpoint_type":    "STANDARD",
+			"budget_policy_id": "budget-456",
+		},
+		HCL: `
+		name             = "abc"
+		endpoint_type    = "STANDARD"
+		budget_policy_id = "budget-123"
+		`,
+	}.ApplyAndExpectData(t, map[string]any{
+		"id":               "abc",
+		"budget_policy_id": "budget-123",
+	})
+}