Store DBR version in context (#2438)

## Changes
This PR also starts storing the DBR versions in the context.

Go patch file used:
```
@@
var x expression
@@
-dbr.MockRuntime(x, true)
+dbr.MockRuntime(x, dbr.Environment{IsDbr: true, Version: "15.4"})

@@
var x expression
@@
-dbr.MockRuntime(x, false)
+dbr.MockRuntime(x, dbr.Environment{})
```

ref: https://github.com/uber-go/gopatch

## Why
This localised all DBR version accesses to `libs/dbr`. Relevant comment:
https://github.com/databricks/cli/pull/2432#discussion_r1982878616

## Tests
Exiting tests are modified.

bundle/config/mutator/apply_source_linked_deployment_preset_test.go

@@ -36,13 +36,13 @@ func TestApplyPresetsSourceLinkedDeployment(t *testing.T) {
 	}{
 		{
 			name:          "preset enabled, bundle in Workspace, databricks runtime",
-			ctx:           dbr.MockRuntime(testContext, true),
+			ctx:           dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}),
 			initialValue:  &enabled,
 			expectedValue: &enabled,
 		},
 		{
 			name: "preset enabled, bundle not in Workspace, databricks runtime",
-			ctx:  dbr.MockRuntime(testContext, true),
+			ctx:  dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}),
 			mutateBundle: func(b *bundle.Bundle) {
 				b.SyncRootPath = "/Users/user.name@company.com"
 			},
@@ -52,26 +52,26 @@ func TestApplyPresetsSourceLinkedDeployment(t *testing.T) {
 		},
 		{
 			name:            "preset enabled, bundle in Workspace, not databricks runtime",
-			ctx:             dbr.MockRuntime(testContext, false),
+			ctx:             dbr.MockRuntime(testContext, dbr.Environment{}),
 			initialValue:    &enabled,
 			expectedValue:   &disabled,
 			expectedWarning: "source-linked deployment is available only in the Databricks Workspace",
 		},
 		{
 			name:          "preset disabled, bundle in Workspace, databricks runtime",
-			ctx:           dbr.MockRuntime(testContext, true),
+			ctx:           dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}),
 			initialValue:  &disabled,
 			expectedValue: &disabled,
 		},
 		{
 			name:          "preset nil, bundle in Workspace, databricks runtime",
-			ctx:           dbr.MockRuntime(testContext, true),
+			ctx:           dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}),
 			initialValue:  nil,
 			expectedValue: nil,
 		},
 		{
 			name: "preset nil, dev mode true, bundle in Workspace, databricks runtime",
-			ctx:  dbr.MockRuntime(testContext, true),
+			ctx:  dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}),
 			mutateBundle: func(b *bundle.Bundle) {
 				b.Config.Bundle.Mode = config.Development
 			},
@@ -80,7 +80,7 @@ func TestApplyPresetsSourceLinkedDeployment(t *testing.T) {
 		},
 		{
 			name: "preset enabled, workspace.file_path is defined by user",
-			ctx:  dbr.MockRuntime(testContext, true),
+			ctx:  dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}),
 			mutateBundle: func(b *bundle.Bundle) {
 				b.Config.Workspace.FilePath = "file_path"
 			},
