Add support for models in direct deployment

acceptance/bundle/deploy/models/basic/output.txt

@@ -0,0 +1,113 @@
+
+>>> export MODEL_NAME=original-name-[UNIQUE_NAME]
+
+>>> export MODEL_DESCRIPTION=original-description-[UNIQUE_NAME]
+
+=== create a model
+>>> [CLI] bundle plan
+create models.my_model
+
+Plan: 1 to add, 0 to change, 0 to delete, 0 unchanged
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-models-basic-[UNIQUE_NAME]/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> [CLI] model-registry get-model original-name-[UNIQUE_NAME]
+{
+  "name": "original-name-[UNIQUE_NAME]",
+  "description": "original-description-[UNIQUE_NAME]",
+  "tags": [
+    {
+      "key": "key1",
+      "value": "value1"
+    }
+  ]
+}
+
+>>> export MODEL_DESCRIPTION=new-description-[UNIQUE_NAME]
+
+=== update the description, this should update the description remotely as well
+>>> [CLI] bundle plan
+update models.my_model
+
+Plan: 0 to add, 1 to change, 0 to delete, 0 unchanged
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-models-basic-[UNIQUE_NAME]/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> [CLI] model-registry get-model original-name-[UNIQUE_NAME]
+{
+  "name": "original-name-[UNIQUE_NAME]",
+  "description": "new-description-[UNIQUE_NAME]",
+  "tags": [
+    {
+      "key": "key1",
+      "value": "value1"
+    }
+  ]
+}
+
+>>> export MODEL_NAME=new-name-[UNIQUE_NAME]
+
+=== update the name, this should recreate the model with the new name
+>>> [CLI] bundle plan
+recreate models.my_model
+
+Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-models-basic-[UNIQUE_NAME]/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> [CLI] model-registry get-model new-name-[UNIQUE_NAME]
+{
+  "name": "new-name-[UNIQUE_NAME]",
+  "description": "new-description-[UNIQUE_NAME]",
+  "tags": [
+    {
+      "key": "key1",
+      "value": "value1"
+    }
+  ]
+}
+
+=== add a new tag, this should be a no-op
+>>> [CLI] bundle plan
+update models.my_model
+
+Plan: 0 to add, 1 to change, 0 to delete, 0 unchanged
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-models-basic-[UNIQUE_NAME]/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> [CLI] model-registry get-model new-name-[UNIQUE_NAME]
+{
+  "name": "new-name-[UNIQUE_NAME]",
+  "description": "new-description-[UNIQUE_NAME]",
+  "tags": [
+    {
+      "key": "key1",
+      "value": "value1"
+    }
+  ]
+}
+
+>>> [CLI] bundle destroy --auto-approve
+The following resources will be deleted:
+  delete model my_model
+
+All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/deploy-models-basic-[UNIQUE_NAME]/default
+
+Deleting files...
+Destroy complete!

acceptance/bundle/deploy/models/basic/script

@@ -0,0 +1,32 @@
+cleanup() {
+    trace $CLI bundle destroy --auto-approve
+}
+trap cleanup EXIT
+
+trace export MODEL_NAME=original-name-$UNIQUE_NAME
+trace export MODEL_DESCRIPTION=original-description-$UNIQUE_NAME
+envsubst < templates/one_tag.tmpl > databricks.yml
+title "create a model"
+trace $CLI bundle plan
+trace $CLI bundle deploy
+trace $CLI model-registry get-model $MODEL_NAME | jq '.registered_model_databricks | {name, description, tags}'
+
+trace export MODEL_DESCRIPTION=new-description-$UNIQUE_NAME
+envsubst < templates/one_tag.tmpl > databricks.yml
+title "update the description, this should update the description remotely as well"
+trace $CLI bundle plan
+trace $CLI bundle deploy
+trace $CLI model-registry get-model $MODEL_NAME | jq '.registered_model_databricks | {name, description, tags}'
+
+trace export MODEL_NAME=new-name-$UNIQUE_NAME
+envsubst < templates/one_tag.tmpl > databricks.yml
+title "update the name, this should recreate the model with the new name"
+trace $CLI bundle plan
+trace $CLI bundle deploy
+trace $CLI model-registry get-model $MODEL_NAME | jq '.registered_model_databricks | {name, description, tags}'
+
+title "add a new tag, this should be a no-op"
+envsubst < templates/two_tag.tmpl > databricks.yml
+trace $CLI bundle plan
+trace $CLI bundle deploy
+trace $CLI model-registry get-model $MODEL_NAME | jq '.registered_model_databricks | {name, description, tags}'

