Add support for Vector Search Endpoint (direct only)

Makefile

@@ -113,6 +113,11 @@ test-update-templates:
 generate-out-test-toml:
 	go test ./acceptance -run '^TestAccept$$' -only-out-test-toml -timeout=${LOCAL_TIMEOUT}
 
+# Regenerate refschema for new resources
+.PHONY: generate-refschema
+generate-refschema:
+	go test ./acceptance -run '^TestAccept/bundle/refschema$$' -update -timeout=${LOCAL_TIMEOUT}
+
 # Updates acceptance test output (integration tests, requires access)
 .PHONY: test-update-aws
 test-update-aws:

NEXT_CHANGELOG.md

@@ -15,6 +15,7 @@
 * Fix resource references not correctly resolved in apps config section ([#4964](https://github.com/databricks/cli/pull/4964))
 * Allow run_as for dashboards with embed_credentials set to false ([#4961](https://github.com/databricks/cli/pull/4961))
 * direct: Pass changed fields into update mask for apps instead of wildcard ([#4963](https://github.com/databricks/cli/pull/4963))
+* engine/direct: Added support for Vector Search Endpoints ([#4887](https://github.com/databricks/cli/pull/4887))
 
 ### Dependency updates
 

acceptance/bundle/deployment/bind/vector_search_endpoint/databricks.yml.tmpl

@@ -0,0 +1,11 @@
+bundle:
+  name: test-bundle-$UNIQUE_NAME
+
+sync:
+  paths: []
+
+resources:
+  vector_search_endpoints:
+    endpoint1:
+      name: $ENDPOINT_NAME
+      endpoint_type: STANDARD

acceptance/bundle/deployment/bind/vector_search_endpoint/output.txt

@@ -0,0 +1,43 @@
+
+>>> [CLI] vector-search-endpoints create-endpoint test-vs-endpoint-[UNIQUE_NAME] STANDARD
+{
+  "id": "[UUID]",
+  "name": "test-vs-endpoint-[UNIQUE_NAME]",
+  "endpoint_type": "STANDARD"
+}
+
+>>> [CLI] bundle deployment bind endpoint1 test-vs-endpoint-[UNIQUE_NAME] --auto-approve
+Updating deployment state...
+Successfully bound vector_search_endpoint with an id 'test-vs-endpoint-[UNIQUE_NAME]'
+Run 'bundle deploy' to deploy changes to your workspace
+
+>>> [CLI] bundle deploy --auto-approve
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle-[UNIQUE_NAME]/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> [CLI] vector-search-endpoints get-endpoint test-vs-endpoint-[UNIQUE_NAME]
+{
+  "id": "[UUID]",
+  "name": "test-vs-endpoint-[UNIQUE_NAME]",
+  "endpoint_type": "STANDARD"
+}
+
+>>> [CLI] bundle deployment unbind endpoint1
+Updating deployment state...
+
+>>> [CLI] bundle destroy --auto-approve
+All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/test-bundle-[UNIQUE_NAME]/default
+
+Deleting files...
+Destroy complete!
+
+>>> [CLI] vector-search-endpoints get-endpoint test-vs-endpoint-[UNIQUE_NAME]
+{
+  "id": "[UUID]",
+  "name": "test-vs-endpoint-[UNIQUE_NAME]",
+  "endpoint_type": "STANDARD"
+}
+
+>>> [CLI] vector-search-endpoints delete-endpoint test-vs-endpoint-[UNIQUE_NAME]

acceptance/bundle/deployment/bind/vector_search_endpoint/script

@@ -0,0 +1,23 @@
+ENDPOINT_NAME="test-vs-endpoint-$UNIQUE_NAME"
+export ENDPOINT_NAME
+envsubst < databricks.yml.tmpl > databricks.yml
+
+cleanup() {
+    trace $CLI vector-search-endpoints delete-endpoint "${ENDPOINT_NAME}"
+}
+trap cleanup EXIT
+
+trace $CLI vector-search-endpoints create-endpoint "${ENDPOINT_NAME}" STANDARD | jq '{id, name, endpoint_type}'
+
+trace $CLI bundle deployment bind endpoint1 "${ENDPOINT_NAME}" --auto-approve
+
+trace $CLI bundle deploy --auto-approve
+
+trace $CLI vector-search-endpoints get-endpoint "${ENDPOINT_NAME}" | jq '{id, name, endpoint_type}'
+
+trace $CLI bundle deployment unbind endpoint1
+
+trace $CLI bundle destroy --auto-approve
+
+# Read the pre-defined endpoint again (expecting it still exists and is not deleted):
+trace $CLI vector-search-endpoints get-endpoint "${ENDPOINT_NAME}" | jq '{id, name, endpoint_type}'

acceptance/bundle/deployment/bind/vector_search_endpoint/test.toml

@@ -0,0 +1,10 @@
+Local = true
+Cloud = true
+
+Ignore = [
+    ".databricks",
+    "databricks.yml",
+]
+
+[EnvMatrix]
+  DATABRICKS_BUNDLE_ENGINE = ["direct"]

acceptance/bundle/invariant/configs/vector_search_endpoint.yml.tmpl

@@ -0,0 +1,8 @@
+bundle:
+  name: test-bundle-$UNIQUE_NAME
+
+resources:
+  vector_search_endpoints:
+    foo:
+      name: test-endpoint-$UNIQUE_NAME
+      endpoint_type: STANDARD

acceptance/bundle/invariant/continue_293/test.toml

@@ -5,3 +5,6 @@ EnvMatrixExclude.no_grant_ref = ["INPUT_CONFIG=schema_grant_ref.yml.tmpl"]
 
 # Model permissions did not work until 0.297.0 https://github.com/databricks/cli/pull/4941
 EnvMatrixExclude.no_model_with_permissions = ["INPUT_CONFIG=model_with_permissions.yml.tmpl"]
+
+# vector_search_endpoints resource is not supported on v0.293.0
+EnvMatrixExclude.no_vector_search_endpoint = ["INPUT_CONFIG=vector_search_endpoint.yml.tmpl"]

acceptance/bundle/invariant/migrate/test.toml

@@ -1,3 +1,6 @@
+# vector_search_endpoints has no terraform converter
+EnvMatrixExclude.no_vector_search_endpoint = ["INPUT_CONFIG=vector_search_endpoint.yml.tmpl"]
+
 # Error: Catalog resources are only supported with direct deployment mode
 EnvMatrixExclude.no_catalog = ["INPUT_CONFIG=catalog.yml.tmpl"]
 EnvMatrixExclude.no_external_location = ["INPUT_CONFIG=external_location.yml.tmpl"]

acceptance/bundle/invariant/test.toml

@@ -52,6 +52,7 @@ EnvMatrix.INPUT_CONFIG = [
     "secret_scope_default_backend_type.yml.tmpl",
     "secret_scope_with_permissions.yml.tmpl",
     "synced_database_table.yml.tmpl",
+    "vector_search_endpoint.yml.tmpl",
     "volume.yml.tmpl",
 ]
 

acceptance/bundle/refschema/out.fields.txt

@@ -2885,6 +2885,38 @@ resources.synced_database_tables.*.spec.source_table_full_name	string	ALL
 resources.synced_database_tables.*.spec.timeseries_key	string	ALL
 resources.synced_database_tables.*.unity_catalog_provisioning_state	database.ProvisioningInfoState	ALL
 resources.synced_database_tables.*.url	string	INPUT
+resources.vector_search_endpoints.*.budget_policy_id	string	INPUT	STATE
+resources.vector_search_endpoints.*.creation_timestamp	int64	REMOTE
+resources.vector_search_endpoints.*.creator	string	REMOTE
+resources.vector_search_endpoints.*.custom_tags	[]vectorsearch.CustomTag	REMOTE
+resources.vector_search_endpoints.*.custom_tags[*]	vectorsearch.CustomTag	REMOTE
+resources.vector_search_endpoints.*.custom_tags[*].key	string	REMOTE
+resources.vector_search_endpoints.*.custom_tags[*].value	string	REMOTE
+resources.vector_search_endpoints.*.effective_budget_policy_id	string	REMOTE
+resources.vector_search_endpoints.*.endpoint_status	*vectorsearch.EndpointStatus	REMOTE
+resources.vector_search_endpoints.*.endpoint_status.message	string	REMOTE
+resources.vector_search_endpoints.*.endpoint_status.state	vectorsearch.EndpointStatusState	REMOTE
+resources.vector_search_endpoints.*.endpoint_type	vectorsearch.EndpointType	ALL
+resources.vector_search_endpoints.*.endpoint_uuid	string	REMOTE
+resources.vector_search_endpoints.*.id	string	INPUT	REMOTE
+resources.vector_search_endpoints.*.last_updated_timestamp	int64	REMOTE
+resources.vector_search_endpoints.*.last_updated_user	string	REMOTE
+resources.vector_search_endpoints.*.lifecycle	resources.Lifecycle	INPUT
+resources.vector_search_endpoints.*.lifecycle.prevent_destroy	bool	INPUT
+resources.vector_search_endpoints.*.min_qps	int64	INPUT	STATE
+resources.vector_search_endpoints.*.modified_status	string	INPUT
+resources.vector_search_endpoints.*.name	string	ALL
+resources.vector_search_endpoints.*.num_indexes	int	REMOTE
+resources.vector_search_endpoints.*.scaling_info	*vectorsearch.EndpointScalingInfo	REMOTE
+resources.vector_search_endpoints.*.scaling_info.requested_min_qps	int64	REMOTE
+resources.vector_search_endpoints.*.scaling_info.state	vectorsearch.ScalingChangeState	REMOTE
+resources.vector_search_endpoints.*.url	string	INPUT
+resources.vector_search_endpoints.*.permissions.object_id	string	ALL
+resources.vector_search_endpoints.*.permissions[*]	dresources.StatePermission	ALL
+resources.vector_search_endpoints.*.permissions[*].group_name	string	ALL
+resources.vector_search_endpoints.*.permissions[*].level	iam.PermissionLevel	ALL
+resources.vector_search_endpoints.*.permissions[*].service_principal_name	string	ALL
+resources.vector_search_endpoints.*.permissions[*].user_name	string	ALL
 resources.volumes.*.access_point	string	REMOTE
 resources.volumes.*.browse_only	bool	REMOTE
 resources.volumes.*.catalog_name	string	ALL

acceptance/bundle/resources/permissions/analyze_requests.py

@@ -3,10 +3,10 @@
 Analyze all requests recorded in subtests to highlight differences between direct and terraform.
 """
 
-import os
-import re
 import json
+import re
 import sys
+import tomllib
 from pathlib import Path
 from difflib import unified_diff
 
@@ -91,6 +91,20 @@ def to_slash(x):
     return str(x).replace("\\", "/")
 
 
+def load_supported_engines(path):
+    current = path
+    while True:
+        for name in ("out.test.toml", "test.toml"):
+            config_file = current / name
+            if config_file.exists():
+                with config_file.open("rb") as fobj:
+                    config = tomllib.load(fobj)
+                return set(config.get("EnvMatrix", {}).get("DATABRICKS_BUNDLE_ENGINE", []))
+        if current == Path("."):
+            return set()
+        current = current.parent
+
+
 def main():
     current_dir = Path(".")
 
@@ -104,10 +118,13 @@ def main():
         terraform_file = direct_file.parent / direct_file.name.replace(".direct.", ".terraform.")
 
         fname = to_slash(direct_file)
+        supported_engines = load_supported_engines(direct_file.parent)
 
         if terraform_file.exists():
             result, diff = compare_files(direct_file, terraform_file)
             print(result + " " + fname + diff)
+        elif "terraform" not in supported_engines:
+            print(f"DIRECT_ONLY {fname}")
         else:
             print(f"ERROR {fname}: Missing terraform file {to_slash(terraform_file)}")
 

acceptance/bundle/resources/permissions/output.txt

@@ -411,3 +411,5 @@ DIFF  target_permissions/out.requests_delete.direct.json
    {
      "body": {
        "job_id": "[NUMID]"
+DIRECT_ONLY vector_search_endpoints/current_can_manage/out.requests.deploy.direct.json
+DIRECT_ONLY vector_search_endpoints/current_can_manage/out.requests.destroy.direct.json

acceptance/bundle/resources/permissions/vector_search_endpoints/current_can_manage/databricks.yml

@@ -0,0 +1,17 @@
+bundle:
+  name: test-bundle
+
+resources:
+  vector_search_endpoints:
+    foo:
+      name: vs-permissions-endpoint
+      endpoint_type: STANDARD
+      permissions:
+        - level: CAN_USE
+          user_name: viewer@example.com
+        - level: CAN_MANAGE
+          group_name: data-team
+        - level: CAN_MANAGE
+          service_principal_name: f37d18cd-98a8-4db5-8112-12dd0a6bfe38
+        - level: CAN_MANAGE
+          user_name: tester@databricks.com

acceptance/bundle/resources/permissions/vector_search_endpoints/current_can_manage/out.plan.direct.json

@@ -0,0 +1,50 @@
+{
+  "plan_version": 2,
+  "cli_version": "[DEV_VERSION]",
+  "plan": {
+    "resources.vector_search_endpoints.foo": {
+      "action": "create",
+      "new_state": {
+        "value": {
+          "endpoint_type": "STANDARD",
+          "name": "vs-permissions-endpoint"
+        }
+      }
+    },
+    "resources.vector_search_endpoints.foo.permissions": {
+      "depends_on": [
+        {
+          "node": "resources.vector_search_endpoints.foo",
+          "label": "${resources.vector_search_endpoints.foo.endpoint_uuid}"
+        }
+      ],
+      "action": "create",
+      "new_state": {
+        "value": {
+          "object_id": "",
+          "__embed__": [
+            {
+              "level": "CAN_USE",
+              "user_name": "viewer@example.com"
+            },
+            {
+              "level": "CAN_MANAGE",
+              "group_name": "data-team"
+            },
+            {
+              "level": "CAN_MANAGE",
+              "service_principal_name": "[UUID]"
+            },
+            {
+              "level": "CAN_MANAGE",
+              "user_name": "[USERNAME]"
+            }
+          ]
+        },
+        "vars": {
+          "object_id": "/vector-search-endpoints/${resources.vector_search_endpoints.foo.endpoint_uuid}"
+        }
+      }
+    }
+  }
+}

acceptance/bundle/resources/permissions/vector_search_endpoints/current_can_manage/out.requests.deploy.direct.json

@@ -0,0 +1,24 @@
+{
+  "method": "PUT",
+  "path": "/api/2.0/permissions/vector-search-endpoints/[UUID]",
+  "body": {
+    "access_control_list": [
+      {
+        "permission_level": "CAN_USE",
+        "user_name": "viewer@example.com"
+      },
+      {
+        "group_name": "data-team",
+        "permission_level": "CAN_MANAGE"
+      },
+      {
+        "permission_level": "CAN_MANAGE",
+        "service_principal_name": "[UUID]"
+      },
+      {
+        "permission_level": "CAN_MANAGE",
+        "user_name": "[USERNAME]"
+      }
+    ]
+  }
+}

acceptance/bundle/resources/permissions/vector_search_endpoints/current_can_manage/out.requests.destroy.direct.json

@@ -0,0 +1,4 @@
+{
+  "method": "DELETE",
+  "path": "/api/2.0/vector-search/endpoints/vs-permissions-endpoint"
+}

acceptance/bundle/resources/permissions/vector_search_endpoints/current_can_manage/output.txt

@@ -0,0 +1,35 @@
+
+>>> [CLI] bundle validate -o json
+[
+  {
+    "level": "CAN_USE",
+    "user_name": "viewer@example.com"
+  },
+  {
+    "group_name": "data-team",
+    "level": "CAN_MANAGE"
+  },
+  {
+    "level": "CAN_MANAGE",
+    "service_principal_name": "[UUID]"
+  },
+  {
+    "level": "CAN_MANAGE",
+    "user_name": "[USERNAME]"
+  }
+]
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> [CLI] bundle destroy --auto-approve
+The following resources will be deleted:
+  delete resources.vector_search_endpoints.foo
+
+All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/test-bundle/default
+
+Deleting files...
+Destroy complete!