Refactor bundle init

bundle/run/app_test.go

@@ -1,7 +1,6 @@
 package run
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"os"
@@ -16,7 +15,6 @@ import (
 	"github.com/databricks/cli/bundle/internal/bundletest"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/dyn"
-	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/vfs"
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
 	"github.com/databricks/databricks-sdk-go/service/apps"
@@ -76,9 +74,7 @@ func setupBundle(t *testing.T) (context.Context, *bundle.Bundle, *mocks.MockWork
 	b.SetWorkpaceClient(mwc.WorkspaceClient)
 	bundletest.SetLocation(b, "resources.apps.my_app", []dyn.Location{{File: "./databricks.yml"}})
 
-	ctx := context.Background()
-	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
-	ctx = cmdio.NewContext(ctx, cmdio.NewLogger(flags.ModeAppend))
+	ctx := cmdio.MockDiscard(context.Background())
 
 	diags := bundle.Apply(ctx, b, bundle.Seq(
 		mutator.DefineDefaultWorkspacePaths(),

bundle/run/job_test.go

@@ -1,7 +1,6 @@
 package run
 
 import (
-	"bytes"
 	"context"
 	"testing"
 	"time"
@@ -159,8 +158,8 @@ func TestJobRunnerRestart(t *testing.T) {
 
 		m := mocks.NewMockWorkspaceClient(t)
 		b.SetWorkpaceClient(m.WorkspaceClient)
-		ctx := context.Background()
-		ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", ""))
+
+		ctx := cmdio.MockDiscard(context.Background())
 		ctx = cmdio.NewContext(ctx, cmdio.NewLogger(flags.ModeAppend))
 
 		jobApi := m.GetMockJobsAPI()
@@ -230,8 +229,8 @@ func TestJobRunnerRestartForContinuousUnpausedJobs(t *testing.T) {
 
 	m := mocks.NewMockWorkspaceClient(t)
 	b.SetWorkpaceClient(m.WorkspaceClient)
-	ctx := context.Background()
-	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
+
+	ctx := cmdio.MockDiscard(context.Background())
 	ctx = cmdio.NewContext(ctx, cmdio.NewLogger(flags.ModeAppend))
 
 	jobApi := m.GetMockJobsAPI()

bundle/run/pipeline_test.go

@@ -1,7 +1,6 @@
 package run
 
 import (
-	"bytes"
 	"context"
 	"testing"
 	"time"
@@ -75,8 +74,8 @@ func TestPipelineRunnerRestart(t *testing.T) {
 		Host: "https://test.com",
 	}
 	b.SetWorkpaceClient(m.WorkspaceClient)
-	ctx := context.Background()
-	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
+
+	ctx := cmdio.MockDiscard(context.Background())
 	ctx = cmdio.NewContext(ctx, cmdio.NewLogger(flags.ModeAppend))
 
 	mockWait := &pipelines.WaitGetPipelineIdle[struct{}]{

cmd/bundle/init.go

@@ -1,175 +1,15 @@
 package bundle
 
 import (
-	"context"
 	"errors"
 	"fmt"
-	"io/fs"
-	"os"
-	"path/filepath"
-	"slices"
-	"strings"
 
 	"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"
 )
 
-var gitUrlPrefixes = []string{
-	"https://",
-	"git@",
-}
-
-type nativeTemplate struct {
-	name        string
-	gitUrl      string
-	description string
-	aliases     []string
-	hidden      bool
-}
-
-const customTemplate = "custom..."
-
-var nativeTemplates = []nativeTemplate{
-	{
-		name:        "default-python",
-		description: "The default Python template for Notebooks / Delta Live Tables / Workflows",
-	},
-	{
-		name:        "default-sql",
-		description: "The default SQL template for .sql files that run with Databricks SQL",
-	},
-	{
-		name:        "dbt-sql",
-		description: "The dbt SQL template (databricks.com/blog/delivering-cost-effective-data-real-time-dbt-and-databricks)",
-	},
-	{
-		name:        "mlops-stacks",
-		gitUrl:      "https://github.com/databricks/mlops-stacks",
-		description: "The Databricks MLOps Stacks template (github.com/databricks/mlops-stacks)",
-		aliases:     []string{"mlops-stack"},
-	},
-	{
-		name:        "default-pydabs",
-		gitUrl:      "https://databricks.github.io/workflows-authoring-toolkit/pydabs-template.git",
-		hidden:      true,
-		description: "The default PyDABs template",
-	},
-	{
-		name:        customTemplate,
-		description: "Bring your own template",
-	},
-}
-
-// Return template descriptions for command-line help
-func nativeTemplateHelpDescriptions() string {
-	var lines []string
-	for _, template := range nativeTemplates {
-		if template.name != customTemplate && !template.hidden {
-			lines = append(lines, fmt.Sprintf("- %s: %s", template.name, template.description))
-		}
-	}
-	return strings.Join(lines, "\n")
-}
-
-// Return template options for an interactive prompt
-func nativeTemplateOptions() []cmdio.Tuple {
-	names := make([]cmdio.Tuple, 0, len(nativeTemplates))
-	for _, template := range nativeTemplates {
-		if template.hidden {
-			continue
-		}
-		tuple := cmdio.Tuple{
-			Name: template.name,
-			Id:   template.description,
-		}
-		names = append(names, tuple)
-	}
-	return names
-}
-
-func getNativeTemplateByDescription(description string) string {
-	for _, template := range nativeTemplates {
-		if template.description == description {
-			return template.name
-		}
-	}
-	return ""
-}
-
-func getUrlForNativeTemplate(name string) string {
-	for _, template := range nativeTemplates {
-		if template.name == name {
-			return template.gitUrl
-		}
-		if slices.Contains(template.aliases, name) {
-			return template.gitUrl
-		}
-	}
-	return ""
-}
-
-func getFsForNativeTemplate(name string) (fs.FS, error) {
-	builtin, err := template.Builtin()
-	if err != nil {
-		return nil, err
-	}
-
-	// If this is a built-in template, the return value will be non-nil.
-	var templateFS fs.FS
-	for _, entry := range builtin {
-		if entry.Name == name {
-			templateFS = entry.FS
-			break
-		}
-	}
-
-	return templateFS, nil
-}
-
-func isRepoUrl(url string) bool {
-	result := false
-	for _, prefix := range gitUrlPrefixes {
-		if strings.HasPrefix(url, prefix) {
-			result = true
-			break
-		}
-	}
-	return result
-}
-
-// Computes the repo name from the repo URL. Treats the last non empty word
-// when splitting at '/' as the repo name. For example: for url git@github.com:databricks/cli.git
-// the name would be "cli.git"
-func repoName(url string) string {
-	parts := strings.Split(strings.TrimRight(url, "/"), "/")
-	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]",
@@ -182,7 +22,7 @@ TEMPLATE_PATH optionally specifies which template to use. It can be one of the f
 - a local file system path with a template directory
 - a Git repository URL, e.g. https://github.com/my/repository
 
-See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more information on templates.`, nativeTemplateHelpDescriptions()),
+See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more information on templates.`, template.HelpDescriptions()),
 	}
 
 	var configFile string
@@ -202,88 +42,32 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 			return errors.New("only one of --tag or --branch can be specified")
 		}
 
-		// Git ref to use for template initialization
-		ref := branch
-		if tag != "" {
-			ref = tag
-		}
-
-		ctx := cmd.Context()
-		var templatePath string
+		var templatePathOrUrl string
 		if len(args) > 0 {
-			templatePath = args[0]
-		} else {
-			var err error
-			if !cmdio.IsPromptSupported(ctx) {
-				return errors.New("please specify a template")
-			}
-			description, err := cmdio.SelectOrdered(ctx, nativeTemplateOptions(), "Template to use")
-			if err != nil {
-				return err
-			}
-			templatePath = getNativeTemplateByDescription(description)
+			templatePathOrUrl = args[0]
 		}
-
-		outputFiler, err := constructOutputFiler(ctx, outputDir)
-		if err != nil {
-			return err
+		r := template.Resolver{
+			TemplatePathOrUrl: templatePathOrUrl,
+			ConfigFile:        configFile,
+			OutputDir:         outputDir,
+			TemplateDir:       templateDir,
+			Tag:               tag,
+			Branch:            branch,
 		}
 
-		if templatePath == customTemplate {
+		ctx := cmd.Context()
+		tmpl, err := r.Resolve(ctx)
+		if errors.Is(err, template.ErrCustomSelected) {
 			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.")
 			return nil
 		}
-
-		// Expand templatePath to a git URL if it's an alias for a known native template
-		// and we know it's git URL.
-		if gitUrl := getUrlForNativeTemplate(templatePath); gitUrl != "" {
-			templatePath = gitUrl
-		}
-
-		if !isRepoUrl(templatePath) {
-			if templateDir != "" {
-				return errors.New("--template-dir can only be used with a Git repository URL")
-			}
-
-			templateFS, err := getFsForNativeTemplate(templatePath)
-			if err != nil {
-				return err
-			}
-
-			// If this is not a built-in template, then it must be a local file system path.
-			if templateFS == nil {
-				templateFS = os.DirFS(templatePath)
-			}
-
-			// 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, outputFiler)
-		}
-
-		// Create a temporary directory with the name of the repository.  The '*'
-		// character is replaced by a random string in the generated temporary directory.
-		repoDir, err := os.MkdirTemp("", repoName(templatePath)+"-*")
-		if err != nil {
-			return err
-		}
-
-		// start the spinner
-		promptSpinner := cmdio.Spinner(ctx)
-		promptSpinner <- "Downloading the template\n"
-
-		// TODO: Add automated test that the downloaded git repo is cleaned up.
-		// Clone the repository in the temporary directory
-		err = git.Clone(ctx, templatePath, ref, repoDir)
-		close(promptSpinner)
 		if err != nil {
 			return err
 		}
+		defer tmpl.Reader.Cleanup()
 
-		// 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, outputFiler)
+		return tmpl.Writer.Materialize(ctx, tmpl.Reader)
 	}
 	return cmd
 }

integration/bundle/helpers_test.go

@@ -16,7 +16,6 @@ import (
 	"github.com/databricks/cli/internal/testutil"
 	"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/folders"
 	"github.com/databricks/cli/libs/template"
@@ -40,10 +39,19 @@ func initTestTemplateWithBundleRoot(t testutil.TestingT, ctx context.Context, te
 	cmd := cmdio.NewIO(ctx, flags.OutputJSON, strings.NewReader(""), os.Stdout, os.Stderr, "", "bundles")
 	ctx = cmdio.InContext(ctx, cmd)
 
-	out, err := filer.NewLocalClient(bundleRoot)
+	r := template.Resolver{
+		TemplatePathOrUrl: templateRoot,
+		ConfigFile:        configFilePath,
+		OutputDir:         bundleRoot,
+	}
+
+	tmpl, err := r.Resolve(ctx)
 	require.NoError(t, err)
-	err = template.Materialize(ctx, configFilePath, os.DirFS(templateRoot), out)
+	defer tmpl.Reader.Cleanup()
+
+	err = tmpl.Writer.Materialize(ctx, tmpl.Reader)
 	require.NoError(t, err)
+
 	return bundleRoot
 }