acceptance/bundle/deploy/models/basic/templates/one_tag.tmpl

@@ -0,0 +1,11 @@
+bundle:
+  name: deploy-models-basic-$UNIQUE_NAME
+
+resources:
+  models:
+    my_model:
+      name: $MODEL_NAME
+      description: $MODEL_DESCRIPTION
+      tags:
+        - key: "key1"
+          value: "value1"

acceptance/bundle/deploy/models/basic/templates/two_tag.tmpl

@@ -0,0 +1,13 @@
+bundle:
+  name: deploy-models-basic-$UNIQUE_NAME
+
+resources:
+  models:
+    my_model:
+      name: $MODEL_NAME
+      description: $MODEL_DESCRIPTION
+      tags:
+        - key: "key1"
+          value: "value1"
+        - key: "key2"
+          value: "value2"

acceptance/bundle/deploy/models/basic/test.toml

@@ -0,0 +1,2 @@
+Cloud = true
+Local = true

acceptance/bundle/refschema/out.fields.txt

@@ -1919,6 +1919,45 @@ resources.jobs.*.webhook_notifications.on_streaming_backlog_exceeded[*].id	strin
 resources.jobs.*.webhook_notifications.on_success	[]jobs.Webhook	INPUT	STATE
 resources.jobs.*.webhook_notifications.on_success[*]	jobs.Webhook	INPUT	STATE
 resources.jobs.*.webhook_notifications.on_success[*].id	string	INPUT	STATE
+resources.models.*.creation_timestamp	int64	REMOTE
+resources.models.*.description	string	ALL
+resources.models.*.id	string	INPUT	REMOTE
+resources.models.*.last_updated_timestamp	int64	REMOTE
+resources.models.*.latest_versions	[]ml.ModelVersion	REMOTE
+resources.models.*.latest_versions[*]	ml.ModelVersion	REMOTE
+resources.models.*.latest_versions[*].creation_timestamp	int64	REMOTE
+resources.models.*.latest_versions[*].current_stage	string	REMOTE
+resources.models.*.latest_versions[*].description	string	REMOTE
+resources.models.*.latest_versions[*].last_updated_timestamp	int64	REMOTE
+resources.models.*.latest_versions[*].name	string	REMOTE
+resources.models.*.latest_versions[*].run_id	string	REMOTE
+resources.models.*.latest_versions[*].run_link	string	REMOTE
+resources.models.*.latest_versions[*].source	string	REMOTE
+resources.models.*.latest_versions[*].status	ml.ModelVersionStatus	REMOTE
+resources.models.*.latest_versions[*].status_message	string	REMOTE
+resources.models.*.latest_versions[*].tags	[]ml.ModelVersionTag	REMOTE
+resources.models.*.latest_versions[*].tags[*]	ml.ModelVersionTag	REMOTE
+resources.models.*.latest_versions[*].tags[*].key	string	REMOTE
+resources.models.*.latest_versions[*].tags[*].value	string	REMOTE
+resources.models.*.latest_versions[*].user_id	string	REMOTE
+resources.models.*.latest_versions[*].version	string	REMOTE
+resources.models.*.lifecycle	resources.Lifecycle	INPUT
+resources.models.*.lifecycle.prevent_destroy	bool	INPUT
+resources.models.*.modified_status	string	INPUT
+resources.models.*.name	string	ALL
+resources.models.*.permission_level	ml.PermissionLevel	REMOTE
+resources.models.*.permissions	[]resources.MlflowModelPermission	INPUT
+resources.models.*.permissions[*]	resources.MlflowModelPermission	INPUT
+resources.models.*.permissions[*].group_name	string	INPUT
+resources.models.*.permissions[*].level	resources.MlflowModelPermissionLevel	INPUT
+resources.models.*.permissions[*].service_principal_name	string	INPUT
+resources.models.*.permissions[*].user_name	string	INPUT
+resources.models.*.tags	[]ml.ModelTag	ALL
+resources.models.*.tags[*]	ml.ModelTag	ALL
+resources.models.*.tags[*].key	string	ALL
+resources.models.*.tags[*].value	string	ALL
+resources.models.*.url	string	INPUT
+resources.models.*.user_id	string	REMOTE
 resources.pipelines.*.allow_duplicate_names	bool	INPUT	STATE
 resources.pipelines.*.budget_policy_id	string	INPUT	STATE
 resources.pipelines.*.catalog	string	INPUT	STATE

