Sync registered model aliases via SetAlias/DeleteAlias in direct mode

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
 

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

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

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

acceptance/bundle/resources/registered_models/aliases/test.toml

@@ -0,0 +1,5 @@
+Cloud = true
+Local = true
+RecordRequests = false
+RunsOnDbr = true
+RequiresUnityCatalog = true

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()
+}