From d7fbf8ab69afc02f7050bde90fd11efc6dfe4895 Mon Sep 17 00:00:00 2001
From: Ilya Kuznetsov <ilya.kuznetsov@databricks.com>
Date: Mon, 23 Dec 2024 12:14:33 +0100
Subject: [PATCH] fix: Handle ${workspace.file_path references} in
 source-linked deployment

---
 bundle/config/mutator/apply_presets.go        |  2 +
 .../mutator/resolve_variable_references.go    | 39 +++++++++++++
 .../resolve_variable_references_test.go       | 55 +++++++++++++++++++
 3 files changed, 96 insertions(+)

diff --git a/bundle/config/mutator/apply_presets.go b/bundle/config/mutator/apply_presets.go
index 381703756..0bc712a7c 100644
--- a/bundle/config/mutator/apply_presets.go
+++ b/bundle/config/mutator/apply_presets.go
@@ -241,6 +241,8 @@ func (m *applyPresets) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnos
 			disabled := false
 			b.Config.Presets.SourceLinkedDeployment = &disabled
 		}
+		enabled := true
+		b.Config.Presets.SourceLinkedDeployment = &enabled
 	}
 
 	return diags
diff --git a/bundle/config/mutator/resolve_variable_references.go b/bundle/config/mutator/resolve_variable_references.go
index 8c207e375..7a600e8ac 100644
--- a/bundle/config/mutator/resolve_variable_references.go
+++ b/bundle/config/mutator/resolve_variable_references.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/variable"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/dyn"
@@ -125,6 +126,19 @@ func (m *resolveVariableReferences) Apply(ctx context.Context, b *bundle.Bundle)
 	varPath := dyn.NewPath(dyn.Key("var"))
 
 	var diags diag.Diagnostics
+
+	if config.IsExplicitlyEnabled(b.Config.Presets.SourceLinkedDeployment) {
+		revertWorkspacePath, err := updateWorkspacePathForSourceLinkedDeployment(b)
+		if err != nil {
+			diags = diags.Extend(diag.FromErr(err))
+		} else {
+			defer func() {
+				err := revertWorkspacePath()
+				diags = diags.Extend(diag.FromErr(err))
+			}()
+		}
+	}
+
 	err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) {
 		// Synthesize a copy of the root that has all fields that are present in the type
 		// but not set in the dynamic value set to their corresponding empty value.
@@ -189,3 +203,28 @@ func (m *resolveVariableReferences) Apply(ctx context.Context, b *bundle.Bundle)
 	}
 	return diags
 }
+
+type cleanup func() error
+
+func updateWorkspacePathForSourceLinkedDeployment(b *bundle.Bundle) (cleanup, error) {
+	originalWorkspacePath := dyn.NilValue
+	fieldName := "workspace.file_path"
+
+	// If source-linked deployment is enabled, update the workspace path to the sync root path so variables will have correct path.
+	err := b.Config.Mutate(func(v dyn.Value) (dyn.Value, error) {
+		path, err := dyn.Get(v, fieldName)
+		if err != nil {
+			return dyn.InvalidValue, err
+		}
+		originalWorkspacePath = path
+		return dyn.Set(v, fieldName, dyn.V(b.SyncRootPath))
+	})
+	if err != nil {
+		return nil, err
+	}
+	return func() error {
+		return b.Config.Mutate(func(v dyn.Value) (dyn.Value, error) {
+			return dyn.Set(v, fieldName, originalWorkspacePath)
+		})
+	}, nil
+}
diff --git a/bundle/config/mutator/resolve_variable_references_test.go b/bundle/config/mutator/resolve_variable_references_test.go
index 7bb6f11a0..833fe7419 100644
--- a/bundle/config/mutator/resolve_variable_references_test.go
+++ b/bundle/config/mutator/resolve_variable_references_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/databricks/databricks-sdk-go/service/pipelines"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -434,3 +435,57 @@ func TestResolveComplexVariableWithVarReference(t *testing.T) {
 	require.NoError(t, diags.Error())
 	require.Equal(t, "cicd_template==1.0.0", b.Config.Resources.Jobs["job1"].JobSettings.Tasks[0].Libraries[0].Pypi.Package)
 }
+
+func TestResolveVariableReferencesWithSourceLinkedDeploymentDisabled(t *testing.T) {
+	testCases := []struct {
+		enabled bool
+		assert  func(t *testing.T, b *bundle.Bundle)
+	}{
+		{
+			true,
+			func(t *testing.T, b *bundle.Bundle) {
+				// Variables that use workspace file path should have SyncRootValue during resolution phase
+				require.Equal(t, "sync/root/path", b.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Configuration["source"])
+
+				// The file path itself should remain the same
+				require.Equal(t, "file/path", b.Config.Workspace.FilePath)
+			},
+		},
+		{
+			false,
+			func(t *testing.T, b *bundle.Bundle) {
+				require.Equal(t, "file/path", b.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Configuration["source"])
+				require.Equal(t, "file/path", b.Config.Workspace.FilePath)
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		b := &bundle.Bundle{
+			SyncRootPath: "sync/root/path",
+			Config: config.Root{
+				Presets: config.Presets{
+					SourceLinkedDeployment: &testCase.enabled,
+				},
+				Workspace: config.Workspace{
+					FilePath: "file/path",
+				},
+				Resources: config.Resources{
+					Pipelines: map[string]*resources.Pipeline{
+						"pipeline1": {
+							PipelineSpec: &pipelines.PipelineSpec{
+								Configuration: map[string]string{
+									"source": "${workspace.file_path}",
+								},
+							},
+						},
+					},
+				},
+			},
+		}
+
+		diags := bundle.Apply(context.Background(), b, ResolveVariableReferences("workspace"))
+		require.NoError(t, diags.Error())
+		testCase.assert(t, b)
+	}
+}