diff --git a/bundle/config/mutator/apply_source_linked_deployment_preset_test.go b/bundle/config/mutator/apply_source_linked_deployment_preset_test.go index 42fda8ea7..0dbe3a441 100644 --- a/bundle/config/mutator/apply_source_linked_deployment_preset_test.go +++ b/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": {}, diff --git a/bundle/config/mutator/configure_wsfs_test.go b/bundle/config/mutator/configure_wsfs_test.go index 6f76293e0..6762a446b 100644 --- a/bundle/config/mutator/configure_wsfs_test.go +++ b/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) diff --git a/bundle/tests/loader.go b/bundle/tests/loader.go index 6748e6409..bd5ba4cba 100644 --- a/bundle/tests/loader.go +++ b/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)) diff --git a/integration/libs/git/git_fetch_test.go b/integration/libs/git/git_fetch_test.go index 0998d775b..73e33de73 100644 --- a/integration/libs/git/git_fetch_test.go +++ 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 diff --git a/libs/dbr/context.go b/libs/dbr/context.go index 7512c0fe2..303e911fb 100644 --- a/libs/dbr/context.go +++ b/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 } diff --git a/libs/dbr/context_test.go b/libs/dbr/context_test.go index fc53cf130..94e155ab5 100644 --- a/libs/dbr/context_test.go +++ b/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{}))) } diff --git a/libs/dbr/detect.go b/libs/dbr/detect.go index d8b4dfe20..a66f2fb77 100644 --- a/libs/dbr/detect.go +++ b/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, + } } diff --git a/libs/dbr/detect_test.go b/libs/dbr/detect_test.go index 3a4a43a73..9716b50f6 100644 --- a/libs/dbr/detect_test.go +++ b/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)) }) } diff --git a/libs/template/renderer_test.go b/libs/template/renderer_test.go index f9306ca31..41439846b 100644 --- a/libs/template/renderer_test.go +++ b/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) diff --git a/libs/template/writer_test.go b/libs/template/writer_test.go index 9d57966ee..a3e10d4fe 100644 --- a/libs/template/writer_test.go +++ b/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"}, })