@@ -90,7 +90,7 @@ func TestApplyPresetsSourceLinkedDeployment(t *testing.T) {
 		},
 		{
 			name: "preset enabled, apps is defined by user",
-			ctx:  dbr.MockRuntime(testContext, true),
+			ctx:  dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}),
 			mutateBundle: func(b *bundle.Bundle) {
 				b.Config.Resources.Apps = map[string]*resources.App{
 					"app": {},

bundle/config/mutator/configure_wsfs_test.go

@@ -47,7 +47,7 @@ func TestConfigureWSFS_SkipsIfNotRunningOnRuntime(t *testing.T) {
 	originalSyncRoot := b.SyncRoot
 
 	ctx := context.Background()
-	ctx = dbr.MockRuntime(ctx, false)
+	ctx = dbr.MockRuntime(ctx, dbr.Environment{})
 	diags := bundle.Apply(ctx, b, mutator.ConfigureWSFS())
 	assert.Empty(t, diags)
 	assert.Equal(t, originalSyncRoot, b.SyncRoot)
@@ -58,7 +58,7 @@ func TestConfigureWSFS_SwapSyncRoot(t *testing.T) {
 	originalSyncRoot := b.SyncRoot
 
 	ctx := context.Background()
-	ctx = dbr.MockRuntime(ctx, true)
+	ctx = dbr.MockRuntime(ctx, dbr.Environment{IsDbr: true, Version: "15.4"})
 	diags := bundle.Apply(ctx, b, mutator.ConfigureWSFS())
 	assert.Empty(t, diags)
 	assert.NotEqual(t, originalSyncRoot, b.SyncRoot)

bundle/tests/loader.go

@@ -69,7 +69,7 @@ func initializeTarget(t *testing.T, path, env string) (*bundle.Bundle, diag.Diag
 	b := load(t, path)
 	configureMock(t, b)
 
-	ctx := dbr.MockRuntime(context.Background(), false)
+	ctx := dbr.MockRuntime(context.Background(), dbr.Environment{})
 	diags := bundle.Apply(ctx, b, mutator.SelectTarget(env))
 	diags = diags.Extend(phases.Initialize(ctx, b))
 

integration/libs/git/git_fetch_test.go

@@ -50,7 +50,7 @@ func TestFetchRepositoryInfoAPI_FromRepo(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 	targetPath := ensureWorkspacePrefix(acc.TemporaryRepo(wt, examplesRepoUrl))
 
-	ctx = dbr.MockRuntime(ctx, true)
+	ctx = dbr.MockRuntime(ctx, dbr.Environment{IsDbr: true, Version: "15.4"})
 
 	for _, inputPath := range []string{
 		path.Join(targetPath, "knowledge_base/dashboard_nyc_taxi"),
@@ -72,7 +72,7 @@ func TestFetchRepositoryInfoAPI_FromNonRepo(t *testing.T) {
 	err := wt.W.Workspace.MkdirsByPath(ctx, path.Join(rootPath, "a/b/c"))
 	require.NoError(t, err)
 
-	ctx = dbr.MockRuntime(ctx, true)
+	ctx = dbr.MockRuntime(ctx, dbr.Environment{IsDbr: true, Version: "15.4"})
 
 	tests := []struct {
 		input string

libs/dbr/context.go

@@ -15,6 +15,11 @@ const (
 	dbrKey = key(1)
 )
 
+type Environment struct {
+	IsDbr   bool
+	Version string
+}
+
 // DetectRuntime detects whether or not the current
 // process is running inside a Databricks Runtime environment.
 // It return a new context with the detection result set.
@@ -27,11 +32,11 @@ func DetectRuntime(ctx context.Context) context.Context {
 
 // MockRuntime is a helper function to mock the detection result.
 // It returns a new context with the detection result set.
-func MockRuntime(ctx context.Context, b bool) context.Context {
+func MockRuntime(ctx context.Context, runtime Environment) context.Context {
 	if v := ctx.Value(dbrKey); v != nil {
 		panic("dbr.MockRuntime called twice on the same context")
 	}
-	return context.WithValue(ctx, dbrKey, b)
+	return context.WithValue(ctx, dbrKey, runtime)
 }
 
 // RunsOnRuntime returns the detection result from the context.
@@ -45,5 +50,14 @@ func RunsOnRuntime(ctx context.Context) bool {
 	if v == nil {
 		panic("dbr.RunsOnRuntime called without calling dbr.DetectRuntime first")
 	}
-	return v.(bool)
+	return v.(Environment).IsDbr
+}
+
+func RuntimeVersion(ctx context.Context) string {
+	v := ctx.Value(dbrKey)
+	if v == nil {
+		panic("dbr.RuntimeVersion called without calling dbr.DetectRuntime first")
+	}
+
+	return v.(Environment).Version
 }

libs/dbr/context_test.go

@@ -23,11 +23,11 @@ func TestContext_MockRuntimePanics(t *testing.T) {
 	ctx := context.Background()
 
 	// Run detection.
-	ctx = MockRuntime(ctx, true)
+	ctx = MockRuntime(ctx, Environment{IsDbr: true, Version: "15.4"})
 
 	// Expect a panic if the mock function is run twice.
 	assert.Panics(t, func() {
-		MockRuntime(ctx, true)
+		MockRuntime(ctx, Environment{IsDbr: true, Version: "15.4"})
 	})
 }
 
@@ -40,6 +40,15 @@ func TestContext_RunsOnRuntimePanics(t *testing.T) {
 	})
 }
 
+func TestContext_RuntimeVersionPanics(t *testing.T) {
+	ctx := context.Background()
+
+	// Expect a panic if the detection is not run.
+	assert.Panics(t, func() {
+		RuntimeVersion(ctx)
+	})
+}
+
 func TestContext_RunsOnRuntime(t *testing.T) {
 	ctx := context.Background()
 
@@ -52,8 +61,26 @@ func TestContext_RunsOnRuntime(t *testing.T) {
 	})
 }
 
+func TestContext_RuntimeVersion(t *testing.T) {
+	ctx := context.Background()
+
+	// Run detection.
+	ctx = DetectRuntime(ctx)
+
+	// Expect no panic because detection has run.
+	assert.NotPanics(t, func() {
+		RuntimeVersion(ctx)
+	})
+}
+
 func TestContext_RunsOnRuntimeWithMock(t *testing.T) {
 	ctx := context.Background()
-	assert.True(t, RunsOnRuntime(MockRuntime(ctx, true)))
-	assert.False(t, RunsOnRuntime(MockRuntime(ctx, false)))
+	assert.True(t, RunsOnRuntime(MockRuntime(ctx, Environment{IsDbr: true, Version: "15.4"})))
+	assert.False(t, RunsOnRuntime(MockRuntime(ctx, Environment{})))
+}
+
+func TestContext_RuntimeVersionWithMock(t *testing.T) {
+	ctx := context.Background()
+	assert.Equal(t, "15.4", RuntimeVersion(MockRuntime(ctx, Environment{IsDbr: true, Version: "15.4"})))
+	assert.Empty(t, RuntimeVersion(MockRuntime(ctx, Environment{})))
 }

libs/dbr/detect.go

@@ -13,23 +13,27 @@ var statFunc = os.Stat
 
 // detect returns true if the current process is running on a Databricks Runtime.
 // Its return value is meant to be cached in the context.
-func detect(ctx context.Context) bool {
+func detect(ctx context.Context) Environment {
 	// Databricks Runtime implies Linux.
 	// Return early on other operating systems.
 	if runtime.GOOS != "linux" {
-		return false
+		return Environment{}
 	}
 
 	// Databricks Runtime always has the DATABRICKS_RUNTIME_VERSION environment variable set.
-	if value, ok := env.Lookup(ctx, "DATABRICKS_RUNTIME_VERSION"); !ok || value == "" {
-		return false
+	version, ok := env.Lookup(ctx, "DATABRICKS_RUNTIME_VERSION")
+	if !ok || version == "" {
+		return Environment{}
 	}
 
 	// Expect to see a "/databricks" directory.
 	if fi, err := statFunc("/databricks"); err != nil || !fi.IsDir() {
-		return false
+		return Environment{}
 	}
 
 	// All checks passed.
-	return true
+	return Environment{
+		IsDbr:   true,
+		Version: version,
+	}
 }

libs/dbr/detect_test.go

@@ -35,7 +35,7 @@ func TestDetect_NotLinux(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	assert.False(t, detect(ctx))
+	assert.Equal(t, Environment{}, detect(ctx))
 }
 
 func TestDetect_Env(t *testing.T) {
@@ -46,17 +46,23 @@ func TestDetect_Env(t *testing.T) {
 
 	t.Run("empty", func(t *testing.T) {
 		ctx := env.Set(context.Background(), "DATABRICKS_RUNTIME_VERSION", "")
-		assert.False(t, detect(ctx))
+		assert.Equal(t, Environment{}, detect(ctx))
 	})
 
 	t.Run("non-empty cluster", func(t *testing.T) {
 		ctx := env.Set(context.Background(), "DATABRICKS_RUNTIME_VERSION", "15.4")
-		assert.True(t, detect(ctx))
+		assert.Equal(t, Environment{
+			IsDbr:   true,
+			Version: "15.4",
+		}, detect(ctx))
 	})
 
 	t.Run("non-empty serverless", func(t *testing.T) {
 		ctx := env.Set(context.Background(), "DATABRICKS_RUNTIME_VERSION", "client.1.13")
-		assert.True(t, detect(ctx))
+		assert.Equal(t, Environment{
+			IsDbr:   true,
+			Version: "client.1.13",
+		}, detect(ctx))
 	})
 }
 
@@ -68,16 +74,19 @@ func TestDetect_Stat(t *testing.T) {
 
 	t.Run("error", func(t *testing.T) {
 		configureStatFunc(t, nil, fs.ErrNotExist)
-		assert.False(t, detect(ctx))
+		assert.Equal(t, Environment{}, detect(ctx))
 	})
 
 	t.Run("not a directory", func(t *testing.T) {
 		configureStatFunc(t, fakefs.FileInfo{}, nil)
-		assert.False(t, detect(ctx))
+		assert.Equal(t, Environment{}, detect(ctx))
 	})
 
 	t.Run("directory", func(t *testing.T) {
 		configureStatFunc(t, fakefs.FileInfo{FakeDir: true}, nil)
-		assert.True(t, detect(ctx))
+		assert.Equal(t, Environment{
+			IsDbr:   true,
+			Version: "non-empty",
+		}, detect(ctx))
 	})
 }

libs/template/renderer_test.go

@@ -43,7 +43,7 @@ func init() {
 }
 
 func assertBuiltinTemplateValid(t *testing.T, template string, settings map[string]any, target string, isServicePrincipal, build bool, tempDir string) {
-	ctx := dbr.MockRuntime(context.Background(), false)
+	ctx := dbr.MockRuntime(context.Background(), dbr.Environment{})
 
 	templateFS, err := fs.Sub(builtinTemplates, path.Join("templates", template))
 	require.NoError(t, err)

libs/template/writer_test.go

@@ -31,7 +31,7 @@ func TestDefaultWriterConfigureOnDBR(t *testing.T) {
 		t.Skip("Skipping test on Windows")
 	}
 
-	ctx := dbr.MockRuntime(context.Background(), true)
+	ctx := dbr.MockRuntime(context.Background(), dbr.Environment{IsDbr: true, Version: "15.4"})
 	ctx = root.SetWorkspaceClient(ctx, &databricks.WorkspaceClient{
 		Config: &workspaceConfig.Config{Host: "https://myhost.com"},
 	})