Fix template initialization when running on Databricks

Author: pietern

cmd/bundle/init.go

@@ -1,6 +1,7 @@
 package bundle
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io/fs"
@@ -11,6 +12,8 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/dbr"
+	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/git"
 	"github.com/databricks/cli/libs/template"
 	"github.com/spf13/cobra"
@@ -147,6 +150,26 @@ func repoName(url string) string {
 	return parts[len(parts)-1]
 }
 
+func constructOutputFiler(ctx context.Context, outputDir string) (filer.Filer, error) {
+	outputDir, err := filepath.Abs(outputDir)
+	if err != nil {
+		return nil, err
+	}
+
+	// If the CLI is running on DBR and we're writing to the workspace file system,
+	// use the extension-aware workspace filesystem filer to instantiate the template.
+	//
+	// It is not possible to write notebooks through the workspace filesystem's FUSE mount.
+	// Therefore this is the only way we can initialize templates that contain notebooks
+	// when running the CLI on DBR and initializing a template to the workspace.
+	//
+	if strings.HasPrefix(outputDir, "/Workspace/") && dbr.RunsOnRuntime(ctx) {
+		return filer.NewWorkspaceFilesExtensionsClient(root.WorkspaceClient(ctx), outputDir)
+	}
+
+	return filer.NewLocalClient(outputDir)
+}
+
 func newInitCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "init [TEMPLATE_PATH]",
@@ -201,6 +224,11 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 			templatePath = getNativeTemplateByDescription(description)
 		}
 
+		outputFiler, err := constructOutputFiler(ctx, outputDir)
+		if err != nil {
+			return err
+		}
+
 		if templatePath == customTemplate {
 			cmdio.LogString(ctx, "Please specify a path or Git repository to use a custom template.")
 			cmdio.LogString(ctx, "See https://docs.databricks.com/en/dev-tools/bundles/templates.html to learn more about custom templates.")
@@ -230,7 +258,7 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 
 			// skip downloading the repo because input arg is not a URL. We assume
 			// it's a path on the local file system in that case
-			return template.Materialize(ctx, configFile, templateFS, outputDir)
+			return template.Materialize(ctx, configFile, templateFS, outputFiler)
 		}
 
 		// Create a temporary directory with the name of the repository.  The '*'
@@ -255,7 +283,7 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 		// Clean up downloaded repository once the template is materialized.
 		defer os.RemoveAll(repoDir)
 		templateFS := os.DirFS(filepath.Join(repoDir, templateDir))
-		return template.Materialize(ctx, configFile, templateFS, outputDir)
+		return template.Materialize(ctx, configFile, templateFS, outputFiler)
 	}
 	return cmd
 }

internal/bundle/helpers.go

@@ -16,6 +16,7 @@ import (
 	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/template"
 	"github.com/databricks/cli/libs/vfs"
@@ -42,7 +43,9 @@ func initTestTemplateWithBundleRoot(t *testing.T, ctx context.Context, templateN
 	cmd := cmdio.NewIO(flags.OutputJSON, strings.NewReader(""), os.Stdout, os.Stderr, "", "bundles")
 	ctx = cmdio.InContext(ctx, cmd)
 
-	err = template.Materialize(ctx, configFilePath, os.DirFS(templateRoot), bundleRoot)
+	out, err := filer.NewLocalClient(bundleRoot)
+	require.NoError(t, err)
+	err = template.Materialize(ctx, configFilePath, os.DirFS(templateRoot), out)
 	return bundleRoot, err
 }
 

libs/template/materialize.go

@@ -21,8 +21,8 @@ const schemaFileName = "databricks_template_schema.json"
 //	ctx: 			context containing a cmdio object. This is used to prompt the user
 //	configFilePath: file path containing user defined config values
 //	templateFS: 	root of the template definition
-//	outputDir: 	root of directory where to initialize the template
-func Materialize(ctx context.Context, configFilePath string, templateFS fs.FS, outputDir string) error {
+//	outputFiler: 	filer to use for writing the initialized template
+func Materialize(ctx context.Context, configFilePath string, templateFS fs.FS, outputFiler filer.Filer) error {
 	if _, err := fs.Stat(templateFS, schemaFileName); errors.Is(err, fs.ErrNotExist) {
 		return fmt.Errorf("not a bundle template: expected to find a template schema file at %s", schemaFileName)
 	}
@@ -73,12 +73,7 @@ func Materialize(ctx context.Context, configFilePath string, templateFS fs.FS, o
 		return err
 	}
 
-	out, err := filer.NewLocalClient(outputDir)
-	if err != nil {
-		return err
-	}
-
-	err = r.persistToDisk(ctx, out)
+	err = r.persistToDisk(ctx, outputFiler)
 	if err != nil {
 		return err
 	}

libs/template/materialize_test.go

@@ -19,6 +19,6 @@ func TestMaterializeForNonTemplateDirectory(t *testing.T) {
 	ctx := root.SetWorkspaceClient(context.Background(), w)
 
 	// Try to materialize a non-template directory.
-	err = Materialize(ctx, "", os.DirFS(tmpDir), "")
+	err = Materialize(ctx, "", os.DirFS(tmpDir), nil)
 	assert.EqualError(t, err, fmt.Sprintf("not a bundle template: expected to find a template schema file at %s", schemaFileName))
 }