[draft] Fix bundle/direct/dresources/experiment.go

acceptance/bundle/resources/experiments/basic/out.plan.txt

@@ -1 +1,3 @@
-Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged
+update experiments.my_experiment
+
+Plan: 0 to add, 1 to change, 0 to delete, 0 unchanged

acceptance/bundle/resources/experiments/basic/output.txt

@@ -133,7 +133,7 @@ Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged
   ]
 }
 
-=== adding a new tag should be a no-op
+=== adding a new tag updates the existing experiment
 >>> envsubst
 
 >>> [CLI] bundle plan
@@ -168,6 +168,10 @@ Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged
     {
       "key": "mlflow.experimentType",
       "value": "MLFLOW_EXPERIMENT"
+    },
+    {
+      "key": "k2",
+      "value": "v2"
     }
   ]
 }

acceptance/bundle/resources/experiments/basic/script

@@ -29,9 +29,11 @@ trace $CLI bundle deploy -qq
 experiment_id=$(read_id.py my_experiment)
 trace $CLI experiments get-experiment $experiment_id | jq '.experiment | {name, artifact_location, tags}'
 
-title "adding a new tag should be a no-op"
+title "adding a new tag updates the existing experiment"
+existing_experiment_id=$experiment_id
 trace envsubst < templates/two_tag.tmpl > databricks.yml
 trace $CLI bundle plan > out.plan.txt
 trace $CLI bundle deploy -qq
 experiment_id=$(read_id.py my_experiment)
+test "$experiment_id" = "$existing_experiment_id"
 trace $CLI experiments get-experiment $experiment_id | jq '.experiment | {name, artifact_location, tags}'

bundle/direct/dresources/configs/experiments.yml

@@ -3,9 +3,6 @@ recreate_on_changes:
   - field: artifact_location
     reason: immutable
 ignore_remote_changes:
-  # Tags updates are not supported by TF. This mirrors that behaviour.
+  # The backend adds tags that the bundle does not manage.
   - field: tags
-    reason: terraform_compat
-ignore_local_changes:
-  - field: tags
-    reason: terraform_compat
+    reason: backend_generated

bundle/direct/dresources/experiment.go

