Skip to content

Commit

Permalink
Fix template initialization when running on Databricks (#1912)
Browse files Browse the repository at this point in the history
## Changes

When running the CLI on Databricks Runtime (DBR), use the
extension-aware filer to write an instantiated template if the instance
path is located in the workspace filesystem.

Notebooks cannot be written through the workspace filesystem's FUSE
mount. As a result, this is the only method for initializing templates
that contain notebooks when running the CLI on DBR and writing to the
workspace filesystem.

Depends on #1910 and #1911.

Supersedes #1744.

## Tests

* Manually confirmed I can initialize a template with notebooks when
running the CLI from the web terminal.
  • Loading branch information
pietern authored Nov 20, 2024
1 parent 75b09ff commit 886e149
Show file tree
Hide file tree
Showing 4 changed files with 38 additions and 12 deletions.
32 changes: 30 additions & 2 deletions cmd/bundle/init.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package bundle

import (
"context"
"errors"
"fmt"
"io/fs"
Expand All @@ -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"
Expand Down Expand Up @@ -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]",
Expand Down Expand Up @@ -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.")
Expand Down Expand Up @@ -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 '*'
Expand All @@ -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
}
5 changes: 4 additions & 1 deletion internal/bundle/helpers.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
}

Expand Down
11 changes: 3 additions & 8 deletions libs/template/materialize.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -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
}
Expand Down
2 changes: 1 addition & 1 deletion libs/template/materialize_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
}

0 comments on commit 886e149

Please sign in to comment.