Fix missing locations in job clusters item (#4683)

## Changes
<!-- Brief summary of your changes that is easy to understand -->

Preserve correct location for sequence items, e.g.
`my_job.job_clusters[0]`

## Why

In config-remote-sync when when we attempt to update
`my_job.job_clusters[0].job_cluster_key` it produces unexpected results
because location is not preserved even though the field is in the config

## Tests
1. Unit test for the dyn helper
2. Add this case to the acceptance tests of the config-remote-sync 

<!-- If your PR needs to be included in the release notes for next
release,
add a separate entry in NEXT_CHANGELOG.md as part of your PR. -->

Author: ilyakuz-db

acceptance/bundle/config-remote-sync/job_fields/databricks.yml.tmpl

@@ -22,8 +22,18 @@ resources:
             environment_version: "3"
             dependencies:
               - ./*.whl
+      job_clusters:
+        - job_cluster_key: test_cluster
+          new_cluster:
+            spark_version: $DEFAULT_SPARK_VERSION
+            node_type_id: $NODE_TYPE_ID
+            num_workers: 1
       tasks:
         - task_key: main
+          notebook_task:
+            notebook_path: /Users/{{workspace_user_name}}/notebook
+          job_cluster_key: test_cluster
+        - task_key: inline_cluster_task
           notebook_task:
             notebook_path: /Users/{{workspace_user_name}}/notebook
           new_cluster:

acceptance/bundle/config-remote-sync/job_fields/output.txt

@@ -12,8 +12,11 @@ Detected changes in 1 resource(s):
 Resource: resources.jobs.my_job
   email_notifications.no_alert_for_skipped_runs: add
   email_notifications.on_failure: add
+  job_clusters[0].job_cluster_key: replace
   parameters: replace
   tags['team']: add
+  tasks[task_key='inline_cluster_task'].new_cluster.num_workers: replace
+  tasks[task_key='main'].job_cluster_key: replace
   trigger.pause_status: add
   trigger.periodic: remove
   trigger.table_update: add
@@ -52,9 +55,25 @@ Resource: resources.jobs.my_job
 +            - samples.nyctaxi.trips
        environments:
          - environment_key: default
-@@ -31,4 +37,6 @@
+@@ -24,5 +30,5 @@
+               - ./*.whl
+       job_clusters:
+-        - job_cluster_key: test_cluster
++        - job_cluster_key: test_cluster_renamed
+           new_cluster:
+             spark_version: [[DEFAULT_SPARK_VERSION]]
+@@ -33,5 +39,5 @@
+           notebook_task:
+             notebook_path: /Users/{{workspace_user_name}}/notebook
+-          job_cluster_key: test_cluster
++          job_cluster_key: test_cluster_renamed
+         - task_key: inline_cluster_task
+           notebook_task:
+@@ -40,5 +46,7 @@
+             spark_version: [[DEFAULT_SPARK_VERSION]]
              node_type_id: [NODE_TYPE_ID]
-             num_workers: 1
+-            num_workers: 1
++            num_workers: 2
 +      tags:
 +        team: data
 

acceptance/bundle/config-remote-sync/job_fields/script

@@ -1,6 +1,8 @@
 #!/bin/bash
 
 envsubst < databricks.yml.tmpl > databricks.yml
+add_repl.py "$DEFAULT_SPARK_VERSION" "[DEFAULT_SPARK_VERSION]"
+add_repl.py "$NODE_TYPE_ID" "[NODE_TYPE_ID]"
 
 cleanup() {
   trace $CLI bundle destroy --auto-approve
@@ -20,6 +22,16 @@ r["email_notifications"]["no_alert_for_skipped_runs"] = True
 r["parameters"].append({"name": "region", "default": "us-east-1"})
 r["trigger"] = {"pause_status": "UNPAUSED", "table_update": {"table_names": ["samples.nyctaxi.trips"]}}
 
+for cluster in r.get("job_clusters", []):
+    if cluster.get("job_cluster_key") == "test_cluster":
+        cluster["job_cluster_key"] = "test_cluster_renamed"
+for task in r.get("tasks", []):
+    if task.get("task_key") == "inline_cluster_task" and "new_cluster" in task:
+        task["new_cluster"]["num_workers"] = 2
+for task in r.get("tasks", []):
+    if task.get("job_cluster_key") == "test_cluster":
+        task["job_cluster_key"] = "test_cluster_renamed"
+
 if "tags" not in r:
     r["tags"] = {}
 r["tags"]["team"] = "data"

libs/dyn/merge/elements_by_key.go

@@ -54,7 +54,10 @@ func (e elementsByKey) doMap(_ dyn.Path, v dyn.Value, mergeFunc func(a, b dyn.Va
 	// Gather resulting elements in natural order.
 	out := make([]dyn.Value, 0, len(keys))
 	for _, key := range keys {
-		nv, err := dyn.Set(seen[key], e.key, dyn.V(key))
+		// Preserve the location from the original key value so that
+		// downstream code can tell this field was defined in the YAML.
+		keyLoc := seen[key].Get(e.key).Locations()
+		nv, err := dyn.Set(seen[key], e.key, dyn.NewValue(key, keyLoc))
 		if err != nil {
 			return dyn.InvalidValue, err
 		}

libs/dyn/merge/elements_by_key_test.go

@@ -49,6 +49,27 @@ func TestElementByKey(t *testing.T) {
 	)
 }
 
+func TestElementByKeyPreservesLocations(t *testing.T) {
+	loc := dyn.Location{File: "config.yml", Line: 10, Column: 5}
+	vin := dyn.V([]dyn.Value{
+		dyn.V(map[string]dyn.Value{
+			"key":   dyn.NewValue("foo", []dyn.Location{loc}),
+			"value": dyn.V(42),
+		}),
+	})
+
+	keyFunc := func(v dyn.Value) string {
+		return v.MustString()
+	}
+
+	vout, err := dyn.MapByPath(vin, dyn.EmptyPath, ElementsByKey("key", keyFunc))
+	require.NoError(t, err)
+
+	// Verify the key field retains its original location.
+	keyValue := vout.Index(0).Get("key")
+	assert.Equal(t, loc, keyValue.Location())
+}
+
 func TestElementByKeyWithOverride(t *testing.T) {
 	vin := dyn.V([]dyn.Value{
 		dyn.V(map[string]dyn.Value{