@@ -2,6 +2,7 @@ package dresources
 
 import (
 	"context"
+	"fmt"
 	"strings"
 
 	"github.com/databricks/cli/bundle/config/resources"
@@ -60,14 +61,38 @@ func (r *ResourceExperiment) DoCreate(ctx context.Context, config *ml.CreateExpe
 	return result.ExperimentId, nil, nil
 }
 
-func (r *ResourceExperiment) DoUpdate(ctx context.Context, id string, config *ml.CreateExperiment, _ *PlanEntry) (*ml.Experiment, error) {
-	updateReq := ml.UpdateExperiment{
-		ExperimentId:    id,
-		NewName:         config.Name,
-		ForceSendFields: utils.FilterFields[ml.UpdateExperiment](config.ForceSendFields),
+func (r *ResourceExperiment) DoUpdate(ctx context.Context, id string, config *ml.CreateExperiment, entry *PlanEntry) (*ml.Experiment, error) {
+	if entry.Changes.HasChange(structpath.MustParsePath("name")) {
+		updateReq := ml.UpdateExperiment{
+			ExperimentId:    id,
+			NewName:         config.Name,
+			ForceSendFields: utils.FilterFields[ml.UpdateExperiment](config.ForceSendFields),
+		}
+		if err := r.client.Experiments.UpdateExperiment(ctx, updateReq); err != nil {
+			return nil, err
+		}
 	}
 
-	return nil, r.client.Experiments.UpdateExperiment(ctx, updateReq)
+	if entry.Changes.HasChange(structpath.MustParsePath("tags")) {
+		remoteTags := make(map[string]string)
+		for _, tag := range entry.RemoteState.(*ml.Experiment).Tags {
+			remoteTags[tag.Key] = tag.Value
+		}
+		for _, tag := range config.Tags {
+			if remoteValue, exists := remoteTags[tag.Key]; exists && remoteValue == tag.Value {
+				continue
+			}
+			if err := r.client.Experiments.SetExperimentTag(ctx, ml.SetExperimentTag{
+				ExperimentId: id,
+				Key:          tag.Key,
+				Value:        tag.Value,
+			}); err != nil {
+				return nil, fmt.Errorf("failed to set experiment %s tag %q: %w", id, tag.Key, err)
+			}
+		}
+	}
+
+	return nil, nil
 }
 
 func (r *ResourceExperiment) DoDelete(ctx context.Context, id string, _ *ml.CreateExperiment) error {

bundle/direct/dresources/experiment_test.go

@@ -0,0 +1,55 @@
+package dresources
+
+import (
+	"testing"
+
+	"github.com/databricks/cli/bundle/deployplan"
+	"github.com/databricks/cli/libs/testserver"
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/ml"
+	"github.com/stretchr/testify/require"
+)
+
+func TestExperimentDoUpdateSetsConfiguredTags(t *testing.T) {
+	server := testserver.New(t)
+	testserver.AddDefaultHandlers(server)
+	client, err := databricks.NewWorkspaceClient(&databricks.Config{
+		Host:  server.URL,
+		Token: "testtoken",
+	})
+	require.NoError(t, err)
+
+	resource := (&ResourceExperiment{}).New(client)
+	ctx := t.Context()
+	id, _, err := resource.DoCreate(ctx, &ml.CreateExperiment{
+		Name: "experiment-tags",
+		Tags: []ml.ExperimentTag{{Key: "mlflow.workload_creation_policy_id", Value: "old-policy"}},
+	})
+	require.NoError(t, err)
+	before, err := resource.DoRead(ctx, id)
+	require.NoError(t, err)
+
+	_, err = resource.DoUpdate(ctx, id, &ml.CreateExperiment{
+		Name: "experiment-tags",
+		Tags: []ml.ExperimentTag{
+			{Key: "mlflow.workload_creation_policy_id", Value: "new-policy"},
+			{Key: "added-tag", Value: "value"},
+		},
+	}, &deployplan.PlanEntry{
+		RemoteState: before,
+		Changes: deployplan.Changes{
+			"tags": {Action: deployplan.Update},
+		},
+	})
+	require.NoError(t, err)
+
+	after, err := resource.DoRead(ctx, id)
+	require.NoError(t, err)
+	require.Equal(t, id, after.ExperimentId)
+	tags := make(map[string]string, len(after.Tags))
+	for _, tag := range after.Tags {
+		tags[tag.Key] = tag.Value
+	}
+	require.Equal(t, "new-policy", tags["mlflow.workload_creation_policy_id"])
+	require.Equal(t, "value", tags["added-tag"])
+}

libs/testserver/experiments.go

@@ -123,6 +123,39 @@ func (s *FakeWorkspace) ExperimentUpdate(req Request) Response {
 	return Response{}
 }
 
+func (s *FakeWorkspace) ExperimentSetTag(req Request) Response {
+	defer s.LockUnlock()()
+
+	var setReq ml.SetExperimentTag
+	if err := json.Unmarshal(req.Body, &setReq); err != nil {
+		return Response{
+			Body:       fmt.Sprintf("request parsing error: %s", err),
+			StatusCode: http.StatusBadRequest,
+		}
+	}
+
+	experiment, exists := s.Experiments[setReq.ExperimentId]
+	if !exists {
+		return Response{
+			StatusCode: http.StatusNotFound,
+			Body:       map[string]string{"message": fmt.Sprintf("Experiment %s not found", setReq.ExperimentId)},
+		}
+	}
+
+	for i := range experiment.Experiment.Tags {
+		if experiment.Experiment.Tags[i].Key == setReq.Key {
+			experiment.Experiment.Tags[i].Value = setReq.Value
+			s.Experiments[setReq.ExperimentId] = experiment
+			return Response{}
+		}
+	}
+	experiment.Experiment.Tags = append(experiment.Experiment.Tags, ml.ExperimentTag{
+		Key: setReq.Key, Value: setReq.Value,
+	})
+	s.Experiments[setReq.ExperimentId] = experiment
+	return Response{}
+}
+
 func (s *FakeWorkspace) ExperimentDelete(req Request) Response {
 	var deleteReq ml.DeleteExperiment
 	if err := json.Unmarshal(req.Body, &deleteReq); err != nil {

libs/testserver/handlers.go

@@ -1355,6 +1355,10 @@ func AddDefaultHandlers(server *Server) {
 		return req.Workspace.ExperimentUpdate(req)
 	})
 
+	server.Handle("POST", "/api/2.0/mlflow/experiments/set-experiment-tag", func(req Request) any {
+		return req.Workspace.ExperimentSetTag(req)
+	})
+
 	server.Handle("POST", "/api/2.0/mlflow/experiments/delete", func(req Request) any {
 		return req.Workspace.ExperimentDelete(req)
 	})