From 356a1c75c118cdcce7d9df62e41074c42d96fdb7 Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Tue, 24 Sep 2024 14:26:20 +0200
Subject: [PATCH] added more patterns to replace

---
 .../mutator/rewrite_workspace_prefix.go       | 21 ++++++++++++-------
 .../mutator/rewrite_workspace_prefix_test.go  |  4 ++--
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/bundle/config/mutator/rewrite_workspace_prefix.go b/bundle/config/mutator/rewrite_workspace_prefix.go
index 8f17ba816e..b739f86279 100644
--- a/bundle/config/mutator/rewrite_workspace_prefix.go
+++ b/bundle/config/mutator/rewrite_workspace_prefix.go
@@ -24,11 +24,15 @@ func (m *rewriteWorkspacePrefix) Name() string {
 
 func (m *rewriteWorkspacePrefix) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
 	diags := diag.Diagnostics{}
-	paths := []string{
-		"/Workspace/${workspace.root_path}",
-		"/Workspace/${workspace.file_path}",
-		"/Workspace/${workspace.artifact_path}",
-		"/Workspace/${workspace.state_path}",
+	paths := map[string]string{
+		"/Workspace/${workspace.root_path}":     "${workspace.root_path}",
+		"/Workspace${workspace.root_path}":      "${workspace.root_path}",
+		"/Workspace/${workspace.file_path}":     "${workspace.file_path}",
+		"/Workspace${workspace.file_path}":      "${workspace.file_path}",
+		"/Workspace/${workspace.artifact_path}": "${workspace.artifact_path}",
+		"/Workspace${workspace.artifact_path}":  "${workspace.artifact_path}",
+		"/Workspace/${workspace.state_path}":    "${workspace.state_path}",
+		"/Workspace${workspace.state_path}":     "${workspace.state_path}",
 	}
 
 	err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) {
@@ -40,17 +44,18 @@ func (m *rewriteWorkspacePrefix) Apply(ctx context.Context, b *bundle.Bundle) di
 				return v, nil
 			}
 
-			for _, path := range paths {
+			for path, replacePath := range paths {
 				if strings.Contains(vv, path) {
+					newPath := strings.Replace(vv, path, replacePath, 1)
 					diags = append(diags, diag.Diagnostic{
 						Severity:  diag.Warning,
-						Summary:   fmt.Sprintf("substring %q found in %q. Please update this to %q. For more information, please refer to: https://docs.databricks.com/en/release-notes/dev-tools/bundles.html#workspace-paths", path, vv, strings.Replace(vv, "/Workspace/", "", 1)),
+						Summary:   fmt.Sprintf("substring %q found in %q. Please update this to %q. For more information, please refer to: https://docs.databricks.com/en/release-notes/dev-tools/bundles.html#workspace-paths", path, vv, newPath),
 						Locations: v.Locations(),
 						Paths:     []dyn.Path{p},
 					})
 
 					// Remove the workspace prefix from the string.
-					return dyn.NewValue(strings.Replace(vv, "/Workspace/", "", 1), v.Locations()), nil
+					return dyn.NewValue(newPath, v.Locations()), nil
 				}
 			}
 
diff --git a/bundle/config/mutator/rewrite_workspace_prefix_test.go b/bundle/config/mutator/rewrite_workspace_prefix_test.go
index b58d6c1b24..847e89f77c 100644
--- a/bundle/config/mutator/rewrite_workspace_prefix_test.go
+++ b/bundle/config/mutator/rewrite_workspace_prefix_test.go
@@ -35,7 +35,7 @@ func TestNoWorkspacePrefixUsed(t *testing.T) {
 								},
 								{
 									NotebookTask: &jobs.NotebookTask{
-										NotebookPath: "/Workspace/${workspace.file_path}/notebook1",
+										NotebookPath: "/Workspace${workspace.file_path}/notebook1",
 									},
 									Libraries: []compute.Library{
 										{
@@ -66,7 +66,7 @@ func TestNoWorkspacePrefixUsed(t *testing.T) {
 
 	expectedErrors := map[string]bool{
 		`substring "/Workspace/${workspace.root_path}" found in "/Workspace/${workspace.root_path}/file1.py". Please update this to "${workspace.root_path}/file1.py". For more information, please refer to: https://docs.databricks.com/en/release-notes/dev-tools/bundles.html#workspace-paths`:             true,
-		`substring "/Workspace/${workspace.file_path}" found in "/Workspace/${workspace.file_path}/notebook1". Please update this to "${workspace.file_path}/notebook1". For more information, please refer to: https://docs.databricks.com/en/release-notes/dev-tools/bundles.html#workspace-paths`:           true,
+		`substring "/Workspace${workspace.file_path}" found in "/Workspace${workspace.file_path}/notebook1". Please update this to "${workspace.file_path}/notebook1". For more information, please refer to: https://docs.databricks.com/en/release-notes/dev-tools/bundles.html#workspace-paths`:             true,
 		`substring "/Workspace/${workspace.artifact_path}" found in "/Workspace/${workspace.artifact_path}/jar1.jar". Please update this to "${workspace.artifact_path}/jar1.jar". For more information, please refer to: https://docs.databricks.com/en/release-notes/dev-tools/bundles.html#workspace-paths`: true,
 	}