From ca6332a5a4325aff1be848536f45d13bd74d93b3 Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Wed, 4 Sep 2024 13:24:55 +0200
Subject: [PATCH] Fixed complex variables are not being correctly merged from
 include files (#1746)

## Changes
Fixes an `Error: no value assigned to required variable <variable>.`
when the main complex variable definition is defined in one file but
target override is defined in separate file which is included in the
main one.

## Tests
Added regression test
---
 bundle/config/root.go                         | 12 +++++++++++-
 bundle/tests/complex_variables_test.go        | 19 +++++++++++++++++++
 .../complex_multiple_files/databricks.yml     | 17 +++++++++++++++++
 .../variables/clusters.yml                    | 11 +++++++++++
 4 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 bundle/tests/variables/complex_multiple_files/databricks.yml
 create mode 100644 bundle/tests/variables/complex_multiple_files/variables/clusters.yml

diff --git a/bundle/config/root.go b/bundle/config/root.go
index 86dc3392..281c5c2a 100644
--- a/bundle/config/root.go
+++ b/bundle/config/root.go
@@ -433,10 +433,20 @@ func rewriteShorthands(v dyn.Value) (dyn.Value, error) {
 				}, variable.Locations()), nil
 
 			case dyn.KindMap, dyn.KindSequence:
+				lookup, err := dyn.Get(variable, "lookup")
+				// If lookup is set, we don't want to rewrite the variable and return it as is.
+				if err == nil && lookup.Kind() != dyn.KindInvalid {
+					return variable, nil
+				}
+
 				// Check if the original definition of variable has a type field.
+				// Type might not be found if the variable overriden in a separate file
+				// and configuration is not merged yet.
 				typeV, err := dyn.GetByPath(v, p.Append(dyn.Key("type")))
 				if err != nil {
-					return variable, nil
+					return dyn.NewValue(map[string]dyn.Value{
+						"default": variable,
+					}, variable.Locations()), nil
 				}
 
 				if typeV.MustString() == "complex" {
diff --git a/bundle/tests/complex_variables_test.go b/bundle/tests/complex_variables_test.go
index 1badea6d..d46d8d8c 100644
--- a/bundle/tests/complex_variables_test.go
+++ b/bundle/tests/complex_variables_test.go
@@ -68,3 +68,22 @@ func TestComplexVariablesOverride(t *testing.T) {
 	require.Equal(t, "", b.Config.Resources.Jobs["my_job"].JobClusters[0].NewCluster.SparkConf["spark.random"])
 	require.Equal(t, "", b.Config.Resources.Jobs["my_job"].JobClusters[0].NewCluster.PolicyId)
 }
+
+func TestComplexVariablesOverrideWithMultipleFiles(t *testing.T) {
+	b, diags := loadTargetWithDiags("variables/complex_multiple_files", "dev")
+	require.Empty(t, diags)
+
+	diags = bundle.Apply(context.Background(), b, bundle.Seq(
+		mutator.SetVariables(),
+		mutator.ResolveVariableReferencesInComplexVariables(),
+		mutator.ResolveVariableReferences(
+			"variables",
+		),
+	))
+	require.NoError(t, diags.Error())
+
+	require.Equal(t, "14.2.x-scala2.11", b.Config.Resources.Jobs["my_job"].JobClusters[0].NewCluster.SparkVersion)
+	require.Equal(t, "Standard_DS3_v2", b.Config.Resources.Jobs["my_job"].JobClusters[0].NewCluster.NodeTypeId)
+	require.Equal(t, 4, b.Config.Resources.Jobs["my_job"].JobClusters[0].NewCluster.NumWorkers)
+	require.Equal(t, "false", b.Config.Resources.Jobs["my_job"].JobClusters[0].NewCluster.SparkConf["spark.speculation"])
+}
diff --git a/bundle/tests/variables/complex_multiple_files/databricks.yml b/bundle/tests/variables/complex_multiple_files/databricks.yml
new file mode 100644
index 00000000..cb5d2439
--- /dev/null
+++ b/bundle/tests/variables/complex_multiple_files/databricks.yml
@@ -0,0 +1,17 @@
+bundle:
+  name: complex-variables-multiple-files
+
+resources:
+  jobs:
+    my_job:
+      job_clusters:
+        - job_cluster_key: key
+          new_cluster: ${var.cluster}
+
+variables:
+  cluster:
+    type: complex
+    description: "A cluster definition"
+
+include:
+  - ./variables/*.yml
diff --git a/bundle/tests/variables/complex_multiple_files/variables/clusters.yml b/bundle/tests/variables/complex_multiple_files/variables/clusters.yml
new file mode 100644
index 00000000..badd4516
--- /dev/null
+++ b/bundle/tests/variables/complex_multiple_files/variables/clusters.yml
@@ -0,0 +1,11 @@
+targets:
+  default:
+  dev:
+    variables:
+      cluster:
+        spark_version: "14.2.x-scala2.11"
+        node_type_id: "Standard_DS3_v2"
+        num_workers: 4
+        spark_conf:
+          spark.speculation: false
+          spark.databricks.delta.retentionDurationCheck.enabled: false