From 90d6490106820c1c16c0e6e6abbeb3bfd4c63062 Mon Sep 17 00:00:00 2001
From: Fabian Jakobs <fabian.jakobs@databricks.com>
Date: Wed, 4 Sep 2024 11:40:15 +0200
Subject: [PATCH] Centralize code to detect if we are running on Databricks

---
 bundle/config/mutator/configure_wsfs.go |  6 ++----
 libs/runtime/detect.go                  | 14 ++++++++++++++
 libs/template/file.go                   |  9 ++-------
 3 files changed, 18 insertions(+), 11 deletions(-)
 create mode 100644 libs/runtime/detect.go

diff --git a/bundle/config/mutator/configure_wsfs.go b/bundle/config/mutator/configure_wsfs.go
index 1d1bec5824..296536d12f 100644
--- a/bundle/config/mutator/configure_wsfs.go
+++ b/bundle/config/mutator/configure_wsfs.go
@@ -6,13 +6,11 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/libs/diag"
-	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/cli/libs/runtime"
 	"github.com/databricks/cli/libs/vfs"
 )
 
-const envDatabricksRuntimeVersion = "DATABRICKS_RUNTIME_VERSION"
-
 type configureWSFS struct{}
 
 func ConfigureWSFS() bundle.Mutator {
@@ -32,7 +30,7 @@ func (m *configureWSFS) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagno
 	}
 
 	// The executable must be running on DBR.
-	if _, ok := env.Lookup(ctx, envDatabricksRuntimeVersion); !ok {
+	if !runtime.RunsOnDatabricks(ctx) {
 		return nil
 	}
 
diff --git a/libs/runtime/detect.go b/libs/runtime/detect.go
new file mode 100644
index 0000000000..cfee420ed1
--- /dev/null
+++ b/libs/runtime/detect.go
@@ -0,0 +1,14 @@
+package runtime
+
+import (
+	"context"
+
+	"github.com/databricks/cli/libs/env"
+)
+
+const envDatabricksRuntimeVersion = "DATABRICKS_RUNTIME_VERSION"
+
+func RunsOnDatabricks(ctx context.Context) bool {
+	_, ok := env.Lookup(ctx, envDatabricksRuntimeVersion)
+	return ok
+}
diff --git a/libs/template/file.go b/libs/template/file.go
index 8553d15f0d..c27857c9ca 100644
--- a/libs/template/file.go
+++ b/libs/template/file.go
@@ -10,8 +10,8 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/cmd/root"
-	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/cli/libs/runtime"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 )
 
@@ -107,13 +107,8 @@ func (f *inMemoryFile) PersistToDisk() error {
 	return writeFile(f.ctx, path, f.content, f.perm)
 }
 
-func runsOnDatabricks(ctx context.Context) bool {
-	_, ok := env.Lookup(ctx, "DATABRICKS_RUNTIME_VERSION")
-	return ok
-}
-
 func shouldUseImportNotebook(ctx context.Context, path string) bool {
-	return strings.HasPrefix(path, "/Workspace/") && runsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb")
+	return strings.HasPrefix(path, "/Workspace/") && runtime.RunsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb")
 }
 
 func writeFile(ctx context.Context, path string, content []byte, perm fs.FileMode) error {