update template library

Author: shreyas-goenka

cmd/bundle/generate/app.go

@@ -78,7 +78,7 @@ per target environment.`,
 			return err
 		}
 
-		outputFiler, err := filer.NewOutputFiler(ctx, w, b.BundleRootPath)
+		outputFiler, err := filer.NewOutputFiler(ctx, b.BundleRootPath)
 		if err != nil {
 			return err
 		}

cmd/bundle/generate/dashboard.go

@@ -364,7 +364,7 @@ func (d *dashboard) initialize(ctx context.Context, b *bundle.Bundle) {
 	d.relativeDashboardDir = filepath.ToSlash(rel)
 
 	// Construct output filer for writing files.
-	outputFiler, err := filer.NewOutputFiler(ctx, b.WorkspaceClient(), b.BundleRootPath)
+	outputFiler, err := filer.NewOutputFiler(ctx, b.BundleRootPath)
 	if err != nil {
 		logdiag.LogError(ctx, err)
 		return

cmd/bundle/generate/job.go

@@ -108,7 +108,7 @@ After generation, you can deploy this job to other targets using:
 			return err
 		}
 
-		outputFiler, err := filer.NewOutputFiler(ctx, w, b.BundleRootPath)
+		outputFiler, err := filer.NewOutputFiler(ctx, b.BundleRootPath)
 		if err != nil {
 			return err
 		}

cmd/bundle/generate/pipeline.go

@@ -79,7 +79,7 @@ like catalogs, schemas, and compute configurations per target.`,
 			return err
 		}
 
-		outputFiler, err := filer.NewOutputFiler(ctx, w, b.BundleRootPath)
+		outputFiler, err := filer.NewOutputFiler(ctx, b.BundleRootPath)
 		if err != nil {
 			return err
 		}

libs/filer/output_filer.go

@@ -5,8 +5,8 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/dbr"
-	"github.com/databricks/databricks-sdk-go"
 )
 
 // NewOutputFiler creates a filer for writing output files.
