Sync registered model aliases via SetAlias/DeleteAlias in direct mode

Signed-off-by: Varun Deep Saini <varun.23bcs10048@ms.sst.scaler.com>

Author: Varun Deep Saini

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,15 @@
+bundle:
+  name: deploy-registered-models-aliases-$UNIQUE_NAME
+
+resources:
+  registered_models:
+    my_registered_model:
+      name: my-registered-model-$UNIQUE_NAME
+      comment: "test model"
+      catalog_name: main
+      schema_name: default
+      aliases:
+        - alias_name: champion
+          version_num: 1
+        - alias_name: staging
+          version_num: 2

acceptance/bundle/resources/registered_models/aliases/script

@@ -0,0 +1,34 @@
+envsubst < databricks.yml.tmpl > databricks.yml
+
+cleanup() {
+    trace $CLI bundle destroy --auto-approve
+}
+trap cleanup EXIT
+
+get_aliases() {
+    registered_model_id=$($CLI bundle summary --output json | jq -r '.resources.registered_models.my_registered_model.id')
+    $CLI registered-models get "${registered_model_id}" --include-aliases | jq '[.aliases[] | {alias_name, version_num}] | sort_by(.alias_name)'
+}
+
+title "create with aliases"
+trace $CLI bundle deploy
+trace print_requests.py //unity-catalog/models --sort > out.create-requests.$DATABRICKS_BUNDLE_ENGINE.json
+trace 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"
+trace $CLI bundle deploy
+trace print_requests.py //unity-catalog/models --sort > out.update-requests.$DATABRICKS_BUNDLE_ENGINE.json
+trace get_aliases
+
+title "remove all aliases"
+update_file.py databricks.yml "      aliases:" ""
+update_file.py databricks.yml "        - alias_name: champion" ""
+update_file.py databricks.yml "          version_num: 5" ""
+update_file.py databricks.yml "        - alias_name: latest" ""
+update_file.py databricks.yml "          version_num: 3" ""
+trace $CLI bundle deploy
+trace print_requests.py //unity-catalog/models --sort > out.remove-requests.$DATABRICKS_BUNDLE_ENGINE.json
+trace get_aliases

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

@@ -0,0 +1,3 @@
+Cloud = true
+Local = 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
 }
@@ -48,10 +52,9 @@ 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,
-		IncludeBrowse:   false,
-		ForceSendFields: nil,
+		FullName:       id,
+		IncludeAliases: true,
+		IncludeBrowse:  false,
 	})
 }
 
@@ -61,10 +64,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 +85,8 @@ 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.
 		BrowseOnly:      config.BrowseOnly,
 		CreatedAt:       config.CreatedAt,
 		CreatedBy:       config.CreatedBy,
@@ -95,6 +104,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 +118,61 @@ 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,
+		})
+		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()
+}

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

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 {

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{}
+}