bundle/direct/dresources/all.go

@@ -11,6 +11,7 @@ var SupportedResources = map[string]any{
 	"pipelines":              (*ResourcePipeline)(nil),
 	"schemas":                (*ResourceSchema)(nil),
 	"volumes":                (*ResourceVolume)(nil),
+	"models":                 (*ResourceMlflowModel)(nil),
 	"apps":                   (*ResourceApp)(nil),
 	"sql_warehouses":         (*ResourceSqlWarehouse)(nil),
 	"database_instances":     (*ResourceDatabaseInstance)(nil),

bundle/direct/dresources/all_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/apps"
 	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/databricks-sdk-go/service/database"
+	"github.com/databricks/databricks-sdk-go/service/ml"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -60,6 +61,18 @@ var testConfig map[string]any = map[string]any{
 			Name: "main.myschema.my_synced_table",
 		},
 	},
+	"models": &resources.MlflowModel{
+		CreateModelRequest: ml.CreateModelRequest{
+			Name:        "my_mlflow_model",
+			Description: "my_mlflow_model_description",
+			Tags: []ml.ModelTag{
+				{
+					Key:   "k1",
+					Value: "v1",
+				},
+			},
+		},
+	},
 }
 
 type prepareWorkspace func(client *databricks.WorkspaceClient) error

bundle/direct/dresources/model.go

