Use filer when running generate on DBR

bundle/generate/downloader.go

@@ -3,22 +3,20 @@ package generate
 import (
 	"context"
 	"fmt"
-	"io"
 	"net/http"
-	"os"
 	"path"
 	"path/filepath"
 	"strings"
 
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/notebook"
 	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/client"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/databricks/databricks-sdk-go/service/pipelines"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 	"golang.org/x/sync/errgroup"
-
-	"github.com/databricks/databricks-sdk-go/client"
 )
 
 type exportFile struct {
@@ -27,11 +25,12 @@ type exportFile struct {
 }
 
 type Downloader struct {
-	files     map[string]exportFile
-	w         *databricks.WorkspaceClient
-	sourceDir string
-	configDir string
-	basePath  string
+	files       map[string]exportFile
+	w           *databricks.WorkspaceClient
+	sourceDir   string
+	configDir   string
+	basePath    string
+	outputFiler filer.Filer
 }
 
 func (n *Downloader) MarkTaskForDownload(ctx context.Context, task *jobs.Task) error {
@@ -194,7 +193,7 @@ func (n *Downloader) relativePath(fullPath string) string {
 func (n *Downloader) FlushToDisk(ctx context.Context, force bool) error {
 	// First check that all files can be written
 	for targetPath := range n.files {
-		info, err := os.Stat(targetPath)
+		info, err := n.outputFiler.Stat(ctx, targetPath)
 		if err == nil {
 			if info.IsDir() {
 				return fmt.Errorf("%s is a directory", targetPath)
@@ -207,42 +206,36 @@ func (n *Downloader) FlushToDisk(ctx context.Context, force bool) error {
 
 	errs, errCtx := errgroup.WithContext(ctx)
 	for targetPath, exportFile := range n.files {
-		// Create parent directories if they don't exist
-		dir := filepath.Dir(targetPath)
-		err := os.MkdirAll(dir, 0o755)
-		if err != nil {
-			return err
-		}
 		errs.Go(func() error {
 			reader, err := n.w.Workspace.Download(errCtx, exportFile.path, workspace.DownloadFormat(exportFile.format))
 			if err != nil {
 				return err
 			}
+			defer reader.Close()
 
-			file, err := os.Create(targetPath)
-			if err != nil {
-				return err
+			mode := []filer.WriteMode{filer.CreateParentDirectories}
+			if force {
+				mode = append(mode, filer.OverwriteIfExists)
 			}
-			defer file.Close()
-
-			_, err = io.Copy(file, reader)
+			err = n.outputFiler.Write(errCtx, targetPath, reader, mode...)
 			if err != nil {
 				return err
 			}
 
 			cmdio.LogString(errCtx, "File successfully saved to "+targetPath)
-			return reader.Close()
+			return nil
 		})
 	}
 
 	return errs.Wait()
 }
 
-func NewDownloader(w *databricks.WorkspaceClient, sourceDir, configDir string) *Downloader {
+func NewDownloader(w *databricks.WorkspaceClient, sourceDir, configDir string, outputFiler filer.Filer) *Downloader {
 	return &Downloader{
-		files:     make(map[string]exportFile),
-		w:         w,
-		sourceDir: sourceDir,
-		configDir: configDir,
+		files:       make(map[string]exportFile),
+		w:           w,
+		sourceDir:   sourceDir,
+		configDir:   configDir,
+		outputFiler: outputFiler,
 	}
 }

bundle/generate/downloader_test.go

@@ -5,6 +5,8 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/databricks/cli/libs/fakefs"
+	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 	"github.com/stretchr/testify/assert"
@@ -18,7 +20,8 @@ func TestDownloader_MarkFileReturnsRelativePath(t *testing.T) {
 	dir := "base/dir/doesnt/matter"
 	sourceDir := filepath.Join(dir, "source")
 	configDir := filepath.Join(dir, "config")
-	downloader := NewDownloader(m.WorkspaceClient, sourceDir, configDir)
+	fakeFiler := filer.NewFakeFiler(map[string]fakefs.FileInfo{})
+	downloader := NewDownloader(m.WorkspaceClient, sourceDir, configDir, fakeFiler)
 
 	var err error
 

cmd/bundle/generate/app.go

@@ -10,6 +10,7 @@ import (
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/cli/libs/dyn/yamlsaver"
+	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/logdiag"
 	"github.com/databricks/cli/libs/textutil"
 	"github.com/databricks/databricks-sdk-go/service/apps"
@@ -77,7 +78,22 @@ per target environment.`,
 			return err
 		}
 
-		downloader := generate.NewDownloader(w, sourceDir, configDir)
+		outputFiler, err := filer.NewOutputFiler(ctx, b.BundleRootPath)
+		if err != nil {
+			return err
+		}
+
+		// Make sourceDir and configDir relative to the bundle root
+		sourceDir, err = makeRelativeToRoot(b.BundleRootPath, sourceDir)
+		if err != nil {
+			return err
+		}
+		configDir, err = makeRelativeToRoot(b.BundleRootPath, configDir)
+		if err != nil {
+			return err
+		}
+
+		downloader := generate.NewDownloader(w, sourceDir, configDir, outputFiler)
 
 		sourceCodePath := app.DefaultSourceCodePath
 		// If the source code path is not set, we don't need to download anything.
@@ -121,7 +137,7 @@ per target environment.`,
 		filename := filepath.Join(configDir, appKey+".app.yml")
 
 		saver := yamlsaver.NewSaver()
-		err = saver.SaveAsYAML(result, filename, force)
+		err = saver.SaveAsYAMLToFiler(ctx, outputFiler, result, filename, force)
 		if err != nil {
 			return err
 		}

cmd/bundle/generate/dashboard.go

@@ -7,7 +7,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"os"
 	"path"
 	"path/filepath"
 	"strings"
@@ -25,6 +24,7 @@ import (
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/cli/libs/dyn/yamlsaver"
+	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/logdiag"
 	"github.com/databricks/cli/libs/textutil"
 	"github.com/databricks/databricks-sdk-go"
@@ -66,6 +66,9 @@ type dashboard struct {
 	// Output and error streams.
 	out io.Writer
 	err io.Writer
+
+	// Output filer for writing files.
+	outputFiler filer.Filer
 }
 
 func (d *dashboard) resolveID(ctx context.Context, b *bundle.Bundle) string {
@@ -165,48 +168,42 @@ func remarshalJSON(data []byte) ([]byte, error) {
 	return buf.Bytes(), nil
 }
 
-func (d *dashboard) saveSerializedDashboard(_ context.Context, b *bundle.Bundle, dashboard *dashboards.Dashboard, filename string) error {
+func (d *dashboard) saveSerializedDashboard(ctx context.Context, dashboard *dashboards.Dashboard, filename string) error {
 	// Unmarshal and remarshal the serialized dashboard to ensure it is formatted correctly.
 	// The result will have alphabetically sorted keys and be indented.
 	data, err := remarshalJSON([]byte(dashboard.SerializedDashboard))
 	if err != nil {
 		return err
 	}
 
-	// Make sure the output directory exists.
-	if err := os.MkdirAll(filepath.Dir(filename), 0o755); err != nil {
-		return err
-	}
-
 	// Clean the filename to ensure it is a valid path (and can be used on this OS).
 	filename = filepath.Clean(filename)
 
-	// Attempt to make the path relative to the bundle root.
-	rel, err := filepath.Rel(b.BundleRootPath, filename)
-	if err != nil {
-		rel = filename
-	}
-
 	// Verify that the file does not already exist.
-	info, err := os.Stat(filename)
+	info, err := d.outputFiler.Stat(ctx, filename)
 	if err == nil {
 		if info.IsDir() {
-			return fmt.Errorf("%s is a directory", rel)
+			return fmt.Errorf("%s is a directory", filename)
 		}
 		if !d.force {
-			return fmt.Errorf("%s already exists. Use --force to overwrite", rel)
+			return fmt.Errorf("%s already exists. Use --force to overwrite", filename)
 		}
 	}
 
-	fmt.Fprintf(d.out, "Writing dashboard to %q\n", rel)
-	return os.WriteFile(filename, data, 0o644)
+	fmt.Fprintf(d.out, "Writing dashboard to %q\n", filename)
+
+	mode := []filer.WriteMode{filer.CreateParentDirectories}
+	if d.force {
+		mode = append(mode, filer.OverwriteIfExists)
+	}
+	return d.outputFiler.Write(ctx, filename, bytes.NewReader(data), mode...)
 }
 
-func (d *dashboard) saveConfiguration(ctx context.Context, b *bundle.Bundle, dashboard *dashboards.Dashboard, key string) error {
+func (d *dashboard) saveConfiguration(ctx context.Context, dashboard *dashboards.Dashboard, key string) error {
 	// Save serialized dashboard definition to the dashboard directory.
 	dashboardBasename := key + ".lvdash.json"
 	dashboardPath := filepath.Join(d.dashboardDir, dashboardBasename)
-	err := d.saveSerializedDashboard(ctx, b, dashboard, dashboardPath)
+	err := d.saveSerializedDashboard(ctx, dashboard, dashboardPath)
 	if err != nil {
 		return err
 	}
@@ -225,25 +222,14 @@ func (d *dashboard) saveConfiguration(ctx context.Context, b *bundle.Bundle, das
 		}),
 	}
 
-	// Make sure the output directory exists.
-	if err := os.MkdirAll(d.resourceDir, 0o755); err != nil {
-		return err
-	}
-
 	// Save the configuration to the resource directory.
 	resourcePath := filepath.Join(d.resourceDir, key+".dashboard.yml")
 	saver := yamlsaver.NewSaverWithStyle(map[string]yaml.Style{
 		"display_name": yaml.DoubleQuotedStyle,
 	})
 
-	// Attempt to make the path relative to the bundle root.
-	rel, err := filepath.Rel(b.BundleRootPath, resourcePath)
-	if err != nil {
-		rel = resourcePath
-	}
-
-	fmt.Fprintf(d.out, "Writing configuration to %q\n", rel)
-	err = saver.SaveAsYAML(result, resourcePath, d.force)
+	fmt.Fprintf(d.out, "Writing configuration to %q\n", resourcePath)
+	err = saver.SaveAsYAMLToFiler(ctx, d.outputFiler, result, resourcePath, d.force)
 	if err != nil {
 		return err
 	}
@@ -306,7 +292,7 @@ func (d *dashboard) updateDashboardForResource(ctx context.Context, b *bundle.Bu
 		}
 
 		if etag != dashboard.Etag {
-			err = d.saveSerializedDashboard(ctx, b, dashboard, dashboardPath)
+			err = d.saveSerializedDashboard(ctx, dashboard, dashboardPath)
 			if err != nil {
 				logdiag.LogError(ctx, err)
 				return
@@ -338,7 +324,7 @@ func (d *dashboard) generateForExisting(ctx context.Context, b *bundle.Bundle, d
 	}
 
 	key := textutil.NormalizeString(dashboard.DisplayName)
-	err = d.saveConfiguration(ctx, b, dashboard, key)
+	err = d.saveConfiguration(ctx, dashboard, key)
 	if err != nil {
 		logdiag.LogError(ctx, err)
 	}
@@ -354,12 +340,18 @@ func (d *dashboard) generateForExisting(ctx context.Context, b *bundle.Bundle, d
 }
 
 func (d *dashboard) initialize(ctx context.Context, b *bundle.Bundle) {
-	// Make the paths absolute if they aren't already.
-	if !filepath.IsAbs(d.resourceDir) {
-		d.resourceDir = filepath.Join(b.BundleRootPath, d.resourceDir)
+	var err error
+
+	// Make paths relative to the bundle root (required for the filer which is rooted there).
+	d.resourceDir, err = makeRelativeToRoot(b.BundleRootPath, d.resourceDir)
+	if err != nil {
+		logdiag.LogError(ctx, err)
+		return
 	}
-	if !filepath.IsAbs(d.dashboardDir) {
-		d.dashboardDir = filepath.Join(b.BundleRootPath, d.dashboardDir)
+	d.dashboardDir, err = makeRelativeToRoot(b.BundleRootPath, d.dashboardDir)
+	if err != nil {
+		logdiag.LogError(ctx, err)
+		return
 	}
 
 	// Make sure we know how the dashboard path is relative to the resource path.
@@ -370,6 +362,14 @@ 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.BundleRootPath)
+	if err != nil {
+		logdiag.LogError(ctx, err)
+		return
+	}
+	d.outputFiler = outputFiler
 }
 
 func (d *dashboard) runForResource(ctx context.Context, b *bundle.Bundle) {