@@ -16,7 +16,7 @@ import (
 //
 // It is not possible to write notebooks through the workspace filesystem's FUSE mount for DBR versions less than 16.4.
 // This function ensures the correct filer is used based on the runtime environment.
-func NewOutputFiler(ctx context.Context, w *databricks.WorkspaceClient, outputDir string) (Filer, error) {
+func NewOutputFiler(ctx context.Context, outputDir string) (Filer, error) {
 	outputDir, err := filepath.Abs(outputDir)
 	if err != nil {
 		return nil, err
@@ -25,6 +25,7 @@ func NewOutputFiler(ctx context.Context, w *databricks.WorkspaceClient, outputDi
 	// If the CLI is running on DBR and we're writing to the workspace file system,
 	// use the extension-aware workspace filesystem filer.
 	if strings.HasPrefix(outputDir, "/Workspace/") && dbr.RunsOnRuntime(ctx) {
+		w := cmdctx.WorkspaceClient(ctx)
 		return NewWorkspaceFilesExtensionsClient(w, outputDir)
 	}
 

libs/filer/output_filer_test.go

@@ -5,6 +5,7 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/dbr"
 	"github.com/databricks/databricks-sdk-go"
 	workspaceConfig "github.com/databricks/databricks-sdk-go/config"
@@ -15,12 +16,8 @@ import (
 func TestNewOutputFilerLocal(t *testing.T) {
 	ctx := dbr.MockRuntime(context.Background(), dbr.Environment{IsDbr: false})
 
-	w := &databricks.WorkspaceClient{
-		Config: &workspaceConfig.Config{Host: "https://myhost.com"},
-	}
-
 	tmpDir := t.TempDir()
-	f, err := NewOutputFiler(ctx, w, tmpDir)
+	f, err := NewOutputFiler(ctx, tmpDir)
 	require.NoError(t, err)
 
 	assert.IsType(t, &LocalClient{}, f)
@@ -35,12 +32,8 @@ func TestNewOutputFilerLocalForNonWorkspacePath(t *testing.T) {
 	// Even on DBR, if path doesn't start with /Workspace/, use local client
 	ctx := dbr.MockRuntime(context.Background(), dbr.Environment{IsDbr: true, Version: "15.4"})
 
-	w := &databricks.WorkspaceClient{
-		Config: &workspaceConfig.Config{Host: "https://myhost.com"},
-	}
-
 	tmpDir := t.TempDir()
-	f, err := NewOutputFiler(ctx, w, tmpDir)
+	f, err := NewOutputFiler(ctx, tmpDir)
 	require.NoError(t, err)
 
 	assert.IsType(t, &LocalClient{}, f)
@@ -53,13 +46,12 @@ func TestNewOutputFilerDBR(t *testing.T) {
 	}
 
 	ctx := dbr.MockRuntime(context.Background(), dbr.Environment{IsDbr: true, Version: "15.4"})
-
-	w := &databricks.WorkspaceClient{
+	ctx = cmdctx.SetWorkspaceClient(ctx, &databricks.WorkspaceClient{
 		Config: &workspaceConfig.Config{Host: "https://myhost.com"},
-	}
+	})
 
 	// On DBR with /Workspace/ path, should use workspace files extensions client
-	f, err := NewOutputFiler(ctx, w, "/Workspace/Users/[email protected]/my-bundle")
+	f, err := NewOutputFiler(ctx, "/Workspace/Users/[email protected]/my-bundle")
 	require.NoError(t, err)
 
 	assert.IsType(t, &WorkspaceFilesExtensionsClient{}, f)

libs/template/resolver_test.go

@@ -5,23 +5,28 @@ import (
 	"testing"
 
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/dbr"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func testContext(t *testing.T) context.Context {
+	return dbr.DetectRuntime(context.Background())
+}
+
 func TestTemplateResolverBothTagAndBranch(t *testing.T) {
 	r := Resolver{
 		Tag:    "tag",
 		Branch: "branch",
 	}
 
-	_, err := r.Resolve(context.Background())
+	_, err := r.Resolve(testContext(t))
 	assert.EqualError(t, err, "only one of tag or branch can be specified")
 }
 
 func TestTemplateResolverErrorsWhenPromptingIsNotSupported(t *testing.T) {
 	r := Resolver{}
-	ctx := cmdio.MockDiscard(context.Background())
+	ctx := cmdio.MockDiscard(testContext(t))
 
 	_, err := r.Resolve(ctx)
 	assert.EqualError(t, err, "prompting is not supported. Please specify the path, name or URL of the template to use")
@@ -38,7 +43,7 @@ func TestTemplateResolverForDefaultTemplates(t *testing.T) {
 				TemplatePathOrUrl: name,
 			}
 
-			tmpl, err := r.Resolve(context.Background())
+			tmpl, err := r.Resolve(testContext(t))
 			require.NoError(t, err)
 
 			assert.Equal(t, &builtinReader{name: name}, tmpl.Reader)
@@ -52,7 +57,7 @@ func TestTemplateResolverForDefaultTemplates(t *testing.T) {
 			ConfigFile:        "/config/file",
 		}
 
-		tmpl, err := r.Resolve(context.Background())
+		tmpl, err := r.Resolve(testContext(t))
 		require.NoError(t, err)
 
 		// Assert reader and writer configuration
@@ -69,7 +74,7 @@ func TestTemplateResolverForCustomUrl(t *testing.T) {
 		ConfigFile:        "/config/file",
 	}
 
-	tmpl, err := r.Resolve(context.Background())
+	tmpl, err := r.Resolve(testContext(t))
 	require.NoError(t, err)
 
 	assert.Equal(t, Custom, tmpl.name)
@@ -89,7 +94,7 @@ func TestTemplateResolverForCustomPath(t *testing.T) {
 		ConfigFile:        "/config/file",
 	}
 
-	tmpl, err := r.Resolve(context.Background())
+	tmpl, err := r.Resolve(testContext(t))
 	require.NoError(t, err)
 
 	assert.Equal(t, Custom, tmpl.name)

libs/template/writer.go

@@ -2,14 +2,10 @@ package template
 
 import (
 	"context"
-	"path/filepath"
 	"sort"
 	"strconv"
-	"strings"
 
-	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
-	"github.com/databricks/cli/libs/dbr"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/jsonschema"
 	"github.com/databricks/cli/libs/telemetry"
@@ -49,20 +45,7 @@ type defaultWriter struct {
 func (tmpl *defaultWriter) Configure(ctx context.Context, configPath, outputDir string) error {
 	tmpl.configPath = configPath
 
-	// Workspace client is only needed when running on DBR and writing to /Workspace/.
-	// We avoid calling cmdctx.WorkspaceClient unconditionally because it panics
-	// if the workspace client is not set in the context.
-	var outputFiler filer.Filer
-	var err error
-	absOutputDir, err := filepath.Abs(outputDir)
-	if err != nil {
-		return err
-	}
-	if strings.HasPrefix(absOutputDir, "/Workspace/") && dbr.RunsOnRuntime(ctx) {
-		outputFiler, err = filer.NewOutputFiler(ctx, cmdctx.WorkspaceClient(ctx), outputDir)
-	} else {
-		outputFiler, err = filer.NewLocalClient(absOutputDir)
-	}
+	outputFiler, err := filer.NewOutputFiler(ctx, outputDir)
 	if err != nil {
 		return err
 	}

libs/template/writer_test.go

@@ -15,9 +15,11 @@ import (
 )
 
 func TestDefaultWriterConfigure(t *testing.T) {
+	ctx := dbr.DetectRuntime(context.Background())
+
 	// Test on local file system.
 	w := &defaultWriter{}
-	err := w.Configure(context.Background(), "/foo/bar", "/out/abc")
+	err := w.Configure(ctx, "/foo/bar", "/out/abc")
 	assert.NoError(t, err)
 
 	assert.Equal(t, "/foo/bar", w.configPath)
@@ -46,7 +48,7 @@ func TestDefaultWriterConfigureOnDBR(t *testing.T) {
 func TestMaterializeForNonTemplateDirectory(t *testing.T) {
 	tmpDir1 := t.TempDir()
 	tmpDir2 := t.TempDir()
-	ctx := context.Background()
+	ctx := dbr.DetectRuntime(context.Background())
 
 	w := &defaultWriter{}
 	err := w.Configure(ctx, "/foo/bar", tmpDir1)