diff --git a/bundle/libraries/filer_volume.go b/bundle/libraries/filer_volume.go index f4b5f51f0c2..fdd6190e019 100644 --- a/bundle/libraries/filer_volume.go +++ b/bundle/libraries/filer_volume.go @@ -10,6 +10,6 @@ import ( func filerForVolume(ctx context.Context, b *bundle.Bundle, uploadPath string) (filer.Filer, string, diag.Diagnostics) { w := b.WorkspaceClient(ctx) - f, err := filer.NewFilesClient(w, uploadPath) + f, err := filer.NewFilesClient(ctx, w, uploadPath) return f, uploadPath, diag.FromErr(err) } diff --git a/cmd/fs/cp.go b/cmd/fs/cp.go index d7ae6517539..db2d615e916 100644 --- a/cmd/fs/cp.go +++ b/cmd/fs/cp.go @@ -18,11 +18,19 @@ import ( "golang.org/x/sync/errgroup" ) -// Default number of concurrent file copy operations. This is a conservative -// default that should be sufficient to fully utilize the available bandwidth -// in most cases. +// defaultConcurrency is the number of files copied in parallel. Each in-flight +// file drives Files API requests (a single-shot PUT, or the multipart +// control-plane calls), which allow only ~10 concurrent requests, so this stays +// conservative regardless of whether multipart upload is enabled. const defaultConcurrency = 8 +// multipartUploadConcurrency is the cloud-transfer budget for large-file +// (multipart) writes: a shared cap on concurrent part uploads, which go to cloud +// storage rather than the rate-limited Files API and so can fan out wider than +// the file-level copy concurrency. It is sized independently of --concurrency and +// applies only to the Volumes target filer. +const multipartUploadConcurrency = 64 + // errInvalidConcurrency is returned when the value of the concurrency // flag is invalid. var errInvalidConcurrency = errors.New("--concurrency must be at least 1") @@ -37,6 +45,11 @@ type copy struct { sourceScheme string targetScheme string + // showProgress renders an upload progress bar. It is set only for a single + // large-file copy to a Volume, not for recursive copies (where many files + // would each fight for the spinner line). + showProgress bool + mu sync.Mutex // protect output from concurrent writes } @@ -123,20 +136,31 @@ func (c *copy) cpFileToFile(ctx context.Context, sourcePath, targetPath string) } defer r.Close() + // For a single large-file copy, attach a progress callback to the context + // that the Files filer forwards to the upload engine, rendering an upload bar. + // The spinner is stopped before any event line is emitted so its final frame + // does not overwrite it. + closeProgress := func() {} + if c.showProgress { + fn, stop := newProgressFunc(ctx) + ctx = filer.WithUploadProgress(ctx, fn) + closeProgress = stop + } + + var writeErr error if c.overwrite { - err = c.targetFiler.Write(ctx, targetPath, r, filer.OverwriteIfExists) - if err != nil { - return err - } + writeErr = c.targetFiler.Write(ctx, targetPath, r, filer.OverwriteIfExists) } else { - err = c.targetFiler.Write(ctx, targetPath, r) - // skip if file already exists - if err != nil && errors.Is(err, fs.ErrExist) { - return c.emitFileSkippedEvent(ctx, sourcePath, targetPath) - } - if err != nil { - return err - } + writeErr = c.targetFiler.Write(ctx, targetPath, r) + } + closeProgress() + + // skip if file already exists + if !c.overwrite && writeErr != nil && errors.Is(writeErr, fs.ErrExist) { + return c.emitFileSkippedEvent(ctx, sourcePath, targetPath) + } + if writeErr != nil { + return writeErr } return c.emitFileCopiedEvent(ctx, sourcePath, targetPath) } @@ -229,9 +253,20 @@ func newCpCommand() *cobra.Command { return err } - // Get target filer and target path without scheme + // The cloud-transfer budget for large-file (multipart) writes is sized + // independently of c.concurrency (the file-level copy parallelism, which + // stays at the Files-API-safe default). A large file's part uploads go to + // cloud storage rather than the rate-limited Files API, so they can fan out + // wider; the budget is shared across all files in a recursive copy and + // applies only to the Volumes target filer. + uploadConcurrency := c.concurrency + if filer.MultipartUploadEnabled(ctx) { + uploadConcurrency = multipartUploadConcurrency + } + + // Get target filer and target path without scheme. fullTargetPath := args[1] - targetFiler, targetPath, err := filerForPath(ctx, fullTargetPath) + targetFiler, targetPath, err := filerForUploadTarget(ctx, fullTargetPath, uploadConcurrency) if err != nil { return err } @@ -259,6 +294,10 @@ func newCpCommand() *cobra.Command { return c.cpDirToDir(ctx, sourcePath, targetPath) } + // A single large file copied to a Volume goes through the multipart engine, + // which reports progress; render an upload bar for it. + c.showProgress = filer.MultipartUploadEnabled(ctx) && strings.HasPrefix(targetPath, "/Volumes/") + // If target path has a trailing separator, trim it and let case 2 handle it if hasTrailingDirSeparator(fullTargetPath) { targetPath = trimTrailingDirSeparators(targetPath) diff --git a/cmd/fs/helpers.go b/cmd/fs/helpers.go index 54d9463ed91..d4f1834c0fa 100644 --- a/cmd/fs/helpers.go +++ b/cmd/fs/helpers.go @@ -14,6 +14,20 @@ import ( ) func filerForPath(ctx context.Context, fullPath string) (filer.Filer, string, error) { + return newFilerForPath(ctx, fullPath, 0) +} + +// filerForUploadTarget is filerForPath for a copy target. When the target +// resolves to a Volumes (Files API) filer, large-file (multipart) uploads use +// uploadConcurrency as the shared transfer budget; other schemes ignore it. +func filerForUploadTarget(ctx context.Context, fullPath string, uploadConcurrency int) (filer.Filer, string, error) { + return newFilerForPath(ctx, fullPath, uploadConcurrency) +} + +// newFilerForPath resolves a filer and the scheme-stripped path. uploadConcurrency +// (when > 0) sizes a Volumes filer's large-file upload budget; it is ignored for +// every other scheme. +func newFilerForPath(ctx context.Context, fullPath string, uploadConcurrency int) (filer.Filer, string, error) { // Split path at : to detect any file schemes parts := strings.SplitN(fullPath, ":", 2) @@ -40,7 +54,11 @@ func filerForPath(ctx context.Context, fullPath string) (filer.Filer, string, er // If the specified path has the "Volumes" prefix, use the Files API. if strings.HasPrefix(path, "/Volumes/") { - f, err := filer.NewFilesClient(w, "/") + var opts []filer.FilesClientOption + if uploadConcurrency > 0 { + opts = append(opts, filer.WithUploadConcurrency(uploadConcurrency)) + } + f, err := filer.NewFilesClient(ctx, w, "/", opts...) return f, path, err } diff --git a/cmd/fs/upload_progress.go b/cmd/fs/upload_progress.go new file mode 100644 index 00000000000..a186fb5fa76 --- /dev/null +++ b/cmd/fs/upload_progress.go @@ -0,0 +1,256 @@ +package fs + +import ( + "context" + "fmt" + "strings" + "sync" + "time" + + "github.com/charmbracelet/lipgloss" + + "github.com/databricks/cli/libs/cmdio" + files "github.com/databricks/sdk-go/files/v2" +) + +// renderThrottle caps how often the interactive bar is redrawn so a fast upload +// does not spend its time repainting the terminal; completion always renders. +const renderThrottle = 100 * time.Millisecond + +// preparingFrame is how long each "Preparing upload" dot frame is shown before +// the upload reports its first completed part. +const preparingFrame = 400 * time.Millisecond + +// barWidth keeps the rendered bar narrow enough that the whole progress line +// stays on a single terminal row, so the spinner's in-place redraw is clean. +const barWidth = 24 + +// rateWindow is the trailing duration over which the transfer rate is averaged. +// A few seconds smooths the burstiness of concurrent part completions while +// still tracking genuine changes in throughput. +const rateWindow = 5 * time.Second + +// Bar styling. Green for the filled portion matches the cmdio spinner glyph; +// the remainder is dimmed. +var ( + barFilledStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("10")) + barEmptyStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("240")) +) + +// barBlocks are partial block glyphs for sub-character precision, so the bar +// advances smoothly rather than in whole-cell jumps. Index 0 is a full cell +// (8/8); index i is (8-i)/8 of a cell. Matches experimental/genie/agentstream. +var barBlocks = []string{"█", "▉", "▊", "▋", "▌", "▍", "▎", "▏"} + +// humanBytes formats a byte count using binary (1024-based) units. +func humanBytes(n int64) string { + const unit = 1024 + if n < unit { + return fmt.Sprintf("%d B", n) + } + units := []string{"KiB", "MiB", "GiB", "TiB", "PiB"} + div, exp := int64(unit), 0 + for n/div >= unit && exp < len(units)-1 { + div *= unit + exp++ + } + return fmt.Sprintf("%.1f %s", float64(n)/float64(div), units[exp]) +} + +// formatSpeed formats a transfer rate in bytes per second. +func formatSpeed(bytesPerSec float64) string { + if bytesPerSec < 0 { + bytesPerSec = 0 + } + return humanBytes(int64(bytesPerSec)) + "/s" +} + +// formatETA formats the estimated time remaining as MM:SS. It returns a +// placeholder when the rate is unknown so the line does not show a misleading +// estimate before any throughput has been measured. +func formatETA(remaining int64, bytesPerSec float64) string { + if bytesPerSec <= 0 || remaining < 0 { + return "--:--" + } + secs := int(float64(remaining) / bytesPerSec) + return fmt.Sprintf("%02d:%02d", secs/60, secs%60) +} + +// formatPlainProgress renders a single non-interactive progress line. +func formatPlainProgress(p files.Progress) string { + if p.Total <= 0 { + return "Uploaded " + humanBytes(p.Transferred) + } + pct := int(float64(p.Transferred) / float64(p.Total) * 100) + return fmt.Sprintf("Uploaded %s / %s (%d%%)", humanBytes(p.Transferred), humanBytes(p.Total), pct) +} + +// renderBar returns a fixed-width progress bar for the given completion ratio. +// The total visible width is always barWidth: full cells, an optional partial +// cell for the fractional remainder, then the dimmed empty track. +func renderBar(ratio float64) string { + ratio = min(max(ratio, 0), 1) + exact := ratio * barWidth + full := int(exact) + + filled := strings.Repeat("█", full) + partial := int((exact - float64(full)) * 8) + if partial > 0 && full < barWidth { + filled += barBlocks[8-partial] + full++ // the partial glyph occupies one cell of the track + } + + return barFilledStyle.Render(filled) + + barEmptyStyle.Render(strings.Repeat("░", barWidth-full)) +} + +// rateSample is a cumulative byte count observed at a point in time. +type rateSample struct { + at time.Time + bytes int64 +} + +// rateMeter computes the transfer rate as a rolling average over a trailing +// time window. Averaging over elapsed wall-clock time, rather than an +// exponential moving average over samples, keeps the reported speed and the +// ETA derived from it stable: the engine reports progress in irregular bursts +// as concurrent parts complete, and a per-sample EMA spikes on every burst. +// Callers pass the current time so the meter stays deterministic and +// unit-testable. +type rateMeter struct { + samples []rateSample +} + +// observe records a cumulative byte count at now and returns the average +// bytes/sec over the trailing rateWindow. It returns 0 until two samples span +// a positive interval (the first sample, or one taken after a stall longer +// than the window, leaves nothing to average against). +func (m *rateMeter) observe(now time.Time, transferred int64) float64 { + // Drop samples that have aged out of the window, then record this one. + // Filtering in place reuses the backing array; at the render cadence the + // window holds at most a few dozen samples, so this stays cheap. + cutoff := now.Add(-rateWindow) + kept := m.samples[:0] + for _, s := range m.samples { + if !s.at.Before(cutoff) { + kept = append(kept, s) + } + } + m.samples = append(kept, rateSample{at: now, bytes: transferred}) + + oldest := m.samples[0] + dt := now.Sub(oldest.at).Seconds() + if dt <= 0 { + return 0 + } + return float64(transferred-oldest.bytes) / dt +} + +// progressRenderer builds the rich single-line progress string shown as the +// spinner suffix in an interactive terminal. +type progressRenderer struct { + meter rateMeter +} + +// newProgressRenderer creates a renderer. +func newProgressRenderer() *progressRenderer { + return &progressRenderer{} +} + +// render returns the progress line for the given progress sample at time now. +func (r *progressRenderer) render(now time.Time, p files.Progress) string { + rate := r.meter.observe(now, p.Transferred) + if p.Total <= 0 { + return fmt.Sprintf("%s %s", humanBytes(p.Transferred), formatSpeed(rate)) + } + ratio := float64(p.Transferred) / float64(p.Total) + return fmt.Sprintf("%s %.0f%% %s/%s %s ETA %s", + renderBar(ratio), ratio*100, + humanBytes(p.Transferred), humanBytes(p.Total), + formatSpeed(rate), formatETA(p.Total-p.Transferred, rate)) +} + +// newProgressFunc returns the upload progress callback for the current terminal +// and a function that stops it. In an interactive terminal it renders a rich +// single-line bar via the spinner; otherwise it logs coarse progress so a long +// upload still shows life. The returned stop function is idempotent (it is safe +// to defer it and also call it before printing a summary line). The callback is +// serialized by the engine, so its captured state needs no locking. +func newProgressFunc(ctx context.Context) (files.ProgressFunc, func()) { + if cmdio.GetInteractiveMode(ctx) == cmdio.InteractiveModeNone { + nextPct := 10 + fn := func(p files.Progress) { + // The final summary line covers completion; only log intermediate steps. + if p.Total <= 0 || p.Transferred >= p.Total { + return + } + pct := int(float64(p.Transferred) / float64(p.Total) * 100) + if pct < nextPct { + return + } + cmdio.LogString(ctx, formatPlainProgress(p)) + for pct >= nextPct { + nextPct += 10 + } + } + return fn, func() {} + } + + sp := cmdio.NewSpinner(ctx, cmdio.WithElapsedTime()) + r := newProgressRenderer() + + // The callback first fires only when a part completes; session initiation, URL + // minting, and the first PUT take a beat, during which it never runs. Animate a + // "Preparing upload" label with cycling dots over that window so the spinner is + // not bare. mu serializes the animator's suffix updates with the callback's, and + // preparing gates the animator off once real progress arrives, so the bar + // replaces the label with no flicker. + var ( + mu sync.Mutex + preparing = true + lastRender time.Time + ) + prepCtx, stopPreparing := context.WithCancel(ctx) + prepDone := make(chan struct{}) + go func() { + defer close(prepDone) + ticker := time.NewTicker(preparingFrame) + defer ticker.Stop() + for dots := 1; ; dots = dots%3 + 1 { + mu.Lock() + if !preparing { + mu.Unlock() + return + } + sp.Update("Preparing upload" + strings.Repeat(".", dots)) + mu.Unlock() + select { + case <-prepCtx.Done(): + return + case <-ticker.C: + } + } + }() + + fn := func(p files.Progress) { + mu.Lock() + defer mu.Unlock() + if preparing { + preparing = false + stopPreparing() // first real progress: stop the animator, the bar takes over + } + now := time.Now() + if p.Total > 0 && p.Transferred < p.Total && now.Sub(lastRender) < renderThrottle { + return + } + lastRender = now + sp.Update(r.render(now, p)) + } + // Stop the animator and wait for it to exit before closing the spinner, so no + // suffix update races the spinner shutdown. + return fn, func() { + stopPreparing() + <-prepDone + sp.Close() + } +} diff --git a/cmd/fs/upload_progress_test.go b/cmd/fs/upload_progress_test.go new file mode 100644 index 00000000000..2786cf0fd40 --- /dev/null +++ b/cmd/fs/upload_progress_test.go @@ -0,0 +1,96 @@ +package fs + +import ( + "testing" + "time" + + "github.com/charmbracelet/lipgloss" +) + +func TestRenderBarWidth(t *testing.T) { + // The bar must always occupy exactly barWidth cells regardless of ratio, + // including the partial-cell and out-of-range cases. + for _, ratio := range []float64{-0.5, 0, 0.0001, 0.123, 0.5, 0.52, 0.999, 1, 1.5} { + if w := lipgloss.Width(renderBar(ratio)); w != barWidth { + t.Errorf("renderBar(%v) width = %d, want %d", ratio, w, barWidth) + } + } +} + +func TestRateMeter(t *testing.T) { + var m rateMeter + base := time.Unix(0, 0) + + if got := m.observe(base, 0); got != 0 { + t.Fatalf("first observe = %v, want 0", got) + } + // A second sample at the same instant spans no interval to average over. + if got := m.observe(base, 1<<20); got != 0 { + t.Fatalf("zero-dt observe = %v, want 0 (no positive interval)", got) + } + + // A steady 1 MiB/s stream averages to 1 MiB/s. + var rate float64 + for i := 1; i <= 50; i++ { + rate = m.observe(base.Add(time.Duration(i)*time.Second), int64(i)<<20) + } + const want = float64(1 << 20) + if rate < want*0.99 || rate > want*1.01 { + t.Errorf("steady rate = %v, want ~%v", rate, want) + } +} + +func TestRateMeterSmoothsBursts(t *testing.T) { + var m rateMeter + base := time.Unix(0, 0) + + // Prime a steady 1 MiB/s stream over the full window. + for i := range 6 { + m.observe(base.Add(time.Duration(i)*time.Second), int64(i)<<20) + } + + // A large part lands almost instantly: 5 MiB in 100ms. A per-sample EMA + // would read this as a ~50 MiB/s spike; the rolling window keeps the + // reported rate near the windowed throughput. + rate := m.observe(base.Add(5100*time.Millisecond), 10<<20) + if rate > 3<<20 { + t.Errorf("post-burst rate = %.0f B/s, want smoothed under %d B/s", rate, 3<<20) + } +} + +func TestFormatSpeed(t *testing.T) { + cases := []struct { + bps float64 + want string + }{ + {0, "0 B/s"}, + {-5, "0 B/s"}, + {512, "512 B/s"}, + {48 << 20, "48.0 MiB/s"}, + } + for _, tc := range cases { + if got := formatSpeed(tc.bps); got != tc.want { + t.Errorf("formatSpeed(%v) = %q, want %q", tc.bps, got, tc.want) + } + } +} + +func TestFormatETA(t *testing.T) { + cases := []struct { + remaining int64 + bps float64 + want string + }{ + {0, 0, "--:--"}, + {1 << 20, 0, "--:--"}, + {-1, 1 << 20, "--:--"}, + {0, 1 << 20, "00:00"}, + {1 << 20, 1 << 20, "00:01"}, + {120 << 20, 1 << 20, "02:00"}, + } + for _, tc := range cases { + if got := formatETA(tc.remaining, tc.bps); got != tc.want { + t.Errorf("formatETA(%d, %v) = %q, want %q", tc.remaining, tc.bps, got, tc.want) + } + } +} diff --git a/experimental/air/cmd/snapshot.go b/experimental/air/cmd/snapshot.go index aba67f6109c..59041f5986d 100644 --- a/experimental/air/cmd/snapshot.go +++ b/experimental/air/cmd/snapshot.go @@ -42,7 +42,7 @@ func snapshotCodeSource(ctx context.Context, w *databricks.WorkspaceClient, snap return snapshotResult{}, err } - up, err := newSnapshotUploader(w, snap, userDir, funcDir, filepath.Base(repoPath)) + up, err := newSnapshotUploader(ctx, w, snap, userDir, funcDir, filepath.Base(repoPath)) if err != nil { return snapshotResult{}, err } @@ -255,7 +255,7 @@ func fileExists(ctx context.Context, store filer.Filer, name string) (bool, erro // newSnapshotUploader builds the uploader for a submission. The tarball store is a // Volume (when remote_volume is set) or the user's repo_snapshots workspace dir; // sidecars always go to the run's funcDir in the workspace. -func newSnapshotUploader(w *databricks.WorkspaceClient, snap *snapshotSourceConfig, userDir, funcDir, dirName string) (snapshotUploader, error) { +func newSnapshotUploader(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, userDir, funcDir, dirName string) (snapshotUploader, error) { sidecarStore, err := filer.NewWorkspaceFilesClient(w, funcDir) if err != nil { return snapshotUploader{}, err @@ -263,7 +263,7 @@ func newSnapshotUploader(w *databricks.WorkspaceClient, snap *snapshotSourceConf if snap.RemoteVolume != nil { tarBase := strings.TrimRight(*snap.RemoteVolume, "/") - tarStore, err := filer.NewFilesClient(w, tarBase) + tarStore, err := filer.NewFilesClient(ctx, w, tarBase) if err != nil { return snapshotUploader{}, err } diff --git a/go.mod b/go.mod index 7d924904c1f..305c226651e 100644 --- a/go.mod +++ b/go.mod @@ -45,6 +45,14 @@ require ( gopkg.in/ini.v1 v1.67.3 // Apache-2.0 ) +require github.com/databricks/sdk-go/core v0.0.1-dev // Apache-2.0 + +require ( + github.com/databricks/sdk-go/auth v0.0.0-dev // Apache-2.0 + github.com/databricks/sdk-go/files v0.0.0-dev.1 // Apache-2.0 + github.com/databricks/sdk-go/options v0.0.0-dev // Apache-2.0 +) + require ( cloud.google.com/go/auth v0.18.1 // indirect cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect diff --git a/go.sum b/go.sum index c7cb717757c..6a56d025ddf 100644 --- a/go.sum +++ b/go.sum @@ -71,6 +71,14 @@ github.com/danieljoos/wincred v1.2.3 h1:v7dZC2x32Ut3nEfRH+vhoZGvN72+dQ/snVXo/vMF github.com/danieljoos/wincred v1.2.3/go.mod h1:6qqX0WNrS4RzPZ1tnroDzq9kY3fu1KwE7MRLQK4X0bs= github.com/databricks/databricks-sdk-go v0.160.0 h1:vwgT/11y2vMw41BxcKbUUqarg45lmoEdukk9yYJg5AM= github.com/databricks/databricks-sdk-go v0.160.0/go.mod h1:C5LNgGe6hGuRrTwoxFmuup3XtQQEaqtq0e+K8IFDIS4= +github.com/databricks/sdk-go/auth v0.0.0-dev h1:SyRGZvExH5TG9ynZEqNn+m+xTY42VOXr2GOxgNOQexk= +github.com/databricks/sdk-go/auth v0.0.0-dev/go.mod h1:Tj09W13MScUaix94cR0He9msAwe40JAQKuS2jP/ofQA= +github.com/databricks/sdk-go/core v0.0.1-dev h1:sIA1hCEJyl8/012vOBUIkDjsycNCD8RwfP76SCM0QbQ= +github.com/databricks/sdk-go/core v0.0.1-dev/go.mod h1:7Ckau34bOsaZhHJE6r7ORNa+/kO33jIs6DfrJL6UbsM= +github.com/databricks/sdk-go/files v0.0.0-dev.1 h1:MAcqxXpTFuM8dV+Yeb+X2+2iQlXWV2MVYEsYKle09+0= +github.com/databricks/sdk-go/files v0.0.0-dev.1/go.mod h1:4xc94QVa2BnkqR3eVC4o73phlwPgB/N2EAPKlK/2/zQ= +github.com/databricks/sdk-go/options v0.0.0-dev h1:+3bgKy5OH6G8VqzAox0eh0Ca28eQnoi9qDi5j/deiRY= +github.com/databricks/sdk-go/options v0.0.0-dev/go.mod h1:+lMasXZ/AfRAUYFNC8KFuT2vCJoF+TYNZ7XE0t26zfk= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= diff --git a/integration/cmd/fs/helpers_test.go b/integration/cmd/fs/helpers_test.go index e1bebb28f42..0f16ac5fdb8 100644 --- a/integration/cmd/fs/helpers_test.go +++ b/integration/cmd/fs/helpers_test.go @@ -30,14 +30,14 @@ func setupDbfsFiler(t testutil.TestingT) (filer.Filer, string) { } func setupUcVolumesFiler(t testutil.TestingT) (filer.Filer, string) { - _, wt := acc.WorkspaceTest(t) + ctx, wt := acc.WorkspaceTest(t) if os.Getenv("TEST_METASTORE_ID") == "" { t.Skip("Skipping tests that require a UC Volume when metastore id is not set.") } tmpdir := acc.TemporaryVolume(wt) - f, err := filer.NewFilesClient(wt.W, tmpdir) + f, err := filer.NewFilesClient(ctx, wt.W, tmpdir) require.NoError(t, err) return f, path.Join("dbfs:/", tmpdir) diff --git a/integration/libs/filer/helpers_test.go b/integration/libs/filer/helpers_test.go index ead83d66b08..baca230b791 100644 --- a/integration/libs/filer/helpers_test.go +++ b/integration/libs/filer/helpers_test.go @@ -58,14 +58,14 @@ func setupDbfsFiler(t testutil.TestingT) (filer.Filer, string) { } func setupUcVolumesFiler(t testutil.TestingT) (filer.Filer, string) { - _, wt := acc.WorkspaceTest(t) + ctx, wt := acc.WorkspaceTest(t) if os.Getenv("TEST_METASTORE_ID") == "" { t.Skip("Skipping tests that require a UC Volume when metastore id is not set.") } tmpdir := acc.TemporaryVolume(wt) - f, err := filer.NewFilesClient(wt.W, tmpdir) + f, err := filer.NewFilesClient(ctx, wt.W, tmpdir) require.NoError(t, err) return f, path.Join("dbfs:/", tmpdir) diff --git a/internal/build/notice_test.go b/internal/build/notice_test.go index 972c7aa9c0c..ea71f1f77a8 100644 --- a/internal/build/notice_test.go +++ b/internal/build/notice_test.go @@ -22,6 +22,10 @@ var moduleToGitHub = map[string]string{ // Modules excluded from NOTICE requirements (Databricks-owned). var noticeExclude = map[string]bool{ "github.com/databricks/databricks-sdk-go": true, + "github.com/databricks/sdk-go/core": true, + "github.com/databricks/sdk-go/auth": true, + "github.com/databricks/sdk-go/files": true, + "github.com/databricks/sdk-go/options": true, } // Additional entries required in the NOTICE file that are not direct go.mod diff --git a/libs/filer/files_client.go b/libs/filer/files_client.go index 232cbe20b9c..4351a3023b9 100644 --- a/libs/filer/files_client.go +++ b/libs/filer/files_client.go @@ -4,26 +4,39 @@ import ( "cmp" "context" "errors" - "fmt" "io" "io/fs" - "maps" "net/http" - "net/url" "path" "slices" - "strings" "time" "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/env" "github.com/databricks/databricks-sdk-go" - "github.com/databricks/databricks-sdk-go/apierr" - "github.com/databricks/databricks-sdk-go/client" - "github.com/databricks/databricks-sdk-go/listing" - "github.com/databricks/databricks-sdk-go/service/files" + "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/sdk-go/core/apierr" + files "github.com/databricks/sdk-go/files/v2" + "github.com/databricks/sdk-go/options/client" "golang.org/x/sync/errgroup" ) +// cloudResponseHeaderTimeout bounds the wait for a cloud storage response header +// on a large-file part transfer. It matches the files/v2 engine's own default. +const cloudResponseHeaderTimeout = 60 * time.Second + +// httpStatus returns the HTTP status code of err if it is (or wraps) an +// [apierr.APIError], and -1 otherwise. The filer keys its error mapping off the +// HTTP status rather than the SDK's canonical codes.Code: the Files API reports +// "path already exists" as 409 Conflict, which the SDK maps to codes.Aborted +// (not codes.AlreadyExists), so matching on the code would miss it. +func httpStatus(err error) int { + if aerr, ok := errors.AsType[*apierr.APIError](err); ok { + return aerr.HTTPStatusCode() + } + return -1 +} + // As of 19th Feb 2024, the Files API backend has a rate limit of 10 concurrent // requests and 100 QPS. We limit the number of concurrent requests to 5 to // avoid hitting the rate limit. @@ -88,112 +101,225 @@ func (e filesApiDirEntry) Info() (fs.FileInfo, error) { return e.i, nil } +// defaultUploadConcurrency bounds large-file (multipart) transfers for a filer +// created without an explicit concurrency. fs cp overrides it from --concurrency. +const defaultUploadConcurrency = 64 + +// multipartUploadEnvVar gates whether large Volumes writes are split into parts +// by the files/v2 upload engine. The engine is new, so multipart is off by +// default: when unset or not truthy, Write sends a single-shot PUT (via the +// files/v2 UploadFile endpoint), leaving fs cp and bundle behavior unchanged. +const multipartUploadEnvVar = "DATABRICKS_EXPERIMENTAL_MULTIPART_UPLOAD" + +// MultipartUploadEnabled reports whether large files written to UC Volumes are +// split into parts by the files/v2 upload engine, gated by +// DATABRICKS_EXPERIMENTAL_MULTIPART_UPLOAD (off by default). Both paths use the +// files/v2 client; the flag only selects the multipart engine over a single-shot +// PUT. +func MultipartUploadEnabled(ctx context.Context) bool { + enabled, _ := env.GetBool(ctx, multipartUploadEnvVar) + return enabled +} + +// uploadProgressKey is the context key for an optional large-file upload +// progress callback. +type uploadProgressKey struct{} + +// WithUploadProgress returns a context carrying a progress callback for +// large-file (multipart) uploads. FilesClient.Write forwards it to the upload +// engine; it has no effect on writes that do not go through the engine (small +// files, non-seekable streams, non-Volumes targets, or when multipart upload is +// disabled). +func WithUploadProgress(ctx context.Context, fn files.ProgressFunc) context.Context { + return context.WithValue(ctx, uploadProgressKey{}, fn) +} + +func uploadProgressFromContext(ctx context.Context) files.ProgressFunc { + fn, _ := ctx.Value(uploadProgressKey{}).(files.ProgressFunc) + return fn +} + // FilesClient implements the [Filer] interface for the Files API backend. type FilesClient struct { - workspaceClient *databricks.WorkspaceClient - apiClient *client.DatabricksClient + client *files.Client // File operations will be relative to this path. root WorkspaceRootPath + + // Large files are uploaded with the multipart engine. The limiter and transfer + // client are shared across every Write on this filer, so concurrent uploads + // (e.g. fs cp -r) draw from one bounded budget and one connection pool. + limiter files.Limiter + transferClient *http.Client +} + +type filesClientConfig struct { + uploadConcurrency int } -func NewFilesClient(w *databricks.WorkspaceClient, root string) (Filer, error) { - apiClient, err := client.New(w.Config) +// FilesClientOption configures a [FilesClient]. +type FilesClientOption func(*filesClientConfig) + +// WithUploadConcurrency sizes the shared limiter and transfer-client connection +// pool used for large-file (multipart) writes on the filer. +func WithUploadConcurrency(n int) FilesClientOption { + return func(c *filesClientConfig) { c.uploadConcurrency = n } +} + +func NewFilesClient(ctx context.Context, w *databricks.WorkspaceClient, root string, opts ...FilesClientOption) (Filer, error) { + cfg := filesClientConfig{uploadConcurrency: defaultUploadConcurrency} + for _, o := range opts { + o(&cfg) + } + + c, err := newFilesAPIClient(ctx, w.Config) if err != nil { return nil, err } return &FilesClient{ - workspaceClient: w, - apiClient: apiClient, + client: c, root: NewWorkspaceRootPath(root), + + limiter: files.NewLimiter(cfg.uploadConcurrency), + transferClient: newTransferClient(cfg.uploadConcurrency), }, nil } -func (w *FilesClient) urlPath(name string) (string, string, error) { - absPath, err := w.root.Join(name) - if err != nil { - return "", "", err - } - - // The user specified part of the path must be escaped. - urlPath := "/api/2.0/fs/files/" + url.PathEscape(strings.TrimLeft(absPath, "/")) +// newTransferClient returns an HTTP client for the cloud-leg part transfers of a +// large-file upload, sized for n concurrent transfers so idle connections are +// reused rather than re-dialed (Go's default of 2 per host would force +// re-dialing). It attaches no Databricks credentials (presigned URLs are +// self-authenticating) and sets no whole-request timeout, which would abort a +// legitimately long transfer; the upload is bounded by its context instead. +// files/v2 builds an equivalent client internally when one is not supplied; this +// filer supplies a shared one so all writes on it draw from a single pool. +func newTransferClient(n int) *http.Client { + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.ResponseHeaderTimeout = cloudResponseHeaderTimeout + transport.MaxIdleConnsPerHost = n + transport.MaxIdleConns = max(transport.MaxIdleConns, n) + return &http.Client{Transport: transport} +} - return absPath, urlPath, nil +// newFilesAPIClient builds the files/v2 client from the CLI's already resolved +// config, reusing its auth instead of re-reading a profile. cfg is passed by +// pointer because config.Config embeds a sync.Mutex and must not be copied. +func newFilesAPIClient(ctx context.Context, cfg *config.Config) (*files.Client, error) { + copts := []client.Option{ + client.WithHost(cfg.Host), + client.WithCredentials(configCredentials{cfg: cfg}), + client.WithoutProfileResolution(), + } + // The workspace routing header is needed on unified ("SPOG") hosts; the CLI's + // "none" sentinel means "no workspace ID", so it is not forwarded. + if id := cfg.WorkspaceID; id != "" && id != auth.WorkspaceIDNone { + copts = append(copts, client.WithWorkspaceID(id)) + } + return files.NewClient(ctx, copts...) } func (w *FilesClient) Write(ctx context.Context, name string, reader io.Reader, mode ...WriteMode) error { - absPath, urlPath, err := w.urlPath(name) + absPath, err := w.root.Join(name) if err != nil { return err } // Check that target path exists if CreateParentDirectories mode is not set if !slices.Contains(mode, CreateParentDirectories) { - err := w.workspaceClient.Files.GetDirectoryMetadataByDirectoryPath(ctx, path.Dir(absPath)) + dir := path.Dir(absPath) + _, err := w.client.GetDirectoryMetadata(ctx, &files.GetDirectoryMetadataRequest{DirectoryPath: &dir}) if err != nil { - aerr, ok := errors.AsType[*apierr.APIError](err) - if !ok { - return err + // This API returns a 404 if the directory doesn't exist. + if httpStatus(err) == http.StatusNotFound { + return noSuchDirectoryError{dir} } - - // This API returns a 404 if the file doesn't exist. - if aerr.StatusCode == http.StatusNotFound { - return noSuchDirectoryError{path.Dir(absPath)} - } - return err } } overwrite := slices.Contains(mode, OverwriteIfExists) - urlPath = fmt.Sprintf("%s?overwrite=%t", urlPath, overwrite) - headers := map[string]string{"Content-Type": "application/octet-stream"} - maps.Copy(headers, auth.WorkspaceIDHeaders(w.workspaceClient.Config)) - err = w.apiClient.Do(ctx, http.MethodPut, urlPath, headers, nil, reader, nil) - // Return early on success. + // When the multipart flag is enabled, seekable uploads go through the files/v2 + // upload engine, which sends small files in a single PUT and splits large ones + // into parts. Non-seekable streams and the flag-off case use the single-shot + // UploadFile endpoint below: the former to avoid buffering the whole stream in + // memory, the latter to keep the default behavior a plain PUT. The engine + // recovers an io.ReaderAt for concurrent positioned reads when the source + // provides one (a local file). + if MultipartUploadEnabled(ctx) && isSeekable(reader) { + opts := []files.UploadOption{ + files.WithOverwrite(overwrite), + files.WithLimiter(w.limiter), + files.WithTransferClient(w.transferClient), + } + // A caller (fs cp) can attach a progress callback via the context to + // render an upload bar; it is absent for other writers (e.g. bundle). + if fn := uploadProgressFromContext(ctx); fn != nil { + opts = append(opts, files.WithProgress(fn)) + } + _, uerr := w.client.Upload(ctx, absPath, reader, opts...) + return mapUploadError(uerr, absPath) + } + + _, err = w.client.UploadFile(ctx, &files.UploadFileRequest{ + FilePath: &absPath, + Contents: io.NopCloser(reader), + Overwrite: &overwrite, + }) if err == nil { return nil } - // Special handling of this error only if it is an API error. - aerr, ok := errors.AsType[*apierr.APIError](err) + // This API returns 409 if a file already exists at the path. + if httpStatus(err) == http.StatusConflict { + return fileAlreadyExistsError{absPath} + } + return err +} + +// isSeekable reports whether r can be seeked, without moving it: a no-op +// Seek(0, io.SeekCurrent) returns the current offset for a working seeker and +// errors for a broken one. Probing this way (rather than seeking to the end and +// back) guarantees the reader keeps its position, so a false result never leaves +// it parked at EOF for the single-shot fallback, which reads from the current +// offset and would otherwise upload a truncated object. The engine sizes the +// stream itself once it takes over. +func isSeekable(r io.Reader) bool { + s, ok := r.(io.Seeker) if !ok { - return err + return false } + _, err := s.Seek(0, io.SeekCurrent) + return err == nil +} - // This API returns 409 if the file already exists, when the object type is file - if aerr.StatusCode == http.StatusConflict && aerr.ErrorCode == "ALREADY_EXISTS" { +// mapUploadError translates the upload engine's already-exists sentinel into the +// filer's error so skip-if-exists logic (which checks fs.ErrExist) keeps working. +// A nil error passes through unchanged. +func mapUploadError(err error, absPath string) error { + if errors.Is(err, files.ErrAlreadyExists) { return fileAlreadyExistsError{absPath} } - return err } func (w *FilesClient) Read(ctx context.Context, name string) (io.ReadCloser, error) { - absPath, urlPath, err := w.urlPath(name) + absPath, err := w.root.Join(name) if err != nil { return nil, err } - var reader io.ReadCloser - err = w.apiClient.Do(ctx, http.MethodGet, urlPath, auth.WorkspaceIDHeaders(w.workspaceClient.Config), nil, nil, &reader) + resp, err := w.client.DownloadFile(ctx, &files.DownloadFileRequest{FilePath: &absPath}) // Return early on success. if err == nil { - return reader, nil - } - - // Special handling of this error only if it is an API error. - aerr, ok := errors.AsType[*apierr.APIError](err) - if !ok { - return nil, err + return resp.Contents, nil } // This API returns a 404 if the specified path does not exist. - if aerr.StatusCode == http.StatusNotFound { + if httpStatus(err) == http.StatusNotFound { // Check if the path is a directory. If so, return not a file error. if _, err := w.statDir(ctx, name); err == nil { return nil, notAFile{absPath} @@ -217,21 +343,15 @@ func (w *FilesClient) deleteFile(ctx context.Context, name string) error { return cannotDeleteRootError{} } - err = w.workspaceClient.Files.DeleteByFilePath(ctx, absPath) + _, err = w.client.DeleteFile(ctx, &files.DeleteFileRequest{FilePath: &absPath}) // Return early on success. if err == nil { return nil } - // Special handling of this error only if it is an API error. - aerr, ok := errors.AsType[*apierr.APIError](err) - if !ok { - return err - } - // This files delete API returns a 404 if the specified path does not exist. - if aerr.StatusCode == http.StatusNotFound { + if httpStatus(err) == http.StatusNotFound { return fileDoesNotExistError{absPath} } @@ -249,27 +369,20 @@ func (w *FilesClient) deleteDirectory(ctx context.Context, name string) error { return cannotDeleteRootError{} } - err = w.workspaceClient.Files.DeleteDirectoryByDirectoryPath(ctx, absPath) + _, err = w.client.DeleteDirectory(ctx, &files.DeleteDirectoryRequest{DirectoryPath: &absPath}) - // Special handling of this error only if it is an API error. - aerr, ok := errors.AsType[*apierr.APIError](err) - if !ok { - return err + // Return early on success. + if err == nil { + return nil } - // The directory delete API returns a 400 if the directory is not empty - if aerr.StatusCode == http.StatusBadRequest { - var reasons []string - details := aerr.ErrorDetails() - if details.ErrorInfo != nil { - reasons = append(reasons, details.ErrorInfo.Reason) - } - // Error code 400 is generic and can be returned for other reasons. Make - // sure one of the reasons for the error is that the directory is not empty. - if !slices.Contains(reasons, "FILES_API_DIRECTORY_IS_NOT_EMPTY") { - return err + // The directory delete API returns a 400 if the directory is not empty. That + // status is generic, so confirm the specific reason before mapping it. + if aerr, ok := errors.AsType[*apierr.APIError](err); ok && aerr.HTTPStatusCode() == http.StatusBadRequest { + if info := aerr.Details().ErrorInfo; info != nil && info.Reason == "FILES_API_DIRECTORY_IS_NOT_EMPTY" { + return directoryNotEmptyError{absPath} } - return directoryNotEmptyError{absPath} + return err } // On GCS-backed storage a directory created implicitly is just a key prefix @@ -277,7 +390,7 @@ func (w *FilesClient) deleteDirectory(ctx context.Context, name string) error { // deleted. The delete API then returns 404 for that already-vanished // directory; treat it as a not-found error so recursive delete can consider // its goal satisfied. - if apierr.IsMissing(err) { + if httpStatus(err) == http.StatusNotFound { return noSuchDirectoryError{absPath} } return err @@ -377,48 +490,37 @@ func (w *FilesClient) ReadDir(ctx context.Context, name string) ([]fs.DirEntry, return nil, err } - iter := w.workspaceClient.Files.ListDirectoryContents(ctx, files.ListDirectoryContentsRequest{ - DirectoryPath: absPath, - }) - - files, err := listing.ToSlice(ctx, iter) - - // Return early on success. - if err == nil { - entries := make([]fs.DirEntry, len(files)) - for i, file := range files { - entries[i] = filesApiDirEntry{ - i: filesApiFileInfo{ - absPath: file.Path, - isDir: file.IsDirectory, - fileSize: file.FileSize, - lastModified: file.LastModified, - }, + var entries []fs.DirEntry + for entry, err := range w.client.ListDirectoryContentsIter(ctx, &files.ListDirectoryContentsRequest{ + DirectoryPath: &absPath, + }) { + if err != nil { + // This API returns a 404 if the specified path does not exist. + if httpStatus(err) == http.StatusNotFound { + // Check if the path is a file. If so, return not a directory error. + if _, ferr := w.statFile(ctx, name); ferr == nil { + return nil, notADirectory{absPath} + } + + // No file or directory exists at the specified path. Return no such directory error. + return nil, noSuchDirectoryError{absPath} } + return nil, err } - // Sort by name for parity with os.ReadDir. - slices.SortFunc(entries, func(a, b fs.DirEntry) int { return cmp.Compare(a.Name(), b.Name()) }) - return entries, nil - } - - // Special handling of this error only if it is an API error. - aerr, ok := errors.AsType[*apierr.APIError](err) - if !ok { - return nil, err + entries = append(entries, filesApiDirEntry{ + i: filesApiFileInfo{ + absPath: derefString(entry.Path), + isDir: derefBool(entry.IsDirectory), + fileSize: int64(derefInt(entry.FileSize)), + lastModified: int64(derefInt(entry.LastModified)), + }, + }) } - // This API returns a 404 if the specified path does not exist. - if aerr.StatusCode == http.StatusNotFound { - // Check if the path is a file. If so, return not a directory error. - if _, err := w.statFile(ctx, name); err == nil { - return nil, notADirectory{absPath} - } - - // No file or directory exists at the specified path. Return no such directory error. - return nil, noSuchDirectoryError{absPath} - } - return nil, err + // Sort by name for parity with os.ReadDir. + slices.SortFunc(entries, func(a, b fs.DirEntry) int { return cmp.Compare(a.Name(), b.Name()) }) + return entries, nil } func (w *FilesClient) Mkdir(ctx context.Context, name string) error { @@ -427,12 +529,11 @@ func (w *FilesClient) Mkdir(ctx context.Context, name string) error { return err } - err = w.workspaceClient.Files.CreateDirectory(ctx, files.CreateDirectoryRequest{ - DirectoryPath: absPath, - }) + _, err = w.client.CreateDirectory(ctx, &files.CreateDirectoryRequest{DirectoryPath: &absPath}) - // Special handling of this error only if it is an API error. - if aerr, ok := errors.AsType[*apierr.APIError](err); ok && aerr.StatusCode == http.StatusConflict { + // This API returns a 409 when a file already exists at the path (the create + // is not idempotent over a file). + if httpStatus(err) == http.StatusConflict { return fileAlreadyExistsError{absPath} } @@ -446,25 +547,19 @@ func (w *FilesClient) statFile(ctx context.Context, name string) (fs.FileInfo, e return nil, err } - fileInfo, err := w.workspaceClient.Files.GetMetadataByFilePath(ctx, absPath) + resp, err := w.client.GetFileMetadata(ctx, &files.GetFileMetadataRequest{FilePath: &absPath}) // If the HEAD requests succeeds, the file exists. if err == nil { return filesApiFileInfo{ absPath: absPath, isDir: false, - fileSize: fileInfo.ContentLength, + fileSize: derefInt64(resp.ContentLength), }, nil } - // Special handling of this error only if it is an API error. - aerr, ok := errors.AsType[*apierr.APIError](err) - if !ok { - return nil, err - } - // This API returns a 404 if the specified path does not exist. - if aerr.StatusCode == http.StatusNotFound { + if httpStatus(err) == http.StatusNotFound { return nil, fileDoesNotExistError{absPath} } @@ -478,21 +573,15 @@ func (w *FilesClient) statDir(ctx context.Context, name string) (fs.FileInfo, er return nil, err } - err = w.workspaceClient.Files.GetDirectoryMetadataByDirectoryPath(ctx, absPath) + _, err = w.client.GetDirectoryMetadata(ctx, &files.GetDirectoryMetadataRequest{DirectoryPath: &absPath}) // If the HEAD requests succeeds, the directory exists. if err == nil { return filesApiFileInfo{absPath: absPath, isDir: true}, nil } - // Special handling of this error only if it is an API error. - aerr, ok := errors.AsType[*apierr.APIError](err) - if !ok { - return nil, err - } - // The directory metadata API returns a 404 if the specified path does not exist. - if aerr.StatusCode == http.StatusNotFound { + if httpStatus(err) == http.StatusNotFound { return nil, noSuchDirectoryError{absPath} } @@ -516,3 +605,31 @@ func (w *FilesClient) Stat(ctx context.Context, name string) (fs.FileInfo, error // Since the path is not a directory, assume that it is a file and issue a stat call. return w.statFile(ctx, name) } + +func derefString(s *string) string { + if s == nil { + return "" + } + return *s +} + +func derefBool(b *bool) bool { + if b == nil { + return false + } + return *b +} + +func derefInt(i *int) int { + if i == nil { + return 0 + } + return *i +} + +func derefInt64(i *int64) int64 { + if i == nil { + return 0 + } + return *i +} diff --git a/libs/filer/files_client_auth.go b/libs/filer/files_client_auth.go new file mode 100644 index 00000000000..1d28d1cacdf --- /dev/null +++ b/libs/filer/files_client_auth.go @@ -0,0 +1,42 @@ +package filer + +import ( + "context" + "net/http" + + "github.com/databricks/databricks-sdk-go/config" + sdkauth "github.com/databricks/sdk-go/auth" +) + +// configCredentials adapts the CLI's resolved SDK config into the credentials +// interface expected by the files/v2 client. It signs a throwaway +// request with config.Authenticate and hands the resulting headers to the +// client, so the files/v2 client reuses the exact auth the rest of the CLI +// already resolved instead of re-reading a profile. +type configCredentials struct { + cfg *config.Config +} + +func (c configCredentials) Name() string { + return "databricks-cli" +} + +func (c configCredentials) AuthHeaders(ctx context.Context) ([]sdkauth.Header, error) { + // The URL is irrelevant: config.Authenticate only reads it to pick the auth + // scheme, and every files/v2 request targets the same workspace host. + req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.cfg.Host, nil) + if err != nil { + return nil, err + } + if err := c.cfg.Authenticate(req); err != nil { + return nil, err + } + + headers := make([]sdkauth.Header, 0, len(req.Header)) + for key, values := range req.Header { + for _, value := range values { + headers = append(headers, sdkauth.Header{Key: key, Value: value}) + } + } + return headers, nil +} diff --git a/libs/filer/files_client_test.go b/libs/filer/files_client_test.go index 57b93f7c6cf..484715bcd96 100644 --- a/libs/filer/files_client_test.go +++ b/libs/filer/files_client_test.go @@ -1,11 +1,17 @@ package filer import ( + "bytes" + "errors" + "fmt" + "io" "io/fs" "testing" + "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" + files "github.com/databricks/sdk-go/files/v2" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -37,7 +43,7 @@ func deleteDirectoryWithError(t *testing.T, statusCode int, errorCode, reason st }) require.NoError(t, err) - f, err := NewFilesClient(client, "/test") + f, err := NewFilesClient(t.Context(), client, "/test") require.NoError(t, err) return f.(*FilesClient).deleteDirectory(t.Context(), "dir") @@ -55,3 +61,114 @@ func TestFilesClientDeleteDirectoryNotEmpty(t *testing.T) { err := deleteDirectoryWithError(t, 400, "INVALID_PARAMETER_VALUE", "FILES_API_DIRECTORY_IS_NOT_EMPTY") assert.ErrorIs(t, err, fs.ErrInvalid) } + +func newTestFilesClient(t *testing.T) Filer { + t.Helper() + + server := testserver.New(t) + testserver.AddDefaultHandlers(server) + + client, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: server.URL, + Token: "testtoken", + }) + require.NoError(t, err) + + f, err := NewFilesClient(t.Context(), client, "/") + require.NoError(t, err) + return f +} + +func TestFilesClientMkdirWhenFileExists(t *testing.T) { + // The Files API reports "a file already exists at this path" as a 409, which + // the SDK maps to codes.Aborted (not codes.AlreadyExists); the filer keys off + // the HTTP status so it still surfaces as fs.ErrExist. + ctx := t.Context() + f := newTestFilesClient(t) + + require.NoError(t, f.Mkdir(ctx, "/Volumes/main/schema/vol")) + require.NoError(t, f.Write(ctx, "/Volumes/main/schema/vol/hello", bytes.NewReader([]byte("abc")))) + + err := f.Mkdir(ctx, "/Volumes/main/schema/vol/hello") + assert.ErrorIs(t, err, fs.ErrExist) +} + +// onlyReader hides the Seek method of an underlying reader, modelling a +// non-seekable stream (e.g. a remote download body). +type onlyReader struct{ io.Reader } + +func TestIsSeekable(t *testing.T) { + in := []byte("hello, files API") + + if !isSeekable(bytes.NewReader(in)) { + t.Fatal("bytes.Reader should be seekable") + } + + // The position must be left at the start so a subsequent read covers every byte. + r := bytes.NewReader(in) + if !isSeekable(r) { + t.Fatal("bytes.Reader should be seekable") + } + b, err := io.ReadAll(r) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(b, in) { + t.Errorf("read %q after isSeekable, want %q (position not left at start)", b, in) + } + + if isSeekable(onlyReader{bytes.NewReader(in)}) { + t.Error("a non-seekable reader should report false") + } +} + +func TestMultipartUploadEnabled(t *testing.T) { + ctx := t.Context() + if MultipartUploadEnabled(ctx) { + t.Error("multipart upload must be disabled by default") + } + for _, on := range []string{"true", "1", "yes", "on"} { + if !MultipartUploadEnabled(env.Set(ctx, multipartUploadEnvVar, on)) { + t.Errorf("value %q should enable multipart upload", on) + } + } + for _, off := range []string{"false", "0", "", "nonsense"} { + if MultipartUploadEnabled(env.Set(ctx, multipartUploadEnvVar, off)) { + t.Errorf("value %q should not enable multipart upload", off) + } + } +} + +func TestMapUploadError(t *testing.T) { + const p = "/Volumes/c/s/v/f.bin" + + if err := mapUploadError(nil, p); err != nil { + t.Errorf("nil error should pass through, got %v", err) + } + + // The engine's already-exists sentinel (even wrapped) must surface as fs.ErrExist + // so skip-if-exists keeps working. + for _, in := range []error{ + files.ErrAlreadyExists, + fmt.Errorf("upload failed: %w", files.ErrAlreadyExists), + } { + got := mapUploadError(in, p) + if !errors.Is(got, fs.ErrExist) { + t.Errorf("mapUploadError(%v) = %v, want errors.Is fs.ErrExist", in, got) + } + } + + // Other errors pass through unchanged. + other := errors.New("boom") + if got := mapUploadError(other, p); got != other { + t.Errorf("mapUploadError(other) = %v, want it unchanged", got) + } +} + +func TestWithUploadConcurrency(t *testing.T) { + var cfg filesClientConfig + WithUploadConcurrency(64)(&cfg) + if cfg.uploadConcurrency != 64 { + t.Errorf("uploadConcurrency = %d, want 64", cfg.uploadConcurrency) + } +} diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index a51a13b4afe..4587137f668 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -219,6 +219,15 @@ func AddDefaultHandlers(server *Server) { defer req.Workspace.LockUnlock()() + // The API rejects creating a directory where a file already exists; it is + // idempotent only over directories. Mirror that so callers observe the 409. + if _, isFile := req.Workspace.files[dirPath]; isFile { + return Response{ + StatusCode: 409, + Body: map[string]string{"message": "The given path points to an existing file. This API does not support operations on files."}, + } + } + // Create directory and all parent directories. for dir := dirPath; dir != "/" && dir != ""; dir = path.Dir(dir) { if _, exists := req.Workspace.directories[dir]; !exists {