From 478b0eb60c29c0418486b978f058ee4f00a0af2b Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Mon, 18 Nov 2024 23:08:07 +0100 Subject: [PATCH] Fix template initialization when running on Databricks --- cmd/bundle/init.go | 32 +++++++++++++++++++++++++++++-- internal/bundle/helpers.go | 5 ++++- libs/template/materialize.go | 11 +++-------- libs/template/materialize_test.go | 2 +- 4 files changed, 38 insertions(+), 12 deletions(-) diff --git a/cmd/bundle/init.go b/cmd/bundle/init.go index d31a702a..687c141e 100644 --- a/cmd/bundle/init.go +++ b/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 } diff --git a/internal/bundle/helpers.go b/internal/bundle/helpers.go index 9740061e..dd9c841c 100644 --- a/internal/bundle/helpers.go +++ b/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 } diff --git a/libs/template/materialize.go b/libs/template/materialize.go index 8338e119..ee30444a 100644 --- a/libs/template/materialize.go +++ b/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 } diff --git a/libs/template/materialize_test.go b/libs/template/materialize_test.go index dc510a30..f7cd916e 100644 --- a/libs/template/materialize_test.go +++ b/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)) }