diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 3988d5160f..46012b8b87 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -6,6 +6,7 @@ ### Bundles * Validate that either source_code_path or git_source is set for apps ([#4632](https://github.com/databricks/cli/pull/4632)) +* direct: Sync registered model aliases via SetAlias/DeleteAlias ([#4637](https://github.com/databricks/cli/pull/4637)) ### Dependency updates diff --git a/acceptance/bundle/resources/registered_models/aliases/databricks.yml.tmpl b/acceptance/bundle/resources/registered_models/aliases/databricks.yml.tmpl new file mode 100644 index 0000000000..2418035eaa --- /dev/null +++ b/acceptance/bundle/resources/registered_models/aliases/databricks.yml.tmpl @@ -0,0 +1,10 @@ +bundle: + name: deploy-registered-models-aliases-$UNIQUE_NAME + +resources: + registered_models: + my_registered_model: + name: $NAME + comment: $COMMENT + catalog_name: $CATALOG_NAME + schema_name: $SCHEMA_NAME diff --git a/acceptance/bundle/resources/registered_models/aliases/out.test.toml b/acceptance/bundle/resources/registered_models/aliases/out.test.toml new file mode 100644 index 0000000000..8d6b9baeb5 --- /dev/null +++ b/acceptance/bundle/resources/registered_models/aliases/out.test.toml @@ -0,0 +1,7 @@ +Local = true +Cloud = true +RequiresUnityCatalog = true +RunsOnDbr = true + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/bundle/resources/registered_models/aliases/output.txt b/acceptance/bundle/resources/registered_models/aliases/output.txt new file mode 100644 index 0000000000..dc5f9fda67 --- /dev/null +++ b/acceptance/bundle/resources/registered_models/aliases/output.txt @@ -0,0 +1,85 @@ + +>>> export NAME=my-registered-model-[UNIQUE_NAME] + +>>> export COMMENT=test model + +>>> export CATALOG_NAME=mycatalog-[UNIQUE_NAME] + +>>> export SCHEMA_NAME=myschema-[UNIQUE_NAME] + +>>> [CLI] catalogs create mycatalog-[UNIQUE_NAME] +{ + "full_name": "mycatalog-[UNIQUE_NAME]" +} + +>>> [CLI] schemas create myschema-[UNIQUE_NAME] mycatalog-[UNIQUE_NAME] +{ + "full_name": "mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME]" +} + +=== create with aliases +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-aliases-[UNIQUE_NAME]/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> [CLI] registered-models get mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME].my-registered-model-[UNIQUE_NAME] --include-aliases +{ + "aliases": [ + { + "alias_name": "champion", + "version_num": 1 + }, + { + "alias_name": "staging", + "version_num": 2 + } + ] +} + +=== update: modify champion version, remove staging, add latest +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-aliases-[UNIQUE_NAME]/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> [CLI] registered-models get mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME].my-registered-model-[UNIQUE_NAME] --include-aliases +{ + "aliases": [ + { + "alias_name": "champion", + "version_num": 5 + }, + { + "alias_name": "latest", + "version_num": 3 + } + ] +} + +=== remove all aliases +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-aliases-[UNIQUE_NAME]/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> [CLI] registered-models get mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME].my-registered-model-[UNIQUE_NAME] --include-aliases +{ + "aliases": [] +} + +>>> [CLI] bundle destroy --auto-approve +The following resources will be deleted: + delete resources.registered_models.my_registered_model + +All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-aliases-[UNIQUE_NAME]/default + +Deleting files... +Destroy complete! + +>>> [CLI] schemas delete mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME] --force + +>>> [CLI] catalogs delete mycatalog-[UNIQUE_NAME] --force diff --git a/acceptance/bundle/resources/registered_models/aliases/script b/acceptance/bundle/resources/registered_models/aliases/script new file mode 100644 index 0000000000..8e4758d55d --- /dev/null +++ b/acceptance/bundle/resources/registered_models/aliases/script @@ -0,0 +1,44 @@ +trace export NAME="my-registered-model-$UNIQUE_NAME" +trace export COMMENT="test model" +trace export CATALOG_NAME="mycatalog-${UNIQUE_NAME}" +trace export SCHEMA_NAME="myschema-${UNIQUE_NAME}" +envsubst < databricks.yml.tmpl > databricks.yml + +trace $CLI catalogs create ${CATALOG_NAME} | jq '{full_name}' +trace $CLI schemas create ${SCHEMA_NAME} ${CATALOG_NAME} | jq '{full_name}' + +# Append aliases to the config. +cat >> databricks.yml <<'YAML' + aliases: + - alias_name: champion + version_num: 1 + - alias_name: staging + version_num: 2 +YAML + +cleanup() { + trace $CLI bundle destroy --auto-approve + trace $CLI schemas delete ${CATALOG_NAME}.${SCHEMA_NAME} --force + trace $CLI catalogs delete ${CATALOG_NAME} --force +} +trap cleanup EXIT + +deploy_and_get_aliases() { + trace $CLI bundle deploy + registered_model_id=$($CLI bundle summary --output json | jq -r '.resources.registered_models.my_registered_model.id') + trace $CLI registered-models get "${registered_model_id}" --include-aliases | jq '{aliases: [.aliases[]? | {alias_name, version_num}] | sort_by(.alias_name)}' +} + +title "create with aliases" +deploy_and_get_aliases + +title "update: modify champion version, remove staging, add latest" +update_file.py databricks.yml "version_num: 1" "version_num: 5" +update_file.py databricks.yml "staging" "latest" +update_file.py databricks.yml "version_num: 2" "version_num: 3" +deploy_and_get_aliases + +title "remove all aliases" +# Replace the aliases block with no aliases. +envsubst < databricks.yml.tmpl > databricks.yml +deploy_and_get_aliases diff --git a/acceptance/bundle/resources/registered_models/aliases/test.toml b/acceptance/bundle/resources/registered_models/aliases/test.toml new file mode 100644 index 0000000000..33508ef327 --- /dev/null +++ b/acceptance/bundle/resources/registered_models/aliases/test.toml @@ -0,0 +1,5 @@ +Cloud = true +Local = true +RecordRequests = false +RunsOnDbr = true +RequiresUnityCatalog = true diff --git a/bundle/direct/dresources/registered_model.go b/bundle/direct/dresources/registered_model.go index 666716b48b..abcbe159eb 100644 --- a/bundle/direct/dresources/registered_model.go +++ b/bundle/direct/dresources/registered_model.go @@ -4,11 +4,15 @@ import ( "context" "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/cli/libs/structs/structpath" "github.com/databricks/cli/libs/utils" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/catalog" + "golang.org/x/sync/errgroup" ) +var pathAliases = structpath.MustParsePath("aliases") + type ResourceRegisteredModel struct { client *databricks.WorkspaceClient } @@ -49,7 +53,7 @@ func (*ResourceRegisteredModel) RemapState(model *catalog.RegisteredModelInfo) * func (r *ResourceRegisteredModel) DoRead(ctx context.Context, id string) (*catalog.RegisteredModelInfo, error) { return r.client.RegisteredModels.Get(ctx, catalog.GetRegisteredModelRequest{ FullName: id, - IncludeAliases: false, + IncludeAliases: true, IncludeBrowse: false, ForceSendFields: nil, }) @@ -61,10 +65,15 @@ func (r *ResourceRegisteredModel) DoCreate(ctx context.Context, config *catalog. return "", nil, err } + // The Create API does not apply aliases, so we sync them separately. + if err := r.syncAliases(ctx, response.FullName, config.Aliases, nil); err != nil { + return "", nil, err + } + return response.FullName, response, nil } -func (r *ResourceRegisteredModel) DoUpdate(ctx context.Context, id string, config *catalog.CreateRegisteredModelRequest, _ Changes) (*catalog.RegisteredModelInfo, error) { +func (r *ResourceRegisteredModel) DoUpdate(ctx context.Context, id string, config *catalog.CreateRegisteredModelRequest, changes Changes) (*catalog.RegisteredModelInfo, error) { updateRequest := catalog.UpdateRegisteredModelRequest{ FullName: id, Comment: config.Comment, @@ -77,7 +86,9 @@ func (r *ResourceRegisteredModel) DoUpdate(ctx context.Context, id string, confi // Note: TF also does not support changing name without a recreate so the current behavior matches TF. NewName: "", - Aliases: config.Aliases, + // Aliases are synced separately via SetAlias/DeleteAlias calls because + // the Update API ignores the Aliases field. + Aliases: nil, BrowseOnly: config.BrowseOnly, CreatedAt: config.CreatedAt, CreatedBy: config.CreatedBy, @@ -95,6 +106,12 @@ func (r *ResourceRegisteredModel) DoUpdate(ctx context.Context, id string, confi return nil, err } + if changes.HasChange(pathAliases) { + if err := r.syncAliases(ctx, id, config.Aliases, nil); err != nil { + return nil, err + } + } + return response, nil } @@ -103,3 +120,63 @@ func (r *ResourceRegisteredModel) DoDelete(ctx context.Context, id string) error FullName: id, }) } + +// syncAliases compares desired and current aliases and calls SetAlias/DeleteAlias +// APIs to reconcile the difference. The Update API ignores the Aliases field, +// so separate API calls are required. +// If current is nil, the current aliases are fetched from the remote. +func (r *ResourceRegisteredModel) syncAliases(ctx context.Context, fullName string, desired, current []catalog.RegisteredModelAlias) error { + if current == nil { + remote, err := r.client.RegisteredModels.Get(ctx, catalog.GetRegisteredModelRequest{ + FullName: fullName, + IncludeAliases: true, + IncludeBrowse: false, + ForceSendFields: nil, + }) + if err != nil { + return err + } + current = remote.Aliases + } + + desiredByName := make(map[string]int, len(desired)) + for _, a := range desired { + desiredByName[a.AliasName] = a.VersionNum + } + + currentByName := make(map[string]int, len(current)) + for _, a := range current { + currentByName[a.AliasName] = a.VersionNum + } + + var eg errgroup.Group + + // Set new or updated aliases. + for name, version := range desiredByName { + if v, ok := currentByName[name]; ok && v == version { + continue + } + eg.Go(func() error { + _, err := r.client.RegisteredModels.SetAlias(ctx, catalog.SetRegisteredModelAliasRequest{ + FullName: fullName, + Alias: name, + VersionNum: version, + }) + return err + }) + } + + // Delete removed aliases. + for name := range currentByName { + if _, ok := desiredByName[name]; !ok { + eg.Go(func() error { + return r.client.RegisteredModels.DeleteAlias(ctx, catalog.DeleteAliasRequest{ + FullName: fullName, + Alias: name, + }) + }) + } + } + + return eg.Wait() +} diff --git a/bundle/direct/dresources/registered_model_test.go b/bundle/direct/dresources/registered_model_test.go new file mode 100644 index 0000000000..0f82d06d8e --- /dev/null +++ b/bundle/direct/dresources/registered_model_test.go @@ -0,0 +1,106 @@ +package dresources + +import ( + "context" + "testing" + + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/catalog" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func setupRegisteredModelTest(t *testing.T) *ResourceRegisteredModel { + server := testserver.New(t) + testserver.AddDefaultHandlers(server) + + client, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: server.URL, + Token: "testtoken", + }) + require.NoError(t, err) + + return (&ResourceRegisteredModel{}).New(client) +} + +func TestSyncAliases_AddsNewAliases(t *testing.T) { + r := setupRegisteredModelTest(t) + ctx := context.Background() + + id, _, err := r.DoCreate(ctx, &catalog.CreateRegisteredModelRequest{ + Name: "my_model", CatalogName: "main", SchemaName: "default", + Aliases: []catalog.RegisteredModelAlias{ + {AliasName: "champion", VersionNum: 1}, + {AliasName: "staging", VersionNum: 2}, + }, + }) + require.NoError(t, err) + + remote, err := r.DoRead(ctx, id) + require.NoError(t, err) + assert.Len(t, remote.Aliases, 2) + + m := aliasMap(remote.Aliases) + assert.Equal(t, 1, m["champion"]) + assert.Equal(t, 2, m["staging"]) +} + +func TestSyncAliases_UpdatesAndDeletesAliases(t *testing.T) { + r := setupRegisteredModelTest(t) + ctx := context.Background() + + id, _, err := r.DoCreate(ctx, &catalog.CreateRegisteredModelRequest{ + Name: "my_model", CatalogName: "main", SchemaName: "default", + Aliases: []catalog.RegisteredModelAlias{ + {AliasName: "champion", VersionNum: 1}, + {AliasName: "staging", VersionNum: 2}, + }, + }) + require.NoError(t, err) + + // Modify champion version, remove staging, add latest. + err = r.syncAliases(ctx, id, []catalog.RegisteredModelAlias{ + {AliasName: "champion", VersionNum: 5}, + {AliasName: "latest", VersionNum: 3}, + }, nil) + require.NoError(t, err) + + remote, err := r.DoRead(ctx, id) + require.NoError(t, err) + assert.Len(t, remote.Aliases, 2) + + m := aliasMap(remote.Aliases) + assert.Equal(t, 5, m["champion"]) + assert.Equal(t, 3, m["latest"]) + _, hasStaging := m["staging"] + assert.False(t, hasStaging) +} + +func TestSyncAliases_RemovesAllAliases(t *testing.T) { + r := setupRegisteredModelTest(t) + ctx := context.Background() + + id, _, err := r.DoCreate(ctx, &catalog.CreateRegisteredModelRequest{ + Name: "my_model", CatalogName: "main", SchemaName: "default", + Aliases: []catalog.RegisteredModelAlias{ + {AliasName: "champion", VersionNum: 1}, + }, + }) + require.NoError(t, err) + + err = r.syncAliases(ctx, id, nil, nil) + require.NoError(t, err) + + remote, err := r.DoRead(ctx, id) + require.NoError(t, err) + assert.Empty(t, remote.Aliases) +} + +func aliasMap(aliases []catalog.RegisteredModelAlias) map[string]int { + m := make(map[string]int, len(aliases)) + for _, a := range aliases { + m[a.AliasName] = a.VersionNum + } + return m +} diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 9e30cb5f0c..4aa77ff140 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -485,6 +485,14 @@ func AddDefaultHandlers(server *Server) { return MapDelete(req.Workspace, req.Workspace.RegisteredModels, req.Vars["full_name"]) }) + server.Handle("PUT", "/api/2.1/unity-catalog/models/{full_name}/aliases/{alias}", func(req Request) any { + return req.Workspace.RegisteredModelsSetAlias(req, req.Vars["full_name"], req.Vars["alias"]) + }) + + server.Handle("DELETE", "/api/2.1/unity-catalog/models/{full_name}/aliases/{alias}", func(req Request) any { + return req.Workspace.RegisteredModelsDeleteAlias(req.Vars["full_name"], req.Vars["alias"]) + }) + // Volumes: server.Handle("GET", "/api/2.1/unity-catalog/volumes/{full_name}", func(req Request) any { diff --git a/libs/testserver/registered_models.go b/libs/testserver/registered_models.go index 74815ae7a8..6063910f5d 100644 --- a/libs/testserver/registered_models.go +++ b/libs/testserver/registered_models.go @@ -84,3 +84,68 @@ func (s *FakeWorkspace) RegisteredModelsUpdate(req Request, fullName string) Res Body: existing, } } + +func (s *FakeWorkspace) RegisteredModelsSetAlias(req Request, fullName, alias string) Response { + defer s.LockUnlock()() + + existing, ok := s.RegisteredModels[fullName] + if !ok { + return Response{ + StatusCode: http.StatusNotFound, + Body: fmt.Sprintf("registered model %s not found", fullName), + } + } + + var setRequest catalog.SetRegisteredModelAliasRequest + if err := json.Unmarshal(req.Body, &setRequest); err != nil { + return Response{ + Body: fmt.Sprintf("internal error: %s", err), + StatusCode: http.StatusInternalServerError, + } + } + + newAlias := catalog.RegisteredModelAlias{ + AliasName: alias, + VersionNum: setRequest.VersionNum, + } + + // Update existing alias or append new one. + found := false + for i, a := range existing.Aliases { + if a.AliasName == alias { + existing.Aliases[i] = newAlias + found = true + break + } + } + if !found { + existing.Aliases = append(existing.Aliases, newAlias) + } + + s.RegisteredModels[fullName] = existing + return Response{ + Body: newAlias, + } +} + +func (s *FakeWorkspace) RegisteredModelsDeleteAlias(fullName, alias string) Response { + defer s.LockUnlock()() + + existing, ok := s.RegisteredModels[fullName] + if !ok { + return Response{ + StatusCode: http.StatusNotFound, + Body: fmt.Sprintf("registered model %s not found", fullName), + } + } + + for i, a := range existing.Aliases { + if a.AliasName == alias { + existing.Aliases = append(existing.Aliases[:i], existing.Aliases[i+1:]...) + break + } + } + + s.RegisteredModels[fullName] = existing + return Response{} +}