@@ -0,0 +1,127 @@
+package dresources
+
+import (
+	"context"
+
+	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/cli/bundle/deployplan"
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/ml"
+)
+
+type ResourceMlflowModel struct {
+	client *databricks.WorkspaceClient
+}
+
+func (*ResourceMlflowModel) New(client *databricks.WorkspaceClient) *ResourceMlflowModel {
+	return &ResourceMlflowModel{client: client}
+}
+
+func (*ResourceMlflowModel) PrepareState(input *resources.MlflowModel) *ml.CreateModelRequest {
+	return &input.CreateModelRequest
+}
+
+func (*ResourceMlflowModel) RemapState(model *ml.ModelDatabricks) *ml.CreateModelRequest {
+	return &ml.CreateModelRequest{
+		Name:            model.Name,
+		Tags:            model.Tags,
+		Description:     model.Description,
+		ForceSendFields: filterFields[ml.CreateModelRequest](model.ForceSendFields),
+	}
+}
+
+func (r *ResourceMlflowModel) DoRefresh(ctx context.Context, id string) (*ml.ModelDatabricks, error) {
+	response, err := r.client.ModelRegistry.GetModel(ctx, ml.GetModelRequest{
+		Name: id,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return response.RegisteredModelDatabricks, nil
+}
+
+func (r *ResourceMlflowModel) DoCreate(ctx context.Context, config *ml.CreateModelRequest) (string, *ml.ModelDatabricks, error) {
+	response, err := r.client.ModelRegistry.CreateModel(ctx, *config)
+	if err != nil {
+		return "", nil, err
+	}
+	// Create API call returns [ml.Model] while DoRefresh returns [ml.ModelDatabricks].
+	// Thus we need to convert the response to the expected type.
+	modelDatabricks := &ml.ModelDatabricks{
+		Name:            response.RegisteredModel.Name,
+		Description:     response.RegisteredModel.Description,
+		Tags:            response.RegisteredModel.Tags,
+		ForceSendFields: filterFields[ml.ModelDatabricks](response.RegisteredModel.ForceSendFields, "CreationTimestamp", "Id", "LastUpdatedTimestamp", "LatestVersions", "PermissionLevel", "UserId"),
+
+		// Coping the fields only to satisfy the linter. These fields are not
+		// part of the configuration tree so they don't need to be copied.
+		// The linter works as a safeguard to ensure we add new fields to the bundle config tree
+		// to the mapping logic here as well.
+		CreationTimestamp:    0,
+		Id:                   "",
+		LastUpdatedTimestamp: 0,
+		LatestVersions:       nil,
+		PermissionLevel:      "",
+		UserId:               "",
+	}
+	return response.RegisteredModel.Name, modelDatabricks, nil
+}
+
+func (r *ResourceMlflowModel) DoUpdate(ctx context.Context, id string, config *ml.CreateModelRequest) (*ml.ModelDatabricks, error) {
+	updateRequest := ml.UpdateModelRequest{
+		Name:            id,
+		Description:     config.Description,
+		ForceSendFields: filterFields[ml.UpdateModelRequest](config.ForceSendFields),
+	}
+
+	response, err := r.client.ModelRegistry.UpdateModel(ctx, updateRequest)
+	if err != nil {
+		return nil, err
+	}
+
+	// Update API call returns [ml.Model] while DoRefresh returns [ml.ModelDatabricks].
+	// Thus we need to convert the response to the expected type.
+	modelDatabricks := &ml.ModelDatabricks{
+		Name:            response.RegisteredModel.Name,
+		Description:     response.RegisteredModel.Description,
+		Tags:            response.RegisteredModel.Tags,
+		ForceSendFields: filterFields[ml.ModelDatabricks](response.RegisteredModel.ForceSendFields, "CreationTimestamp", "Id", "LastUpdatedTimestamp", "LatestVersions", "PermissionLevel", "UserId"),
+
+		// Coping the fields only to satisfy the linter. These fields are not
+		// part of the configuration tree so they don't need to be copied.
+		// The linter works as a safeguard to ensure we add new fields to the bundle config tree
+		// to the mapping logic here as well.
+		CreationTimestamp:    0,
+		Id:                   "",
+		LastUpdatedTimestamp: 0,
+		LatestVersions:       nil,
+		PermissionLevel:      "",
+		UserId:               "",
+	}
+	return modelDatabricks, nil
+}
+
+func (r *ResourceMlflowModel) DoDelete(ctx context.Context, id string) error {
+	return r.client.ModelRegistry.DeleteModel(ctx, ml.DeleteModelRequest{
+		Name: id,
+	})
+}
+
+func (*ResourceMlflowModel) FieldTriggers() map[string]deployplan.ActionType {
+	return map[string]deployplan.ActionType{
+		// Recreate matches current behavior of Terraform. It is possible to rename without recreate
+		// but that would require dynamic select of the method during update since
+		// the [ml.RenameModel] needs to be called instead of [ml.UpdateModel].
+		//
+		// We might reasonably choose to never fix this because this is a legacy resource.
+		"name": deployplan.ActionTypeRecreate,
+
+		// Allowing updates for tags requires dynamic selection of the method since
+		// tags can only be updated by calling [ml.SetModelTag] or [ml.DeleteModelTag] methods.
+		//
+		// Skip annotation matches the current behavior of Terraform where tags changes are showed
+		// in plan but are just ignored / not applied. Since this is a legacy resource we might
+		// reasonably choose to not fix it here as well.
+		"tags": deployplan.ActionTypeSkip,
+	}
+}