From 39fe34792a67c7126c547ebec012db13bacc1d19 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 22 Jul 2026 04:47:02 +0000 Subject: [PATCH] fs: reimplement UC Volumes filer on the files/v2 client Reworks the `fs` command's UC Volumes filer (also used by bundle library uploads) to run on the new `sdk-go/files/v2` client instead of hand-rolled Files API calls plus a separate CLI-side multipart upload engine. `files/v2` is not published to the Go proxy yet, so it is vendored under `libs/tmp/` along with the two unpublished sibling modules it needs (`auth` and `options`); the published `sdk-go/core` module stays a normal dependency. The only edits to the vendored copy are the import-path rewrite and two fixes to the generated client (see below). `libs/tmp/README.md` records the copy and the swap-and-delete steps for when the module is published. The filer now delegates upload, download, listing, mkdir, delete, and stat to `files/v2`, which brings its own multipart/resumable engine. The CLI's own `libs/upload` engine is therefore removed. `DATABRICKS_EXPERIMENTAL_MULTIPART_UPLOAD` still gates whether large Volumes uploads are split into parts, but now selects multipart-vs-single-shot within files/v2 rather than a separate engine. This pull request and its description were written by Isaac. Co-authored-by: Isaac --- .github/OWNERS | 1 + .golangci.yaml | 10 + bundle/libraries/filer_volume.go | 2 +- cmd/fs/cp.go | 73 +- cmd/fs/helpers.go | 20 +- cmd/fs/upload_progress.go | 256 ++++ cmd/fs/upload_progress_test.go | 96 ++ experimental/air/cmd/snapshot.go | 6 +- go.mod | 2 + go.sum | 2 + integration/cmd/fs/helpers_test.go | 4 +- integration/libs/filer/helpers_test.go | 4 +- internal/build/notice_test.go | 1 + libs/filer/files_client.go | 401 +++-- libs/filer/files_client_auth.go | 42 + libs/filer/files_client_test.go | 119 +- libs/testserver/handlers.go | 9 + libs/tmp/README.md | 48 + libs/tmp/auth/auth.go | 98 ++ libs/tmp/auth/auth_test.go | 16 + libs/tmp/auth/cache.go | 245 ++++ libs/tmp/auth/cache_test.go | 608 ++++++++ libs/tmp/auth/retrying.go | 66 + libs/tmp/auth/retrying_test.go | 150 ++ libs/tmp/files/internal/version.go | 5 + libs/tmp/files/v2/client.go | 1303 +++++++++++++++++ libs/tmp/files/v2/ext_limiter.go | 55 + libs/tmp/files/v2/ext_limiter_test.go | 187 +++ libs/tmp/files/v2/ext_multipart.go | 450 ++++++ libs/tmp/files/v2/ext_resumable.go | 187 +++ libs/tmp/files/v2/ext_resumable_test.go | 21 + libs/tmp/files/v2/ext_straggler.go | 81 + libs/tmp/files/v2/ext_straggler_test.go | 101 ++ libs/tmp/files/v2/ext_tunables.go | 171 +++ libs/tmp/files/v2/ext_upload.go | 568 +++++++ libs/tmp/files/v2/ext_upload_control.go | 367 +++++ libs/tmp/files/v2/ext_upload_control_test.go | 206 +++ libs/tmp/files/v2/ext_upload_helpers.go | 116 ++ libs/tmp/files/v2/ext_upload_test.go | 726 +++++++++ libs/tmp/files/v2/genhelper.go | 250 ++++ .../files/v2/internal/cloudstorage/body.go | 55 + .../v2/internal/cloudstorage/body_test.go | 57 + .../v2/internal/cloudstorage/cloudstorage.go | 320 ++++ .../cloudstorage/cloudstorage_test.go | 174 +++ .../files/v2/internal/cloudstorage/retries.go | 76 + .../v2/internal/cloudstorage/retries_test.go | 43 + libs/tmp/files/v2/model.go | 190 +++ libs/tmp/files/v2/wire.go | 854 +++++++++++ libs/tmp/options/call/call.go | 53 + libs/tmp/options/call/call_test.go | 58 + libs/tmp/options/client/client.go | 119 ++ libs/tmp/options/client/client_test.go | 57 + libs/tmp/options/internal/version.go | 5 + .../internaloptions/internaloptions.go | 113 ++ .../internaloptions/internaloptions_test.go | 41 + tools/check_deadcode.py | 1 + 56 files changed, 9121 insertions(+), 168 deletions(-) create mode 100644 cmd/fs/upload_progress.go create mode 100644 cmd/fs/upload_progress_test.go create mode 100644 libs/filer/files_client_auth.go create mode 100644 libs/tmp/README.md create mode 100644 libs/tmp/auth/auth.go create mode 100644 libs/tmp/auth/auth_test.go create mode 100644 libs/tmp/auth/cache.go create mode 100644 libs/tmp/auth/cache_test.go create mode 100644 libs/tmp/auth/retrying.go create mode 100644 libs/tmp/auth/retrying_test.go create mode 100644 libs/tmp/files/internal/version.go create mode 100644 libs/tmp/files/v2/client.go create mode 100644 libs/tmp/files/v2/ext_limiter.go create mode 100644 libs/tmp/files/v2/ext_limiter_test.go create mode 100644 libs/tmp/files/v2/ext_multipart.go create mode 100644 libs/tmp/files/v2/ext_resumable.go create mode 100644 libs/tmp/files/v2/ext_resumable_test.go create mode 100644 libs/tmp/files/v2/ext_straggler.go create mode 100644 libs/tmp/files/v2/ext_straggler_test.go create mode 100644 libs/tmp/files/v2/ext_tunables.go create mode 100644 libs/tmp/files/v2/ext_upload.go create mode 100644 libs/tmp/files/v2/ext_upload_control.go create mode 100644 libs/tmp/files/v2/ext_upload_control_test.go create mode 100644 libs/tmp/files/v2/ext_upload_helpers.go create mode 100644 libs/tmp/files/v2/ext_upload_test.go create mode 100644 libs/tmp/files/v2/genhelper.go create mode 100644 libs/tmp/files/v2/internal/cloudstorage/body.go create mode 100644 libs/tmp/files/v2/internal/cloudstorage/body_test.go create mode 100644 libs/tmp/files/v2/internal/cloudstorage/cloudstorage.go create mode 100644 libs/tmp/files/v2/internal/cloudstorage/cloudstorage_test.go create mode 100644 libs/tmp/files/v2/internal/cloudstorage/retries.go create mode 100644 libs/tmp/files/v2/internal/cloudstorage/retries_test.go create mode 100644 libs/tmp/files/v2/model.go create mode 100644 libs/tmp/files/v2/wire.go create mode 100644 libs/tmp/options/call/call.go create mode 100644 libs/tmp/options/call/call_test.go create mode 100644 libs/tmp/options/client/client.go create mode 100644 libs/tmp/options/client/client_test.go create mode 100644 libs/tmp/options/internal/version.go create mode 100644 libs/tmp/options/internaloptions/internaloptions.go create mode 100644 libs/tmp/options/internaloptions/internaloptions_test.go diff --git a/.github/OWNERS b/.github/OWNERS index 76bc081f0eb..865d29ff3f4 100644 --- a/.github/OWNERS +++ b/.github/OWNERS @@ -39,6 +39,7 @@ /cmd/sync/ team:platform /libs/filer/ team:platform /libs/sync/ team:platform +/libs/tmp/ team:platform # Core CLI infrastructure /cmd/root/ team:platform diff --git a/.golangci.yaml b/.golangci.yaml index 565f46a5ca8..e2f03cf9db3 100644 --- a/.golangci.yaml +++ b/.golangci.yaml @@ -163,6 +163,11 @@ linters: default-signifies-exhaustive: true exclusions: generated: lax + # libs/tmp is a temporary verbatim copy of the unpublished sdk-go/files/v2 + # module and its deps; keep it byte-for-byte with upstream rather than making + # it satisfy this repo's linters. Remove when the module is published. + paths: + - libs/tmp presets: - comments - common-false-positives @@ -199,3 +204,8 @@ formatters: extra-rules: true exclusions: generated: lax + # See the note on linters.exclusions.paths above: libs/tmp is a temporary + # verbatim vendored copy and is intentionally not reformatted to this repo's + # style. + paths: + - libs/tmp 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..47a5461b1eb --- /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/cli/libs/tmp/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..f2924fac769 100644 --- a/go.mod +++ b/go.mod @@ -45,6 +45,8 @@ require ( gopkg.in/ini.v1 v1.67.3 // Apache-2.0 ) +require github.com/databricks/sdk-go/core v0.0.0 // 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..a1994be2135 100644 --- a/go.sum +++ b/go.sum @@ -71,6 +71,8 @@ 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/core v0.0.0 h1:1Zg54LjihKnFTfcVwtEsbv/nm446QVOtenS8dSPjQ9U= +github.com/databricks/sdk-go/core v0.0.0/go.mod h1:7Ckau34bOsaZhHJE6r7ORNa+/kO33jIs6DfrJL6UbsM= 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..4adbc2670d0 100644 --- a/internal/build/notice_test.go +++ b/internal/build/notice_test.go @@ -22,6 +22,7 @@ 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, } // 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..cca9b7a47c6 100644 --- a/libs/filer/files_client.go +++ b/libs/filer/files_client.go @@ -4,26 +4,40 @@ 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" + files "github.com/databricks/cli/libs/tmp/files/v2" + "github.com/databricks/cli/libs/tmp/options/client" "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" "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; +// see the transfer-client note in libs/tmp/files/v2. +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 +102,226 @@ 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 +} + +// 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(w *databricks.WorkspaceClient, root string) (Filer, error) { - apiClient, err := client.New(w.Config) +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 vendored 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 file doesn't exist. - if aerr.StatusCode == http.StatusNotFound { - return noSuchDirectoryError{path.Dir(absPath)} + // This API returns a 404 if the directory doesn't exist. + if httpStatus(err) == http.StatusNotFound { + return noSuchDirectoryError{dir} } - 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 +345,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 +371,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) + // 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} } - // 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 - } - return directoryNotEmptyError{absPath} + return err } // On GCS-backed storage a directory created implicitly is just a key prefix @@ -277,7 +392,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 +492,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 +531,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 +549,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 +575,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 +607,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..5d0d30afae9 --- /dev/null +++ b/libs/filer/files_client_auth.go @@ -0,0 +1,42 @@ +package filer + +import ( + "context" + "net/http" + + tmpauth "github.com/databricks/cli/libs/tmp/auth" + "github.com/databricks/databricks-sdk-go/config" +) + +// configCredentials adapts the CLI's resolved SDK config into the credentials +// interface expected by the vendored 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) ([]tmpauth.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([]tmpauth.Header, 0, len(req.Header)) + for key, values := range req.Header { + for _, value := range values { + headers = append(headers, tmpauth.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..7008e6858a9 100644 --- a/libs/filer/files_client_test.go +++ b/libs/filer/files_client_test.go @@ -1,10 +1,16 @@ package filer import ( + "bytes" + "errors" + "fmt" + "io" "io/fs" "testing" + "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testserver" + files "github.com/databricks/cli/libs/tmp/files/v2" "github.com/databricks/databricks-sdk-go" "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 { diff --git a/libs/tmp/README.md b/libs/tmp/README.md new file mode 100644 index 00000000000..70337684528 --- /dev/null +++ b/libs/tmp/README.md @@ -0,0 +1,48 @@ +# libs/tmp + +Temporary in-repo copy of the not-yet-published `github.com/databricks/sdk-go` +`files/v2` client and the two sibling modules it depends on. It exists only so +the CLI can build against `files/v2` before that module is published to the +Databricks Go proxy. + +## Contents + +| Path | Upstream source (in the `universe` monorepo) | +| ------------------- | ------------------------------------------------ | +| `files/` | `deco/oss/repos/sdk-go/files` | +| `auth/` | `deco/oss/repos/sdk-go/auth` (top-level package) | +| `options/` | `deco/oss/repos/sdk-go/options` | + +Only the packages `files/v2` actually imports were copied: the top-level `auth` +package (not its `oidc`/`credentials`/`transport` subpackages) and +`options/{call,client,internaloptions,internal}`. The published +`github.com/databricks/sdk-go/core` module is a real dependency and is NOT +copied here; the vendored code imports it directly. + +The only edits to the copied sources are the import-path rewrite from +`github.com/databricks/sdk-go/{files,auth,options}` to +`github.com/databricks/cli/libs/tmp/{files,auth,options}` (`core` imports are +left unchanged), plus two local fixes to the generated +`files/v2/client.go` that the upstream generator still gets wrong: + +- Empty response bodies: the generated methods called `json.Unmarshal` on every + response body, which fails with "unexpected end of JSON input" on the HEAD + metadata calls and 204 PUT/DELETE responses that the Files API returns with no + body. They now go through `unmarshalResponse` (in `files/v2/genhelper.go`), + which treats an empty body as a zero-value response. +- HTTP method casing: the two directory/file metadata calls used the method + string `"head"`; Go sends the method verbatim, so it is now `"HEAD"`. + +Both are bugs in the upstream generator; report them there so the published +module needs no patching, and drop these edits when swapping back. + +## Removing this copy + +When `github.com/databricks/sdk-go/files/v2` is published: + +1. Add it (and `auth`/`options` if they publish separately) to `go.mod`. +2. Rewrite every `github.com/databricks/cli/libs/tmp/{files,auth,options}` + import back to `github.com/databricks/sdk-go/{files,auth,options}`. The + `files/v2` package keeps its `v2` path element, so consumers only change the + module prefix. +3. Delete `libs/tmp/`. diff --git a/libs/tmp/auth/auth.go b/libs/tmp/auth/auth.go new file mode 100644 index 00000000000..fac53e7db42 --- /dev/null +++ b/libs/tmp/auth/auth.go @@ -0,0 +1,98 @@ +// Package auth is an internal package that provides authentication utilities. +// +// IMPORTANT: This package is not meant to be used directly by consumers of the +// SDK and is subject to change without notice. +package auth + +import ( + "context" + "fmt" + "time" +) + +// Header represents a header that can be used to sign requests. +type Header struct { + Key string + Value string +} + +// Credentials represents a named source of authentication headers. +type Credentials interface { + Name() string + + AuthHeaders(context.Context) ([]Header, error) +} + +// Token represents a token that can be used to sign requests. +type Token struct { + // Value is the raw value to sign requests with. It typically is an + // access token but can represent other types of tokens (e.g. ID tokens). + Value string + + // Type is the type of token. If Type is empty, the token type is + // assumed to be "Bearer". + Type string + + // Expiry is the time at which the token expires. If Expiry is zero, the + // token is considered to be valid indefinitely. + Expiry time.Time +} + +// A TokenProvider is anything that can return a token. +type TokenProvider interface { + // Token returns a token or an error. The returned Token should be + // considered immutable and should not be modified. + Token(context.Context) (*Token, error) +} + +// TokenProviderFn is an adapter to allow the use of ordinary functions as +// TokenProvider. +// +// Example: +// +// ts := TokenProviderFn(func(ctx context.Context) (*Token, error) { +// return &Token{}, nil +// }) +type TokenProviderFn func(context.Context) (*Token, error) + +func (fn TokenProviderFn) Token(ctx context.Context) (*Token, error) { + return fn(ctx) +} + +type TokenCredentials interface { + TokenProvider + Credentials +} + +// NewTokenCredentials returns TokenCredentials with the given name that use the +// given TokenProvider to return authentication headers. +func NewTokenCredentials(name string, p TokenProvider) TokenCredentials { + return &tokenCredentials{name: name, p: p} +} + +type tokenCredentials struct { + name string + p TokenProvider +} + +func (c *tokenCredentials) Name() string { + return c.name +} + +func (c *tokenCredentials) AuthHeaders(ctx context.Context) ([]Header, error) { + token, err := c.p.Token(ctx) + if err != nil { + return nil, err + } + tt := token.Type + if tt == "" { + tt = "Bearer" + } + return []Header{ + {Key: "Authorization", Value: fmt.Sprintf("%s %s", tt, token.Value)}, + }, nil +} + +func (c *tokenCredentials) Token(ctx context.Context) (*Token, error) { + return c.p.Token(ctx) +} diff --git a/libs/tmp/auth/auth_test.go b/libs/tmp/auth/auth_test.go new file mode 100644 index 00000000000..1294e404e44 --- /dev/null +++ b/libs/tmp/auth/auth_test.go @@ -0,0 +1,16 @@ +package auth + +import ( + "context" + "testing" +) + +func TestNewTokenCredentials_name(t *testing.T) { + credentials := NewTokenCredentials("test-strategy", TokenProviderFn(func(context.Context) (*Token, error) { + return &Token{Value: "token"}, nil + })) + + if got, want := credentials.Name(), "test-strategy"; got != want { + t.Errorf("Name() = %q, want %q", got, want) + } +} diff --git a/libs/tmp/auth/cache.go b/libs/tmp/auth/cache.go new file mode 100644 index 00000000000..96e69c1d133 --- /dev/null +++ b/libs/tmp/auth/cache.go @@ -0,0 +1,245 @@ +package auth + +import ( + "context" + "sync" + "time" +) + +// Maximum time before expiry when async refreshes may start. This value is +// chosen to provide robustness for standard OAuth tokens, supporting 99.95% +// availability, adding breathing room over the existing 99.99% SLA. +const maxAsyncRefreshLeadTime = 20 * time.Minute + +// Minimum wait time before retrying a failed async refresh. This prevents +// hammering a failing server while still recovering quickly from transient +// errors (e.g., brief network issues). With a 20-minute max refresh window, +// this allows up to ~20 retries before the token expires and forces a +// blocking call. +const asyncRefreshRetryBackoff = 1 * time.Minute + +// Maximum time an async refresh is allowed to run before the child context +// is canceled. +const asyncRefreshTimeout = 1 * time.Minute + +// computeAsyncRefreshLeadTime calculates how long before expiry async +// refreshes may start for a token with the given remaining TTL. +// +// Formula: min(TTL × 0.5, maxAsyncRefreshLeadTime) +// +// Edge cases: +// - TTL <= 0 (expired or no expiry): returns 0. +// - Very short TTL (e.g., 2 seconds): returns TTL / 2 without minimum enforcement. +// - Standard OAuth (60 minutes): returns 20 minutes (capped at max). +// - FastPath (10 minutes): returns 5 minutes. +func computeAsyncRefreshLeadTime(ttl time.Duration) time.Duration { + if ttl <= 0 { + return 0 + } + + leadTime := ttl / 2 + if leadTime > maxAsyncRefreshLeadTime { + return maxAsyncRefreshLeadTime + } + return leadTime +} + +type Option func(*cachedTokenProvider) + +// WithCachedToken sets the initial token to be used by a cached token source. +func WithCachedToken(t *Token) Option { + return func(cts *cachedTokenProvider) { + cts.cachedToken = t + } +} + +// WithAsyncRefresh enables or disables the asynchronous token refresh. +func WithAsyncRefresh(b bool) Option { + return func(cts *cachedTokenProvider) { + cts.disableAsync = !b + } +} + +// NewCachedTokenProvider wraps a [TokenProvider] to cache the tokens it returns. +// By default, the cache will refresh tokens asynchronously a few minutes before +// they expire. +// +// The token cache is safe for concurrent use by multiple goroutines and will +// guarantee that only one token refresh is triggered at a time. +// +// The token cache does not take care of retries in case the token source +// returns and error; it is the responsibility of the provided token source to +// handle retries appropriately. +// +// If the TokenProvider is already a cached token source (obtained by calling this +// function), it is returned as is. +func NewCachedTokenProvider(ts TokenProvider, opts ...Option) TokenProvider { + // This is meant as a niche optimization to avoid double caching of the + // token source in situations where the caller needs caching guarantees + // but does not know if the token source is already cached. + if cts, ok := ts.(*cachedTokenProvider); ok { + return cts + } + + cts := &cachedTokenProvider{ + TokenProvider: ts, + timeNow: time.Now, + } + + for _, opt := range opts { + opt(cts) + } + + // If an initial token was provided via WithCachedToken, compute the time + // after which async refreshes may be attempted based on the token TTL + // available at construction time. + if cts.cachedToken != nil { + cts.updateNextAsyncRefresh() + } + + return cts +} + +type cachedTokenProvider struct { + // The token source to obtain tokens from. + TokenProvider TokenProvider + + // If true, only refresh the token with a blocking call when it is expired. + disableAsync bool + + // Earliest time when an async refresh may be attempted for cachedToken. + // Computed from the token TTL at the moment the token is cached. If an + // async refresh fails, this timestamp is pushed forward by + // asyncRefreshRetryBackoff to delay the next retry. + nextAsyncRefresh time.Time + + mu sync.Mutex + cachedToken *Token + + // Indicates that an async refresh is in progress. This is used to prevent + // multiple async refreshes from being triggered at the same time. + isRefreshing bool + + // Monotonic counter incremented each time cachedToken is replaced at + // runtime. Async refresh goroutines snapshot this value before starting + // and skip their result if it has changed, indicating that another path + // (e.g. a blocking refresh) already updated the token. + tokenGeneration uint64 + + timeNow func() time.Time // for testing +} + +// Token returns a token from the cache or fetches a new one if the current +// token is expired. +func (cts *cachedTokenProvider) Token(ctx context.Context) (*Token, error) { + if cts.disableAsync { + return cts.blockingToken(ctx) + } + return cts.asyncToken(ctx) +} + +// updateNextAsyncRefresh recomputes nextAsyncRefresh from the current +// cachedToken's TTL. It must be called immediately after cachedToken is set. +// When called on a shared instance (after construction), the lock must be +// held. +func (cts *cachedTokenProvider) updateNextAsyncRefresh() { + if cts.cachedToken == nil || cts.cachedToken.Expiry.IsZero() { + cts.nextAsyncRefresh = time.Time{} + return + } + ttl := cts.cachedToken.Expiry.Sub(cts.timeNow()) + cts.nextAsyncRefresh = cts.cachedToken.Expiry.Add(-computeAsyncRefreshLeadTime(ttl)) +} + +// tokenExpired reports whether the current token is missing or unusable. The +// function is not thread-safe and should be called with the lock held. +func (cts *cachedTokenProvider) tokenExpired() bool { + if cts.cachedToken == nil { + return true + } + if cts.cachedToken.Expiry.IsZero() { + return false // zero expiry means that token is permanently valid. + } + return cts.timeNow().After(cts.cachedToken.Expiry) +} + +// canRefreshAsync reports whether the current token has entered the async +// refresh window. The function is not thread-safe and should be called with +// the lock held. +func (cts *cachedTokenProvider) canRefreshAsync() bool { + if cts.cachedToken == nil || cts.cachedToken.Expiry.IsZero() { + return false + } + return cts.timeNow().After(cts.nextAsyncRefresh) +} + +func (cts *cachedTokenProvider) asyncToken(ctx context.Context) (*Token, error) { + cts.mu.Lock() + t := cts.cachedToken + tokenExpired := cts.tokenExpired() + canRefreshAsync := !tokenExpired && cts.canRefreshAsync() + cts.mu.Unlock() + + if tokenExpired { + return cts.blockingToken(ctx) + } + if canRefreshAsync { + cts.triggerAsyncRefresh(ctx) + } + return t, nil +} + +func (cts *cachedTokenProvider) blockingToken(ctx context.Context) (*Token, error) { + cts.mu.Lock() + + // The lock is kept for the entire operation to ensure that only one + // blockingToken operation is running at a time. + defer cts.mu.Unlock() + + // It's possible that the token got refreshed (either by a blockingToken or + // an asyncRefresh call) while this particular call was waiting to acquire + // the mutex. This check avoids refreshing the token again in such cases. + if !cts.tokenExpired() { + return cts.cachedToken, nil + } + + t, err := cts.TokenProvider.Token(ctx) + if err != nil { + return nil, err + } + cts.cachedToken = t + cts.tokenGeneration++ + cts.updateNextAsyncRefresh() + return t, nil +} + +func (cts *cachedTokenProvider) triggerAsyncRefresh(ctx context.Context) { + cts.mu.Lock() + defer cts.mu.Unlock() + if cts.isRefreshing || cts.tokenExpired() || !cts.canRefreshAsync() { + return + } + + genAtSubmit := cts.tokenGeneration + cts.isRefreshing = true + go func() { + refreshCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), asyncRefreshTimeout) + defer cancel() + + t, err := cts.TokenProvider.Token(refreshCtx) + + cts.mu.Lock() + defer cts.mu.Unlock() + cts.isRefreshing = false + if cts.tokenGeneration != genAtSubmit { + return // Token was already updated by another path. + } + if err != nil { + cts.nextAsyncRefresh = cts.timeNow().Add(asyncRefreshRetryBackoff) + return + } + cts.cachedToken = t + cts.tokenGeneration++ + cts.updateNextAsyncRefresh() + }() +} diff --git a/libs/tmp/auth/cache_test.go b/libs/tmp/auth/cache_test.go new file mode 100644 index 00000000000..e23d847a9a2 --- /dev/null +++ b/libs/tmp/auth/cache_test.go @@ -0,0 +1,608 @@ +package auth + +import ( + "context" + "fmt" + "reflect" + "sync" + "sync/atomic" + "testing" + "time" +) + +func waitForAsyncRefreshToComplete(t *testing.T, cts *cachedTokenProvider) { + t.Helper() + + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + cts.mu.Lock() + isRefreshing := cts.isRefreshing + cts.mu.Unlock() + + if !isRefreshing { + return + } + + time.Sleep(1 * time.Millisecond) + } + + t.Fatal("timed out waiting for async refresh to complete") +} + +func TestNewCachedTokenProvider_noCaching(t *testing.T) { + want := &cachedTokenProvider{} + got := NewCachedTokenProvider(want, nil) + if got != want { + t.Errorf("NewCachedTokenProvider() = %v, want %v", got, want) + } +} + +func TestNewCachedTokenProvider_default(t *testing.T) { + ts := TokenProviderFn(func(_ context.Context) (*Token, error) { + return nil, nil + }) + + got, ok := NewCachedTokenProvider(ts).(*cachedTokenProvider) + if !ok { + t.Fatalf("NewCachedTokenProvider() = %T, want *cachedTokenProvider", got) + } + + if !got.nextAsyncRefresh.IsZero() { + t.Errorf("NewCachedTokenProvider() nextAsyncRefresh = %v, want zero value", got.nextAsyncRefresh) + } + if got.disableAsync != false { + t.Errorf("NewCachedTokenProvider() disableAsync = %v, want %v", got.disableAsync, false) + } + if got.cachedToken != nil { + t.Errorf("NewCachedTokenProvider() cachedToken = %v, want nil", got.cachedToken) + } +} + +func TestNewCachedTokenProvider_options(t *testing.T) { + ts := TokenProviderFn(func(_ context.Context) (*Token, error) { + return nil, nil + }) + + wantDisableAsync := false + wantCachedToken := &Token{Expiry: time.Unix(42, 0)} + + opts := []Option{ + WithAsyncRefresh(!wantDisableAsync), + WithCachedToken(wantCachedToken), + } + + got, ok := NewCachedTokenProvider(ts, opts...).(*cachedTokenProvider) + if !ok { + t.Fatalf("NewCachedTokenProvider() = %T, want *cachedTokenProvider", got) + } + + if got.disableAsync != wantDisableAsync { + t.Errorf("NewCachedTokenProvider(): disableAsync = %v, want %v", got.disableAsync, wantDisableAsync) + } + if got.cachedToken != wantCachedToken { + t.Errorf("NewCachedTokenProvider(): cachedToken = %v, want %v", got.cachedToken, wantCachedToken) + } +} + +func TestCachedTokenProvider_updateNextAsyncRefresh(t *testing.T) { + now := time.Unix(1337, 0) + + testCases := []struct { + name string + tokenTTL time.Duration + wantAllowedAfter time.Duration + }{ + { + name: "standard OAuth token with 60-minute TTL", + tokenTTL: 60 * time.Minute, + wantAllowedAfter: 40 * time.Minute, + }, + { + name: "short-lived token with 10-minute TTL", + tokenTTL: 10 * time.Minute, + wantAllowedAfter: 5 * time.Minute, + }, + { + name: "very short token with 90-second TTL", + tokenTTL: 90 * time.Second, + wantAllowedAfter: 45 * time.Second, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + cts := &cachedTokenProvider{ + cachedToken: &Token{ + Expiry: now.Add(tc.tokenTTL), + }, + timeNow: func() time.Time { return now }, + } + + cts.updateNextAsyncRefresh() + + want := now.Add(tc.wantAllowedAfter) + if cts.nextAsyncRefresh != want { + t.Errorf("nextAsyncRefresh = %v, want %v", cts.nextAsyncRefresh, want) + } + }) + } +} + +func TestCachedTokenProvider_tokenExpired(t *testing.T) { + now := time.Unix(1337, 0) // mock value for time.Now() + + testCases := []struct { + name string + token *Token + want bool + }{ + { + name: "nil token", + token: nil, + want: true, + }, + { + name: "expired token", + token: &Token{ + Expiry: now.Add(-1 * time.Second), + }, + want: true, + }, + { + name: "token expiring now", + token: &Token{ + Expiry: now, + }, + want: false, + }, + { + name: "future token", + token: &Token{ + Expiry: now.Add(1 * time.Hour), + }, + want: false, + }, + { + name: "token without expiry", + token: &Token{ + Expiry: time.Time{}, + }, + want: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + cts := &cachedTokenProvider{ + cachedToken: tc.token, + timeNow: func() time.Time { return now }, + } + + got := cts.tokenExpired() + + if got != tc.want { + t.Errorf("tokenExpired() = %v, want %v", got, tc.want) + } + }) + } +} + +func TestCachedTokenProvider_canRefreshAsync(t *testing.T) { + now := time.Unix(1337, 0) // mock value for time.Now() + + testCases := []struct { + name string + token *Token + nextAsyncRefresh time.Time + want bool + }{ + { + name: "nil token", + want: false, + }, + { + name: "token without expiry", + token: &Token{ + Expiry: time.Time{}, + }, + want: false, + }, + { + name: "before async refresh window", + token: &Token{ + Expiry: now.Add(1 * time.Hour), + }, + nextAsyncRefresh: now.Add(1 * time.Minute), + want: false, + }, + { + name: "exactly at async refresh boundary", + token: &Token{ + Expiry: now.Add(1 * time.Hour), + }, + nextAsyncRefresh: now, + want: false, + }, + { + name: "after async refresh boundary", + token: &Token{ + Expiry: now.Add(1 * time.Hour), + }, + nextAsyncRefresh: now.Add(-1 * time.Second), + want: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + cts := &cachedTokenProvider{ + cachedToken: tc.token, + nextAsyncRefresh: tc.nextAsyncRefresh, + timeNow: func() time.Time { return now }, + } + + got := cts.canRefreshAsync() + + if got != tc.want { + t.Errorf("canRefreshAsync() = %v, want %v", got, tc.want) + } + }) + } +} + +func TestCachedTokenProvider_AsyncRefreshRetry(t *testing.T) { + failTime := time.Unix(1337, 0) + currentToken := &Token{ + Value: "current-token", + Expiry: failTime.Add(5 * time.Minute), + } + refreshedToken := &Token{ + Value: "refreshed-token", + Expiry: failTime.Add(1 * time.Hour), + } + + testCases := []struct { + name string + now time.Time + wantCalls int32 + wantCache *Token + }{ + { + name: "no async refresh allowed during backoff", + now: failTime.Add(asyncRefreshRetryBackoff - 1*time.Second), + wantCalls: 0, + wantCache: currentToken, + }, + { + name: "async refresh allowed after backoff", + now: failTime.Add(asyncRefreshRetryBackoff + 1*time.Second), + wantCalls: 1, + wantCache: refreshedToken, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + gotCalls := int32(0) + cts := &cachedTokenProvider{ + cachedToken: currentToken, + nextAsyncRefresh: failTime.Add(asyncRefreshRetryBackoff), + timeNow: func() time.Time { return tc.now }, + TokenProvider: TokenProviderFn(func(_ context.Context) (*Token, error) { + atomic.AddInt32(&gotCalls, 1) + return refreshedToken, nil + }), + } + + gotToken, err := cts.Token(context.Background()) + if err != nil { + t.Fatalf("Token() error = %v", err) + } + if gotToken != currentToken { + t.Errorf("Token() = %v, want %v", gotToken, currentToken) + } + + waitForAsyncRefreshToComplete(t, cts) + + if got := atomic.LoadInt32(&gotCalls); got != tc.wantCalls { + t.Fatalf("token source calls = %d, want %d", got, tc.wantCalls) + } + if cts.cachedToken != tc.wantCache { + t.Errorf("cachedToken = %v, want %v", cts.cachedToken, tc.wantCache) + } + }) + } +} + +func TestCachedTokenProvider_Token(t *testing.T) { + now := time.Unix(1337, 0) // mock value for time.Now() + nTokenCalls := 10 // number of goroutines calling Token() + testCases := []struct { + desc string // description of the test case + cachedToken *Token // token cached before calling Token() + disableAsync bool // whether async refreshes are disabled + nextAsyncRefresh time.Time + + returnedToken *Token // token returned by the token source + returnedError error // error returned by the token source + + wantCalls int // expected number of calls to the token source + wantToken *Token // expected token in the cache + }{ + { + desc: "[Blocking] no cached token", + disableAsync: true, + returnedToken: &Token{Expiry: now.Add(1 * time.Hour)}, + wantCalls: 1, + wantToken: &Token{Expiry: now.Add(1 * time.Hour)}, + }, + { + desc: "[Blocking] expired cached token", + disableAsync: true, + cachedToken: &Token{Expiry: now.Add(-1 * time.Second)}, + returnedToken: &Token{Expiry: now.Add(1 * time.Hour)}, + wantCalls: 1, + wantToken: &Token{Expiry: now.Add(1 * time.Hour)}, + }, + { + desc: "[Blocking] fresh cached token", + disableAsync: true, + cachedToken: &Token{Expiry: now.Add(1 * time.Hour)}, + wantCalls: 0, + wantToken: &Token{Expiry: now.Add(1 * time.Hour)}, + }, + { + desc: "[Blocking] stale cached token", + disableAsync: true, + cachedToken: &Token{Expiry: now.Add(1 * time.Minute)}, + wantCalls: 0, + wantToken: &Token{Expiry: now.Add(1 * time.Minute)}, + }, + { + desc: "[Blocking] refresh error", + disableAsync: true, + returnedError: fmt.Errorf("test error"), + wantCalls: 10, + }, + { + desc: "[Blocking] recover from error", + disableAsync: true, + cachedToken: &Token{Expiry: now.Add(-1 * time.Minute)}, + returnedToken: &Token{Expiry: now.Add(-1 * time.Hour)}, + wantCalls: 10, + wantToken: &Token{Expiry: now.Add(-1 * time.Hour)}, + }, + { + desc: "[Async] no cached token", + returnedToken: &Token{Expiry: now.Add(1 * time.Hour)}, + wantCalls: 1, + wantToken: &Token{Expiry: now.Add(1 * time.Hour)}, + }, + { + desc: "[Async] expired cached token", + cachedToken: &Token{Expiry: now.Add(-1 * time.Second)}, + returnedToken: &Token{Expiry: now.Add(1 * time.Hour)}, + wantCalls: 1, + wantToken: &Token{Expiry: now.Add(1 * time.Hour)}, + }, + { + desc: "[Async] fresh cached token", + cachedToken: &Token{Expiry: now.Add(1 * time.Hour)}, + nextAsyncRefresh: now.Add(1 * time.Minute), + wantCalls: 0, + wantToken: &Token{Expiry: now.Add(1 * time.Hour)}, + }, + { + desc: "[Async] stale cached token", + cachedToken: &Token{Expiry: now.Add(1 * time.Minute)}, + nextAsyncRefresh: now.Add(-1 * time.Second), + returnedToken: &Token{Expiry: now.Add(1 * time.Hour)}, + wantCalls: 1, + wantToken: &Token{Expiry: now.Add(1 * time.Hour)}, + }, + { + desc: "[Async] refresh error", + cachedToken: &Token{Expiry: now.Add(1 * time.Minute)}, + nextAsyncRefresh: now.Add(-1 * time.Second), + returnedError: fmt.Errorf("test error"), + wantCalls: 1, + wantToken: &Token{Expiry: now.Add(1 * time.Minute)}, + }, + { + desc: "[Async] stale cached token, expired token returned", + cachedToken: &Token{Expiry: now.Add(1 * time.Minute)}, + nextAsyncRefresh: now.Add(-1 * time.Second), + returnedToken: &Token{Expiry: now.Add(-1 * time.Second)}, + wantCalls: 1, + wantToken: &Token{Expiry: now.Add(-1 * time.Second)}, + }, + { + desc: "[Async] recover from error", + cachedToken: &Token{Expiry: now.Add(-1 * time.Minute)}, + returnedToken: &Token{Expiry: now.Add(-1 * time.Hour)}, + wantCalls: 10, + wantToken: &Token{Expiry: now.Add(-1 * time.Hour)}, + }, + } + + for _, tc := range testCases { + t.Run(tc.desc, func(t *testing.T) { + gotCalls := int32(0) + cts := &cachedTokenProvider{ + disableAsync: tc.disableAsync, + cachedToken: tc.cachedToken, + nextAsyncRefresh: tc.nextAsyncRefresh, + timeNow: func() time.Time { return now }, + TokenProvider: TokenProviderFn(func(_ context.Context) (*Token, error) { + atomic.AddInt32(&gotCalls, 1) + time.Sleep(10 * time.Millisecond) + return tc.returnedToken, tc.returnedError + }), + } + + wg := sync.WaitGroup{} + for range nTokenCalls { + wg.Add(1) + go func() { + defer wg.Done() + _, _ = cts.Token(context.Background()) + }() + } + + wg.Wait() + + waitForAsyncRefreshToComplete(t, cts) + + if got := int(atomic.LoadInt32(&gotCalls)); got != tc.wantCalls { + t.Errorf("want %d calls to cts.TokenProvider.Token(), got %d", tc.wantCalls, got) + } + if !reflect.DeepEqual(tc.wantToken, cts.cachedToken) { + t.Errorf("want cached token %v, got %v", tc.wantToken, cts.cachedToken) + } + }) + } +} + +func TestCachedTokenProvider_BlockingRefreshUpdatesNextAsyncRefresh(t *testing.T) { + now := time.Unix(1337, 0) + refreshedToken := &Token{Expiry: now.Add(1 * time.Hour)} + + cts := &cachedTokenProvider{ + disableAsync: true, + nextAsyncRefresh: now.Add(-2 * time.Minute), + timeNow: func() time.Time { return now }, + cachedToken: &Token{Expiry: now.Add(-1 * time.Second)}, // expired + TokenProvider: TokenProviderFn(func(_ context.Context) (*Token, error) { + return refreshedToken, nil + }), + } + + _, err := cts.Token(context.Background()) + if err != nil { + t.Fatalf("Token() error = %v", err) + } + + want := refreshedToken.Expiry.Add(-maxAsyncRefreshLeadTime) + if cts.nextAsyncRefresh != want { + t.Errorf("nextAsyncRefresh = %v, want %v", cts.nextAsyncRefresh, want) + } +} + +func TestCachedTokenProvider_AsyncRefreshSkipsOlderToken(t *testing.T) { + now := time.Unix(1337, 0) + cachedToken := &Token{ + Value: "fresh-from-blocking", + Expiry: now.Add(1 * time.Hour), + } + olderToken := &Token{ + Value: "stale-from-async", + Expiry: now.Add(30 * time.Minute), + } + + cts := &cachedTokenProvider{ + cachedToken: &Token{Value: "original", Expiry: now.Add(2 * time.Minute)}, + nextAsyncRefresh: now.Add(-1 * time.Second), + timeNow: func() time.Time { return now }, + } + cts.TokenProvider = TokenProviderFn(func(_ context.Context) (*Token, error) { + // Simulate blocking refresh completing first by swapping in the + // fresh token and bumping the generation before this goroutine returns. + cts.mu.Lock() + cts.cachedToken = cachedToken + cts.tokenGeneration++ + cts.mu.Unlock() + return olderToken, nil + }) + + _, err := cts.Token(context.Background()) + if err != nil { + t.Fatalf("Token() error = %v", err) + } + + waitForAsyncRefreshToComplete(t, cts) + + if cts.cachedToken != cachedToken { + t.Errorf("cachedToken = %v, want %v (fresher token from blocking refresh)", cts.cachedToken, cachedToken) + } +} + +func TestCachedTokenProvider_AsyncRefreshAcceptsShorterTTL(t *testing.T) { + now := time.Unix(1337, 0) + originalToken := &Token{ + Value: "original", + Expiry: now.Add(5 * time.Minute), + } + shorterTTLToken := &Token{ + Value: "shorter-ttl", + Expiry: now.Add(2 * time.Minute), + } + + cts := &cachedTokenProvider{ + cachedToken: originalToken, + nextAsyncRefresh: now.Add(-1 * time.Second), + timeNow: func() time.Time { return now }, + TokenProvider: TokenProviderFn(func(_ context.Context) (*Token, error) { + return shorterTTLToken, nil + }), + } + + gotToken, err := cts.Token(context.Background()) + if err != nil { + t.Fatalf("Token() error = %v", err) + } + if gotToken != originalToken { + t.Errorf("Token() = %v, want %v (stale token returned immediately)", gotToken, originalToken) + } + + waitForAsyncRefreshToComplete(t, cts) + + if cts.cachedToken != shorterTTLToken { + t.Errorf("cachedToken = %v, want %v (shorter-TTL token should be accepted)", cts.cachedToken, shorterTTLToken) + } +} + +func TestCachedTokenProvider_LateAsyncFailureDoesNotBackoff(t *testing.T) { + now := time.Unix(1337, 0) + staleToken := &Token{ + Value: "stale", + Expiry: now.Add(5 * time.Minute), + } + blockingResult := &Token{ + Value: "blocking-new", + Expiry: now.Add(1 * time.Hour), + } + + cts := &cachedTokenProvider{ + cachedToken: staleToken, + nextAsyncRefresh: now.Add(-1 * time.Second), + timeNow: func() time.Time { return now }, + } + cts.TokenProvider = TokenProviderFn(func(_ context.Context) (*Token, error) { + // Simulate a blocking refresh completing first: update the token + // and bump the generation before this async goroutine returns an error. + cts.mu.Lock() + cts.cachedToken = blockingResult + cts.tokenGeneration++ + cts.updateNextAsyncRefresh() + cts.mu.Unlock() + return nil, fmt.Errorf("late async failure") + }) + + _, err := cts.Token(context.Background()) + if err != nil { + t.Fatalf("Token() error = %v", err) + } + + waitForAsyncRefreshToComplete(t, cts) + + // The late async failure must not have pushed nextAsyncRefresh forward + // because the token generation changed while it was in flight. + wantNextAsync := blockingResult.Expiry.Add(-maxAsyncRefreshLeadTime) + if cts.nextAsyncRefresh != wantNextAsync { + t.Errorf("nextAsyncRefresh = %v, want %v (late failure should not apply backoff to newer generation)", + cts.nextAsyncRefresh, wantNextAsync) + } + if cts.cachedToken != blockingResult { + t.Errorf("cachedToken = %v, want %v", cts.cachedToken, blockingResult) + } +} diff --git a/libs/tmp/auth/retrying.go b/libs/tmp/auth/retrying.go new file mode 100644 index 00000000000..8944d3aa6fa --- /dev/null +++ b/libs/tmp/auth/retrying.go @@ -0,0 +1,66 @@ +package auth + +import ( + "context" + "net/http" + "time" + + "github.com/databricks/sdk-go/core/apiretry" + "github.com/databricks/sdk-go/core/ops" +) + +// defaultRetriableStatuses is the set of HTTP status codes the default token +// retrier treats as transient. +var defaultRetriableStatuses = []int{ + http.StatusTooManyRequests, // 429 + http.StatusBadGateway, // 502 + http.StatusServiceUnavailable, // 503 + http.StatusGatewayTimeout, // 504 +} + +// NewRetryingTokenProvider wraps a [TokenProvider] with retry logic for +// transient failures. +// +// By default, the return TokenProvider will retry on transient error +// failures (network errors, timeouts, ...) as well as the following HTTP +// status codes: +// +// - 429 Too Many Requests +// - 502 Bad Gateway +// - 503 Service Unavailable +// - 504 Gateway Timeout +// +// The provided options are applied after the defaults, allowing callers to +// override the default timeout and retry behavior. +func NewRetryingTokenProvider(inner TokenProvider, opts ...ops.Option) TokenProvider { + defaults := []ops.Option{ + ops.WithTimeout(1 * time.Minute), + ops.WithRetrier(func() ops.Retrier { + return apiretry.NewRetrier( + ops.BackoffPolicy{ + Initial: 200 * time.Millisecond, + Maximum: 10 * time.Second, + }, + apiretry.RetrierConfig{StatusCodes: defaultRetriableStatuses}, + ) + }), + } + return &retryingTokenProvider{ + inner: inner, + opts: append(defaults, opts...), + } +} + +// retryingTokenProvider wraps a TokenProvider with retry logic for transient +// failures during token acquisition. Each retry calls the underlying Token() +// method, which creates a fresh HTTP request. +type retryingTokenProvider struct { + inner TokenProvider + opts []ops.Option +} + +// Token returns a token from the underlying provider, retrying on transient +// errors. +func (r *retryingTokenProvider) Token(ctx context.Context) (*Token, error) { + return ops.ExecuteWithResult(ctx, r.inner.Token, r.opts...) +} diff --git a/libs/tmp/auth/retrying_test.go b/libs/tmp/auth/retrying_test.go new file mode 100644 index 00000000000..93209270fc3 --- /dev/null +++ b/libs/tmp/auth/retrying_test.go @@ -0,0 +1,150 @@ +package auth + +import ( + "context" + "errors" + "net/http" + "testing" + "time" + + "github.com/databricks/sdk-go/core/apierr" + "github.com/databricks/sdk-go/core/ops" +) + +type mockRetrier struct { + fn func(error) (time.Duration, bool) +} + +func (m *mockRetrier) IsRetriable(err error) (time.Duration, bool) { + return m.fn(err) +} + +func TestRetryingTokenProvider_PlainErrorNotRetried(t *testing.T) { + wantErr := errors.New("boom") + calls := 0 + tp := TokenProviderFn(func(ctx context.Context) (*Token, error) { + calls++ + return nil, wantErr + }) + + got := NewRetryingTokenProvider(tp) + _, err := got.Token(context.Background()) + + if !errors.Is(err, wantErr) { + t.Errorf("Token() error = %v, want %v", err, wantErr) + } + if calls != 1 { + t.Errorf("calls = %d, want 1", calls) + } +} + +func TestRetryingTokenProvider_RetriesUntilSuccess(t *testing.T) { + retriable := errors.New("retriable") + wantToken := &Token{Value: "ok"} + calls := 0 + tp := TokenProviderFn(func(ctx context.Context) (*Token, error) { + calls++ + if calls < 3 { + return nil, retriable + } + return wantToken, nil + }) + + retrier := &mockRetrier{fn: func(err error) (time.Duration, bool) { + return 0, errors.Is(err, retriable) + }} + + got := NewRetryingTokenProvider(tp, + ops.WithRetrier(func() ops.Retrier { return retrier }), + ) + gotToken, err := got.Token(context.Background()) + + if err != nil { + t.Fatalf("Token() error = %v, want nil", err) + } + if gotToken != wantToken { + t.Errorf("Token() = %v, want %v", gotToken, wantToken) + } + if calls != 3 { + t.Errorf("calls = %d, want 3", calls) + } +} + +func TestRetryingTokenProvider_NonRetriableErrorPropagates(t *testing.T) { + retriable := errors.New("retriable") + nonRetriable := errors.New("non-retriable") + calls := 0 + tp := TokenProviderFn(func(ctx context.Context) (*Token, error) { + calls++ + return nil, nonRetriable + }) + + retrier := &mockRetrier{fn: func(err error) (time.Duration, bool) { + return 0, errors.Is(err, retriable) + }} + + got := NewRetryingTokenProvider(tp, + ops.WithRetrier(func() ops.Retrier { return retrier }), + ) + _, err := got.Token(context.Background()) + + if !errors.Is(err, nonRetriable) { + t.Errorf("Token() error = %v, want %v", err, nonRetriable) + } + if calls != 1 { + t.Errorf("calls = %d, want 1", calls) + } +} + +func TestRetryingTokenProvider_ContextCancellation(t *testing.T) { + retriable := errors.New("retriable") + tp := TokenProviderFn(func(ctx context.Context) (*Token, error) { + return nil, retriable + }) + + retrier := &mockRetrier{fn: func(err error) (time.Duration, bool) { + return 1 * time.Hour, errors.Is(err, retriable) + }} + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + got := NewRetryingTokenProvider(tp, + ops.WithRetrier(func() ops.Retrier { return retrier }), + ) + _, err := got.Token(ctx) + + if !errors.Is(err, context.Canceled) { + t.Errorf("Token() error = %v, want context.Canceled", err) + } +} + +// TestRetryingTokenProvider_DefaultRetrier verifies the default retrier +// retries on an APIError with a retriable HTTP status. The bridge from +// oauth2.RetrieveError to APIError happens inside fetchToken; at this layer +// the provider only deals with apierr taxonomy. +func TestRetryingTokenProvider_DefaultRetrier(t *testing.T) { + transient := apierr.FromHTTPError(http.StatusServiceUnavailable, nil, nil) + wantToken := &Token{Value: "ok"} + calls := 0 + tp := TokenProviderFn(func(ctx context.Context) (*Token, error) { + calls++ + if calls < 3 { + return nil, transient + } + return wantToken, nil + }) + + got := NewRetryingTokenProvider(tp) + gotToken, err := got.Token(context.Background()) + + if err != nil { + t.Fatalf("Token() error = %v, want nil", err) + } + if gotToken != wantToken { + t.Errorf("Token() = %v, want %v", gotToken, wantToken) + } + if calls != 3 { + t.Errorf("calls = %d, want 3", calls) + } +} diff --git a/libs/tmp/files/internal/version.go b/libs/tmp/files/internal/version.go new file mode 100644 index 00000000000..4b332087817 --- /dev/null +++ b/libs/tmp/files/internal/version.go @@ -0,0 +1,5 @@ +package internal + +const ModuleName = "sdk-go-files" + +const Version = "0.0.0-dev" diff --git a/libs/tmp/files/v2/client.go b/libs/tmp/files/v2/client.go new file mode 100644 index 00000000000..01ac14d5ca8 --- /dev/null +++ b/libs/tmp/files/v2/client.go @@ -0,0 +1,1303 @@ +// Code generated by Databricks SDK Generator. DO NOT EDIT. + +package files + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "iter" + "log/slog" + "net/http" + "net/url" + "strconv" + "sync" + + "github.com/databricks/cli/libs/tmp/auth" + "github.com/databricks/cli/libs/tmp/files/internal" + "github.com/databricks/cli/libs/tmp/options/call" + "github.com/databricks/cli/libs/tmp/options/client" + "github.com/databricks/cli/libs/tmp/options/internaloptions" + "github.com/databricks/sdk-go/core/clientinfo" +) + +type Client struct { + internalClient + + // extensions is a per-client store for hand-written mixins to attach + // lazily-built state. Keys should be unexported types private to the + // mixin, following the context.Value convention, to avoid collisions. + // The zero value is ready to use; it is safe for concurrent use. + extensions sync.Map +} + +type internalClient struct { + httpClient *http.Client + credentials auth.Credentials + logger *slog.Logger + userAgent string + host string + workspaceID string + accountID string +} + +func NewClient(ctx context.Context, opts ...client.Option) (*Client, error) { + cfg := internaloptions.ClientOptions{} + for _, opt := range opts { + if err := opt(&cfg); err != nil { + return nil, err + } + } + if err := cfg.Resolve(); err != nil { + return nil, err + } + info, err := clientinfo.Default().With( + internal.ModuleName, internal.Version, + "auth", cfg.Credentials.Name(), + ) + if err != nil { + return nil, err + } + + return &Client{ + internalClient: internalClient{ + httpClient: cfg.HTTPClient, + credentials: cfg.Credentials, + logger: cfg.Logger, + userAgent: info.String(), + host: cfg.Host, + workspaceID: cfg.WorkspaceID, + accountID: cfg.AccountID, + }, + }, nil +} + +// Appends a block of data to the stream specified by the input handle. If the +// handle does not exist, this call will throw an exception with +// “RESOURCE_DOES_NOT_EXIST“. +// +// If the block of data exceeds 1 MB, this call will throw an exception with +// “MAX_BLOCK_SIZE_EXCEEDED“. +func (c *internalClient) AddBlock(ctx context.Context, req *AddBlockRequest, opts ...call.Option) (*AddBlockResponse, error) { + body, err := json.Marshal(addBlockRequestToWire(req)) + if err != nil { + return nil, err + } + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + baseURL.Path = "/api/2.0/dbfs/add-block" + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *AddBlockResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "POST", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + Body: bytes.NewBuffer(body), + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp addBlockResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = addBlockResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Closes the stream specified by the input handle. If the handle does not +// exist, this call throws an exception with “RESOURCE_DOES_NOT_EXIST“. +func (c *internalClient) Close(ctx context.Context, req *CloseRequest, opts ...call.Option) (*CloseResponse, error) { + body, err := json.Marshal(closeRequestToWire(req)) + if err != nil { + return nil, err + } + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + baseURL.Path = "/api/2.0/dbfs/close" + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *CloseResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "POST", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + Body: bytes.NewBuffer(body), + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp closeResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = closeResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Opens a stream to write to a file and returns a handle to this stream. There +// is a 10 minute idle timeout on this handle. If a file or directory already +// exists on the given path and __overwrite__ is set to false, this call will +// throw an exception with “RESOURCE_ALREADY_EXISTS“. +// +// A typical workflow for file upload would be: +// +// 1. Issue a “create“ call and get a handle. 2. Issue one or more +// “add-block“ calls with the handle you have. 3. Issue a “close“ call with +// the handle you have. +func (c *internalClient) Create(ctx context.Context, req *CreateRequest, opts ...call.Option) (*CreateResponse, error) { + body, err := json.Marshal(createRequestToWire(req)) + if err != nil { + return nil, err + } + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + baseURL.Path = "/api/2.0/dbfs/create" + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *CreateResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "POST", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + Body: bytes.NewBuffer(body), + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp createResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = createResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Delete the file or directory (optionally recursively delete all files in the +// directory). This call throws an exception with `IO_ERROR` if the path is a +// non-empty directory and `recursive` is set to `false` or on other similar +// errors. +// +// When you delete a large number of files, the delete operation is done in +// increments. The call returns a response after approximately 45 seconds with +// an error message (503 Service Unavailable) asking you to re-invoke the delete +// operation until the directory structure is fully deleted. +// +// For operations that delete more than 10K files, we discourage using the DBFS +// REST API, but advise you to perform such operations in the context of a +// cluster, using the [File system utility +// (dbutils.fs)](/dev-tools/databricks-utils.html#dbutils-fs). `dbutils.fs` +// covers the functional scope of the DBFS REST API, but from notebooks. Running +// such operations using notebooks provides better control and manageability, +// such as selective deletes, and the possibility to automate periodic delete +// jobs. +func (c *internalClient) Delete(ctx context.Context, req *DeleteRequest, opts ...call.Option) (*DeleteResponse, error) { + body, err := json.Marshal(deleteRequestToWire(req)) + if err != nil { + return nil, err + } + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + baseURL.Path = "/api/2.0/dbfs/delete" + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *DeleteResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "POST", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + Body: bytes.NewBuffer(body), + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp deleteResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = deleteResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Gets the file information for a file or directory. If the file or directory +// does not exist, this call throws an exception with `RESOURCE_DOES_NOT_EXIST`. +func (c *internalClient) GetStatus(ctx context.Context, req *GetStatusRequest, opts ...call.Option) (*GetStatusResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + baseURL.Path = "/api/2.0/dbfs/get-status" + queryParams := url.Values{} + if err := addQueryValue(queryParams, "path", req.Path); err != nil { + return nil, err + } + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *GetStatusResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "GET", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp getStatusResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = getStatusResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// List the contents of a directory, or details of the file. If the file or +// directory does not exist, this call throws an exception with +// `RESOURCE_DOES_NOT_EXIST`. +// +// When calling list on a large directory, the list operation will time out +// after approximately 60 seconds. We strongly recommend using list only on +// directories containing less than 10K files and discourage using the DBFS REST +// API for operations that list more than 10K files. Instead, we recommend that +// you perform such operations in the context of a cluster, using the [File +// system utility (dbutils.fs)](/dev-tools/databricks-utils.html#dbutils-fs), +// which provides the same functionality without timing out. +func (c *internalClient) List(ctx context.Context, req *ListStatusRequest, opts ...call.Option) (*ListStatusResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + baseURL.Path = "/api/2.0/dbfs/list" + queryParams := url.Values{} + if err := addQueryValue(queryParams, "path", req.Path); err != nil { + return nil, err + } + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *ListStatusResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "GET", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp listStatusResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = listStatusResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Creates the given directory and necessary parent directories if they do not +// exist. If a file (not a directory) exists at any prefix of the input path, +// this call throws an exception with `RESOURCE_ALREADY_EXISTS`. **Note**: If +// this operation fails, it might have succeeded in creating some of the +// necessary parent directories. +func (c *internalClient) Mkdirs(ctx context.Context, req *MkDirsRequest, opts ...call.Option) (*MkDirsResponse, error) { + body, err := json.Marshal(mkDirsRequestToWire(req)) + if err != nil { + return nil, err + } + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + baseURL.Path = "/api/2.0/dbfs/mkdirs" + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *MkDirsResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "POST", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + Body: bytes.NewBuffer(body), + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp mkDirsResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = mkDirsResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Moves a file from one location to another location within DBFS. If the source +// file does not exist, this call throws an exception with +// `RESOURCE_DOES_NOT_EXIST`. If a file already exists in the destination path, +// this call throws an exception with `RESOURCE_ALREADY_EXISTS`. If the given +// source path is a directory, this call always recursively moves all files. +func (c *internalClient) Move(ctx context.Context, req *MoveRequest, opts ...call.Option) (*MoveResponse, error) { + body, err := json.Marshal(moveRequestToWire(req)) + if err != nil { + return nil, err + } + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + baseURL.Path = "/api/2.0/dbfs/move" + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *MoveResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "POST", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + Body: bytes.NewBuffer(body), + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp moveResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = moveResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Uploads a file through the use of multipart form post. It is mainly used for +// streaming uploads, but can also be used as a convenient single call for data +// upload. +// +// Alternatively you can pass contents as base64 string. +// +// The amount of data that can be passed (when not streaming) using the +// __contents__ parameter is limited to 1 MB. `MAX_BLOCK_SIZE_EXCEEDED` will be +// thrown if this limit is exceeded. +// +// If you want to upload large files, use the streaming upload. For details, see +// :method:dbfs/create, :method:dbfs/addBlock, :method:dbfs/close. +func (c *internalClient) Put(ctx context.Context, req *PutRequest, opts ...call.Option) (*PutResponse, error) { + body, err := json.Marshal(putRequestToWire(req)) + if err != nil { + return nil, err + } + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + baseURL.Path = "/api/2.0/dbfs/put" + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *PutResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "POST", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + Body: bytes.NewBuffer(body), + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp putResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = putResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Returns the contents of a file. If the file does not exist, this call throws +// an exception with `RESOURCE_DOES_NOT_EXIST`. If the path is a directory, the +// read length is negative, or if the offset is negative, this call throws an +// exception with `INVALID_PARAMETER_VALUE`. If the read length exceeds 1 MB, +// this call throws an exception with `MAX_READ_SIZE_EXCEEDED`. +// +// If `offset + length` exceeds the number of bytes in a file, it reads the +// contents until the end of file. +func (c *internalClient) Read(ctx context.Context, req *ReadRequest, opts ...call.Option) (*ReadResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + baseURL.Path = "/api/2.0/dbfs/read" + queryParams := url.Values{} + if err := addQueryValue(queryParams, "path", req.Path); err != nil { + return nil, err + } + if err := addQueryValue(queryParams, "offset", req.Offset); err != nil { + return nil, err + } + if err := addQueryValue(queryParams, "length", req.Length); err != nil { + return nil, err + } + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *ReadResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "GET", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp readResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = readResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Creates an empty directory. If necessary, also creates any parent directories +// of the new, empty directory (like the shell command `mkdir -p`). If called on +// an existing directory, returns a success response; this method is idempotent +// (it will succeed if the directory already exists). +func (c *internalClient) CreateDirectory(ctx context.Context, req *CreateDirectoryRequest, opts ...call.Option) (*CreateDirectoryResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + pb := pathBuilder{} + pb.literal("/api/2.0/fs/directories") + pb.multiSegments(*req.DirectoryPath) + baseURL.Path, baseURL.RawPath = pb.build() + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *CreateDirectoryResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "PUT", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp createDirectoryResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = createDirectoryResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Deletes an empty directory. +// +// To delete a non-empty directory, first delete all of its contents. This can +// be done by listing the directory contents and deleting each file and +// subdirectory recursively. +func (c *internalClient) DeleteDirectory(ctx context.Context, req *DeleteDirectoryRequest, opts ...call.Option) (*DeleteDirectoryResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + pb := pathBuilder{} + pb.literal("/api/2.0/fs/directories") + pb.multiSegments(*req.DirectoryPath) + baseURL.Path, baseURL.RawPath = pb.build() + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *DeleteDirectoryResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "DELETE", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp deleteDirectoryResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = deleteDirectoryResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Deletes a file. If the request is successful, there is no response body. +func (c *internalClient) DeleteFile(ctx context.Context, req *DeleteFileRequest, opts ...call.Option) (*DeleteFileResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + pb := pathBuilder{} + pb.literal("/api/2.0/fs/files") + pb.multiSegments(*req.FilePath) + baseURL.Path, baseURL.RawPath = pb.build() + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *DeleteFileResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "DELETE", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp deleteFileResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = deleteFileResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Downloads a file. The file contents are the response body. This is a standard +// HTTP file download, not a JSON RPC. It supports the Range and +// If-Unmodified-Since HTTP headers. +func (c *internalClient) DownloadFile(ctx context.Context, req *DownloadFileRequest, opts ...call.Option) (*DownloadFileResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + headers.Set("Accept", "application/octet-stream") + if req.Range != nil { + headers.Set("Range", fmt.Sprintf("%v", *req.Range)) + } + if req.IfUnmodifiedSince != nil { + headers.Set("If-Unmodified-Since", fmt.Sprintf("%v", *req.IfUnmodifiedSince)) + } + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + pb := pathBuilder{} + pb.literal("/api/2.0/fs/files") + pb.multiSegments(*req.FilePath) + baseURL.Path, baseURL.RawPath = pb.build() + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *DownloadFileResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "GET", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + }) + if err != nil { + return err + } + + httpResp, err := executeStreamingHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + resp = &DownloadFileResponse{} + resp.Contents = httpResp.Body + if v := httpResp.Header.Get("content-length"); v != "" { + n, err := strconv.ParseInt(v, 10, 64) + if err != nil { + httpResp.Body.Close() // avoid leaking the connection + return err + } + resp.ContentLength = &n + } + if v := httpResp.Header.Get("content-type"); v != "" { + h := v + resp.ContentType = &h + } + if v := httpResp.Header.Get("last-modified"); v != "" { + h := v + resp.LastModified = &h + } + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Get the metadata of a directory. The response HTTP headers contain the +// metadata. There is no response body. +// +// This method is useful to check if a directory exists and the caller has +// access to it. +// +// If you wish to ensure the directory exists, you can instead use `PUT`, which +// will create the directory if it does not exist, and is idempotent (it will +// succeed if the directory already exists). +func (c *internalClient) GetDirectoryMetadata(ctx context.Context, req *GetDirectoryMetadataRequest, opts ...call.Option) (*GetDirectoryMetadataResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + pb := pathBuilder{} + pb.literal("/api/2.0/fs/directories") + pb.multiSegments(*req.DirectoryPath) + baseURL.Path, baseURL.RawPath = pb.build() + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *GetDirectoryMetadataResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "HEAD", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp getDirectoryMetadataResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = getDirectoryMetadataResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Get the metadata of a file. The response HTTP headers contain the metadata. +// There is no response body. +func (c *internalClient) GetFileMetadata(ctx context.Context, req *GetFileMetadataRequest, opts ...call.Option) (*GetFileMetadataResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + if req.Range != nil { + headers.Set("Range", fmt.Sprintf("%v", *req.Range)) + } + if req.IfUnmodifiedSince != nil { + headers.Set("If-Unmodified-Since", fmt.Sprintf("%v", *req.IfUnmodifiedSince)) + } + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + pb := pathBuilder{} + pb.literal("/api/2.0/fs/files") + pb.multiSegments(*req.FilePath) + baseURL.Path, baseURL.RawPath = pb.build() + queryParams := url.Values{} + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *GetFileMetadataResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "HEAD", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + }) + if err != nil { + return err + } + + respBody, respHeader, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + _ = respBody + resp = &GetFileMetadataResponse{} + if v := respHeader.Get("content-length"); v != "" { + n, err := strconv.ParseInt(v, 10, 64) + if err != nil { + return err + } + resp.ContentLength = &n + } + if v := respHeader.Get("content-type"); v != "" { + h := v + resp.ContentType = &h + } + if v := respHeader.Get("last-modified"); v != "" { + h := v + resp.LastModified = &h + } + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// Returns the contents of a directory. If there is no directory at the +// specified path, the API returns an HTTP 404 error. +func (c *internalClient) ListDirectoryContents(ctx context.Context, req *ListDirectoryContentsRequest, opts ...call.Option) (*ListDirectoryResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/json") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + pb := pathBuilder{} + pb.literal("/api/2.0/fs/directories") + pb.multiSegments(*req.DirectoryPath) + baseURL.Path, baseURL.RawPath = pb.build() + queryParams := url.Values{} + if err := addQueryValue(queryParams, "page_size", req.PageSize); err != nil { + return nil, err + } + if err := addQueryValue(queryParams, "page_token", req.PageToken); err != nil { + return nil, err + } + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *ListDirectoryResponse + + call := func(ctx context.Context) error { + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "GET", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp listDirectoryResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = listDirectoryResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} + +// ListDirectoryContentsIter returns an iterator that iterates +// over the results of ListDirectoryContents. +// +// For example: +// +// for item, err := range c.ListDirectoryContentsIter(ctx, &ListDirectoryContentsRequest{}) { +// if err != nil { +// return err +// } +// fmt.Println(item) +// } +// +// Options opts are passed to each ListDirectoryContents call +// made by the iterator under the hood. +// +// Callers who need custom pagination logic should use +// ListDirectoryContents directly. +func (c *internalClient) ListDirectoryContentsIter(ctx context.Context, req *ListDirectoryContentsRequest, opts ...call.Option) iter.Seq2[*DirectoryEntry, error] { + return func(yield func(*DirectoryEntry, error) bool) { + // Deep copy the request via JSON round-trip to avoid modifying the original. + reqBody, err := json.Marshal(req) + if err != nil { + yield(nil, err) + return + } + pageReq := ListDirectoryContentsRequest{} + if err := json.Unmarshal(reqBody, &pageReq); err != nil { + yield(nil, err) + return + } + for { + resp, err := c.ListDirectoryContents(ctx, &pageReq, opts...) + if err != nil { + yield(nil, err) + return + } + for i := range resp.Contents { + if !yield(&resp.Contents[i], nil) { + return + } + } + if resp.NextPageToken == nil || *resp.NextPageToken == "" { + return + } + pageReq.PageToken = resp.NextPageToken + } + } +} + +// Uploads a file of up to 5 GiB. The file contents should be sent as the +// request body as raw bytes (an octet stream); do not encode or otherwise +// modify the bytes before sending. The contents of the resulting file will be +// exactly the bytes sent in the request body. If the request is successful, +// there is no response body. +func (c *internalClient) UploadFile(ctx context.Context, req *UploadFileRequest, opts ...call.Option) (*UploadFileResponse, error) { + + headers := http.Header{} + if c.userAgent != "" { + headers.Set("User-Agent", c.userAgent) + } + headers.Set("Content-Type", "application/octet-stream") + + baseURL, err := url.Parse(c.host) + if err != nil { + return nil, err + } + pb := pathBuilder{} + pb.literal("/api/2.0/fs/files") + pb.multiSegments(*req.FilePath) + baseURL.Path, baseURL.RawPath = pb.build() + queryParams := url.Values{} + if err := addQueryValue(queryParams, "overwrite", req.Overwrite); err != nil { + return nil, err + } + baseURL.RawQuery = queryParams.Encode() + urlStr := baseURL.String() + + var resp *UploadFileResponse + + // A streaming request body is consumed on the first attempt and cannot be + // replayed, so retries default to disabled to avoid resending a partial + // body. Prepended (not appended) so a caller who explicitly passes a + // retry option still overrides this default. Set before the call closure + // below shadows the call package identifier. + opts = append([]call.Option{call.WithDisableRetry()}, opts...) + + call := func(ctx context.Context) error { + var reqBody io.Reader + if req.Contents != nil { + reqBody = req.Contents + } + httpReq, err := newHTTPRequest(ctx, httpRequestOptions{ + Method: "PUT", + URL: urlStr, + Credentials: c.credentials, + Headers: headers, + Body: reqBody, + }) + if err != nil { + return err + } + + respBody, _, err := executeHTTPCall(httpCallOptions{ + req: httpReq, + client: c.httpClient, + logger: c.logger, + }) + if err != nil { + return err + } + var wireResp uploadFileResponseWire + if err := unmarshalResponse(respBody, &wireResp); err != nil { + return err + } + resp = uploadFileResponseFromWire(&wireResp) + return nil + } + + if err := executeCall(ctx, call, opts); err != nil { + return nil, err + } + return resp, nil +} diff --git a/libs/tmp/files/v2/ext_limiter.go b/libs/tmp/files/v2/ext_limiter.go new file mode 100644 index 00000000000..e485e617f10 --- /dev/null +++ b/libs/tmp/files/v2/ext_limiter.go @@ -0,0 +1,55 @@ +package files + +import "context" + +// Limiter bounds the number of concurrent cloud-leg transfers an upload may run. +// The engine acquires one unit before each transfer (every multipart part PUT, +// the single-shot PUT, and every resumable chunk) and releases it when that +// transfer returns. Pass the same Limiter to multiple Upload calls -- for example +// when copying many files at once -- to cap their combined concurrency. +// Implementations must be safe for concurrent use. +type Limiter interface { + // Acquire blocks until a slot is free or ctx is cancelled; on cancellation it + // returns ctx.Err() and the caller must not Release. + Acquire(ctx context.Context) error + // Release returns a slot taken by a successful Acquire. + Release() +} + +// NewLimiter returns a Limiter permitting at most n concurrent transfers. A value +// of n <= 0 yields an unlimited limiter whose Acquire never blocks. +func NewLimiter(n int) Limiter { + if n <= 0 { + return unlimitedLimiter{} + } + return make(chanLimiter, n) +} + +// unlimitedLimiter imposes no bound. It is the default when no limiter is set, so +// the transfer sites can call the limiter unconditionally. +type unlimitedLimiter struct{} + +func (unlimitedLimiter) Acquire(context.Context) error { return nil } +func (unlimitedLimiter) Release() {} + +// chanLimiter is a counting semaphore backed by a buffered channel: a send takes a +// slot, a receive returns one. +type chanLimiter chan struct{} + +func (c chanLimiter) Acquire(ctx context.Context) error { + select { + case c <- struct{}{}: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (c chanLimiter) Release() { <-c } + +// WithLimiter bounds concurrent transfers via l, shared across Upload calls that +// pass the same Limiter. When unset, an upload's own parallelism governs its +// concurrency and there is no cross-upload bound. +func WithLimiter(l Limiter) UploadOption { + return func(c *uploadConfig) { c.limiter = l } +} diff --git a/libs/tmp/files/v2/ext_limiter_test.go b/libs/tmp/files/v2/ext_limiter_test.go new file mode 100644 index 00000000000..e695ffe239e --- /dev/null +++ b/libs/tmp/files/v2/ext_limiter_test.go @@ -0,0 +1,187 @@ +package files + +import ( + "bytes" + "context" + "errors" + "sync" + "testing" + "time" +) + +func TestLimiterCapacityAndCancel(t *testing.T) { + lim := NewLimiter(1) + ctx := t.Context() + if err := lim.Acquire(ctx); err != nil { + t.Fatalf("first acquire: %v", err) + } + + // A second acquire blocks at capacity; cancelling its context unblocks it with + // the context error, and the caller must not Release. + cctx, cancel := context.WithCancel(ctx) + done := make(chan error, 1) + go func() { done <- lim.Acquire(cctx) }() + select { + case err := <-done: + t.Fatalf("acquire should block at capacity, returned %v", err) + case <-time.After(50 * time.Millisecond): + } + cancel() + if err := <-done; !errors.Is(err, context.Canceled) { + t.Fatalf("blocked acquire = %v, want context.Canceled", err) + } + + // Releasing the first slot lets a fresh acquire succeed. + lim.Release() + if err := lim.Acquire(ctx); err != nil { + t.Fatalf("acquire after release: %v", err) + } + lim.Release() +} + +func TestUnlimitedLimiter(t *testing.T) { + lim := NewLimiter(0) // <= 0 means unlimited + for range 5 { + if err := lim.Acquire(t.Context()); err != nil { + t.Fatalf("unlimited acquire: %v", err) + } + } + for range 5 { + lim.Release() + } +} + +// countingLimiter enforces a real bound (via an inner Limiter) and records the +// peak concurrency and the acquire/release counts so tests can assert the engine +// holds exactly one slot per transfer and never exceeds the bound. +type countingLimiter struct { + inner Limiter + + mu sync.Mutex + cur, max int + acquires int + releases int +} + +func newCountingLimiter(n int) *countingLimiter { + return &countingLimiter{inner: NewLimiter(n)} +} + +func (l *countingLimiter) Acquire(ctx context.Context) error { + if err := l.inner.Acquire(ctx); err != nil { + return err + } + l.mu.Lock() + l.cur++ + l.acquires++ + l.max = max(l.max, l.cur) + l.mu.Unlock() + return nil +} + +func (l *countingLimiter) Release() { + l.mu.Lock() + l.cur-- + l.releases++ + l.mu.Unlock() + l.inner.Release() +} + +// slowParts makes every cloud part/chunk PUT take d so concurrent transfers +// overlap and the peak-concurrency assertion is meaningful. +func slowParts(f *fakeServer, d time.Duration) { + f.partHook = func(n, attempt int) (int, string, []byte) { + time.Sleep(d) + return 0, "", nil + } +} + +func TestMultipartLimiterBoundsConcurrency(t *testing.T) { + f := newFakeServer(t, "multipart") + slowParts(f, 15*time.Millisecond) + e := shrunkEngine(t, f) + + const limit = 2 + lim := newCountingLimiter(limit) + in := data(8 * 1024) // 8 parts at the 1024 part size + _, err := e.upload(t.Context(), "/Volumes/c/s/v/m.bin", bytes.NewReader(in), + WithParallelism(8), WithLimiter(lim)) + noErr(t, err) + eqBytes(t, f.assembled(), in) + + if lim.max > limit { + t.Errorf("peak concurrency %d exceeded limit %d", lim.max, limit) + } + if lim.max != limit { + t.Errorf("peak concurrency %d, want it to reach the limit %d (8 parts, 8 workers)", lim.max, limit) + } + if lim.acquires != 8 { + t.Errorf("acquires = %d, want 8 (one per part)", lim.acquires) + } + if lim.acquires != lim.releases { + t.Errorf("unbalanced: %d acquires, %d releases", lim.acquires, lim.releases) + } +} + +func TestSingleShotAcquiresOneSlot(t *testing.T) { + f := newFakeServer(t, "multipart") + e := shrunkEngine(t, f) + + lim := newCountingLimiter(4) + in := data(500) // below the 1024 single-shot threshold + _, err := e.upload(t.Context(), "/Volumes/c/s/v/s.bin", bytes.NewReader(in), WithLimiter(lim)) + noErr(t, err) + eqBytes(t, f.assembled(), in) + + if lim.acquires != 1 || lim.releases != 1 { + t.Errorf("single-shot: acquires=%d releases=%d, want 1/1", lim.acquires, lim.releases) + } + if lim.max != 1 { + t.Errorf("single-shot peak concurrency = %d, want 1", lim.max) + } +} + +func TestResumableAcquiresPerChunk(t *testing.T) { + f := newFakeServer(t, "resumable") + e := shrunkEngine(t, f) + + lim := newCountingLimiter(4) + in := data(4 * 1024) // several resumable chunks + _, err := e.upload(t.Context(), "/Volumes/c/s/v/r.bin", bytes.NewReader(in), WithLimiter(lim)) + noErr(t, err) + eqBytes(t, f.assembled(), in) + + if lim.acquires < 2 { + t.Errorf("resumable acquires = %d, want >= 2 (one per chunk)", lim.acquires) + } + if lim.acquires != lim.releases { + t.Errorf("unbalanced: %d acquires, %d releases", lim.acquires, lim.releases) + } + if lim.max != 1 { + t.Errorf("resumable is sequential; peak concurrency = %d, want 1", lim.max) + } +} + +func TestLimiterCancelNoLeak(t *testing.T) { + f := newFakeServer(t, "multipart") + slowParts(f, 50*time.Millisecond) + e := shrunkEngine(t, f) + + lim := newCountingLimiter(2) + ctx, cancel := context.WithCancel(t.Context()) + errc := make(chan error, 1) + go func() { + _, err := e.upload(ctx, "/Volumes/c/s/v/c.bin", bytes.NewReader(data(16*1024)), + WithParallelism(8), WithLimiter(lim)) + errc <- err + }() + time.Sleep(20 * time.Millisecond) + cancel() + <-errc // upload returns (with a context error) + + lim.mu.Lock() + defer lim.mu.Unlock() + if lim.acquires != lim.releases { + t.Errorf("cancellation leaked slots: %d acquires, %d releases", lim.acquires, lim.releases) + } +} diff --git a/libs/tmp/files/v2/ext_multipart.go b/libs/tmp/files/v2/ext_multipart.go new file mode 100644 index 00000000000..a882d84e715 --- /dev/null +++ b/libs/tmp/files/v2/ext_multipart.go @@ -0,0 +1,450 @@ +package files + +// Multipart upload for AWS and Azure: the part-based protocol and its +// single-threaded and parallel (producer/worker) drivers. The control-plane +// calls it relies on live in upload_control.go; the part transfers go through +// uc.cloud. + +import ( + "context" + "errors" + "fmt" + "io" + "maps" + "net/http" + "sync" + "time" + + "github.com/databricks/sdk-go/core/apierr" + "github.com/databricks/cli/libs/tmp/files/v2/internal/cloudstorage" +) + +// uploadPart is the unit of work handed to an upload worker: a part number, the +// source of its bytes (an in-memory buffer for a streamed part, a section of the +// open file for a file part), and an optionally pre-minted presigned URL. When +// url is empty the worker mints one on demand; the file path pre-mints URLs in +// batches so workers go straight to the cloud PUT. +type uploadPart struct { + partNumber int + body cloudstorage.Body + url string // pre-minted presigned URL; "" => mint on demand + headers map[string]string // PUT headers accompanying url +} + +// isCloudStatus reports whether err is a cloud-storage HTTP status error (an +// *apierr.APIError returned by Send) rather than a transport failure. A status +// rejection before any part has landed triggers the single-shot fallback; a +// transport failure does not. +func isCloudStatus(err error) bool { + _, ok := errors.AsType[*apierr.APIError](err) + return ok +} + +// performMultipartUpload uploads a stream in parts on a single goroutine using +// presigned URLs (AWS/Azure). +func (e *engine) performMultipartUpload(ctx context.Context, uc *uploadContext, token string, r io.Reader, preRead []byte) error { + currentPart := 1 + etags := map[int]string{} + chunkOffset := int64(0) + // AWS and Azure require each part size up front, so we buffer a part before + // uploading it; this also lets us replay a part on retry. The buffer starts + // with the bytes already read while deciding single-shot vs multipart. + buffer := preRead + retryCount := 0 + eof := false + + for !eof { + var err error + if buffer, err = fillBuffer(buffer, uc.partSize, r); err != nil { + return err + } + if len(buffer) == 0 { + break + } + + urls, err := e.createPartURLs(ctx, uc.targetPath, token, currentPart, uc.batchSize) + if err != nil { + if chunkOffset == 0 { + return &fallbackToFilesAPI{buffer: buffer, reason: fmt.Sprintf("failed to obtain upload URLs: %v", err)} + } + return err + } + + for _, purl := range urls { + if buffer, err = fillBuffer(buffer, uc.partSize, r); err != nil { + return err + } + if len(buffer) == 0 { + eof = true + break + } + chunkLen := min(int64(len(buffer)), uc.partSize) + headers := octetStreamHeaders(purl.Headers) + + partStart := time.Now() + if err := uc.limiter.Acquire(ctx); err != nil { + return err + } + resp, rerr := uc.cloud.Send(ctx, http.MethodPut, purl.URL, headers, cloudstorage.BytesBody(buffer[:chunkLen])) + uc.limiter.Release() + switch { + case rerr == nil: + chunkOffset += chunkLen + etags[currentPart] = resp.Header.Get("ETag") + buffer = buffer[chunkLen:] + retryCount = 0 + uc.progress.add(chunkLen) + e.c.logger.DebugContext(ctx, "chunk uploaded", + "part", currentPart, "bytes", chunkLen, "offset", chunkOffset-chunkLen, + "duration_ms", time.Since(partStart).Milliseconds()) + case cloudstorage.IsURLExpired(rerr): + if retryCount >= e.tun.maxRetries { + return errUploadURLExpired + } + retryCount++ + // Preserve the buffer: the same bytes are uploaded under the next + // part number using a fresh URL. + case chunkOffset == 0 && isCloudStatus(rerr): + // The cloud rejected the first chunk (e.g. a firewall 403); fall back + // to a single-shot upload through the Files API. + return &fallbackToFilesAPI{buffer: buffer, reason: fmt.Sprintf("first chunk upload failed: %v", rerr)} + default: + return rerr + } + currentPart++ + } + } + + return e.completeMultipart(ctx, uc.targetPath, token, etags) +} + +// doUploadOnePart uploads a single part. It uses part.url when the caller +// pre-minted one (the file path mints in batches); otherwise, and whenever a URL +// expires or a slow attempt is re-issued, it mints a fresh one (count=1). While no +// part has completed yet, it converts a cloud rejection or URL-mint failure into a +// fallback signal so the caller can retry the whole upload through a single-shot +// Files API PUT (the firewall/auth case); once any part has landed, the storage +// endpoint is known good and such a failure is returned as a real error. +// +// isFirstPart marks the stream path's synchronous, uncontended first part, which +// gets its own tight soft deadline; it does not affect the fallback gating above. +func (e *engine) doUploadOnePart(ctx context.Context, uc *uploadContext, part uploadPart, token string, isFirstPart bool) (string, error) { + start := time.Now() + urlRetries, slowRetries := 0, 0 + url, headers := part.url, part.headers + for { + if url == "" { + urls, err := e.createPartURLs(ctx, uc.targetPath, token, part.partNumber, 1) + if err != nil { + if !uc.completed.Load() { + return "", &fallbackToFilesAPI{reason: fmt.Sprintf("failed to obtain upload URL for part %d: %v", part.partNumber, err)} + } + return "", err + } + url, headers = urls[0].URL, octetStreamHeaders(urls[0].Headers) + } + + attemptStart := time.Now() + resp, slow, rerr := e.sendPart(ctx, uc, url, headers, part.body, func() time.Duration { + return uc.attemptDeadline(isFirstPart, slowRetries) + }) + switch { + case rerr == nil: + uc.completed.Store(true) + uc.slowGuard.record(time.Since(attemptStart)) + uc.progress.add(part.body.Size()) + e.c.logger.DebugContext(ctx, "part uploaded", + "part", part.partNumber, "bytes", part.body.Size(), "first", isFirstPart, + "duration_ms", time.Since(start).Milliseconds()) + return resp.Header.Get("ETag"), nil + case slow: + // A wedged connection: re-mint and re-issue on a fresh URL and connection. + // attemptDeadline disarms after slowAttemptMax tries, so a part that keeps + // drawing slow connections rides out the normal timeouts rather than failing. + slowRetries++ + url = "" + e.c.logger.DebugContext(ctx, "part slow, re-issuing on a fresh connection", + "part", part.partNumber, "attempt", slowRetries+1) + case cloudstorage.IsURLExpired(rerr): + if urlRetries >= e.tun.maxRetries { + return "", errUploadURLExpired + } + urlRetries++ + url = "" // re-mint the expired URL + case isCloudStatus(rerr) && !uc.completed.Load(): + // The cloud rejected this part before any part has landed (e.g. a firewall + // 403 on the storage host); fall back to a single-shot upload through the + // Files API. + return "", &fallbackToFilesAPI{reason: fmt.Sprintf("part %d upload failed before any part completed: %v", part.partNumber, rerr)} + default: + return "", rerr + } + } +} + +// sendPart performs one part PUT, cancelling it (slow=true) if it outlives the +// soft deadline so the caller can re-issue on a fresh connection (a wedged +// connection trickles bytes, escaping the idle and response timeouts). deadline +// is re-evaluated while the attempt is in flight, so a part that started under +// the cold-start deadline is caught at the tighter p95 deadline as soon as the +// guard warms up rather than waiting out the floor; deadline returning <= 0 +// leaves the attempt unbounded. A ctx cancellation from the caller is reported +// as not-slow so it propagates instead of looping. +func (e *engine) sendPart(ctx context.Context, uc *uploadContext, url string, headers map[string]string, body cloudstorage.Body, deadline func() time.Duration) (resp *cloudstorage.Response, slow bool, err error) { + // Hold one transfer slot for the duration of this attempt. A re-issued attempt + // (slow/expiry) acquires a fresh slot, so a part never holds more than one. + if err := uc.limiter.Acquire(ctx); err != nil { + return nil, false, err + } + defer uc.limiter.Release() + + attemptCtx, cancel := context.WithCancelCause(ctx) + defer cancel(nil) + + // The deferred cancel above fires on every return path, so attemptCtx.Done() + // is the goroutine's stop signal once Send returns. + start := time.Now() + go func() { + ticker := time.NewTicker(uc.tun.slowCheckInterval) + defer ticker.Stop() + for { + select { + case <-attemptCtx.Done(): + return + case <-ticker.C: + if d := deadline(); d > 0 && time.Since(start) >= d { + cancel(errSlowAttempt) + return + } + } + } + }() + + resp, err = uc.cloud.Send(attemptCtx, http.MethodPut, url, headers, body) + slow = ctx.Err() == nil && errors.Is(context.Cause(attemptCtx), errSlowAttempt) + return resp, slow, err +} + +// startPartWorkers spawns parallelism workers that drain tasks, upload each part +// from its source, and record the ETag into seed (which the caller may +// pre-populate with parts it already uploaded, e.g. the stream path's synchronous +// first part). The first failure records the error and cancels +// the rest via cancel, which a stream producer also observes to stop feeding. +// The returned wait blocks for all workers and returns the collected ETags or +// the first error. The caller owns producing and closing tasks. +func (e *engine) startPartWorkers(pctx context.Context, cancel context.CancelFunc, uc *uploadContext, token string, parallelism int, seed map[int]string, tasks <-chan uploadPart) func() (map[int]string, error) { + results := &uploadResults{etags: seed} + var wg sync.WaitGroup + for range parallelism { + wg.Go(func() { + for part := range tasks { + if pctx.Err() != nil { + return + } + etag, e := e.doUploadOnePart(pctx, uc, part, token, false) + if e != nil { + results.setErr(e) + cancel() + return + } + results.put(part.partNumber, etag) + } + }) + } + return func() (map[int]string, error) { + wg.Wait() + if err := results.err(); err != nil { + return nil, err + } + return results.snapshot(), nil + } +} + +// parallelMultipartFromStream uploads a non-seekable stream in parts using a +// bounded producer/consumer. The first part is uploaded synchronously so an +// auth/firewall rejection surfaces before workers spin up and can fall back. +func (e *engine) parallelMultipartFromStream(ctx context.Context, uc *uploadContext, token string, r io.Reader) error { + first, err := readUpTo(r, uc.partSize) + if err != nil { + return err + } + if len(first) == 0 { + return &fallbackToFilesAPI{buffer: nil, reason: "empty input stream"} + } + etag1, err := e.doUploadOnePart(ctx, uc, uploadPart{partNumber: 1, body: cloudstorage.BytesBody(first)}, token, true) + if err != nil { + if fb, ok := errors.AsType[*fallbackToFilesAPI](err); ok { + fb.buffer = first + return fb + } + return err + } + if int64(len(first)) < uc.partSize { + return e.completeMultipart(ctx, uc.targetPath, token, map[int]string{1: etag1}) + } + + pctx, cancel := context.WithCancel(ctx) + defer cancel() + + // Hand parts straight to workers over an unbuffered channel so the bytes held + // in memory stay bounded to the in-flight set: at most parallelism workers each + // holding one partSize buffer, plus the one the producer is handing off. The + // stream size is unknown, so the worker count cannot be capped by the part + // count (as the file path does); workers beyond the number of parts observe the + // closed channel and exit. + tasks := make(chan uploadPart) + wait := e.startPartWorkers(pctx, cancel, uc, token, uc.parallelism, map[int]string{1: etag1}, tasks) + + // The calling goroutine is the single producer. It reads sequential parts + // and stops on EOF, an error, or cancellation. + var producerErr error + func() { + defer close(tasks) + partNumber := 2 + for { + if pctx.Err() != nil { + return + } + buf, rerr := readUpTo(r, uc.partSize) + if rerr != nil { + producerErr = rerr + cancel() + return + } + if len(buf) == 0 { + return + } + select { + case tasks <- uploadPart{partNumber: partNumber, body: cloudstorage.BytesBody(buf)}: + case <-pctx.Done(): + return + } + partNumber++ + if int64(len(buf)) < uc.partSize { + return // short read => end of stream + } + } + }() + + etags, err := wait() + if producerErr != nil { + return producerErr + } + if err != nil { + return err + } + return e.completeMultipart(ctx, uc.targetPath, token, etags) +} + +// parallelMultipartFromReaderAt uploads a known-size, randomly-readable source in +// parts. Each part streams its own section of the shared io.ReaderAt on demand, +// so the resident bytes stay bounded to the transport's per-connection write +// buffers rather than parallelism*partSize. The caller owns the reader's +// lifetime; it must stay readable until this returns (wait() joins all workers +// first). +// +// Unlike the stream path, the first part is not uploaded synchronously: a +// randomly-readable source lets the single-shot fallback re-read from the start, +// so part 1 is just another worker task rather than a canary that blocks the +// whole upload at 0%. The fallback is instead gated on no part having completed +// (see doUploadOnePart). +func (e *engine) parallelMultipartFromReaderAt(ctx context.Context, uc *uploadContext, token string, ra io.ReaderAt) error { + fileSize := uc.contentLength + partSize := uc.partSize + numParts := max(int((fileSize+partSize-1)/partSize), 1) + + pctx, cancel := context.WithCancel(ctx) + defer cancel() + + // A bounded channel lets the minting producer (this goroutine) run about a + // batch ahead of the workers without holding many presigned URLs outstanding + // at once, and gives backpressure on the slowest leg. + tasks := make(chan uploadPart, uc.batchSize) + wait := e.startPartWorkers(pctx, cancel, uc, token, min(uc.parallelism, numParts), map[int]string{}, tasks) + + // Mint presigned URLs in batches and hand each part to the workers pre-minted, + // so a worker goes straight to the cloud PUT instead of minting its own URL: + // one control-plane call per batch rather than per part, and no parallelism-wide + // mint burst at startup. doUploadOnePart still re-mints on its own when a URL + // expires or a slow attempt is re-issued. + var producerErr error + func() { + defer close(tasks) + for next := 1; next <= numParts; { + if pctx.Err() != nil { + return + } + count := min(uc.batchSize, numParts-next+1) + urls, err := e.createPartURLs(pctx, uc.targetPath, token, next, count) + if err != nil { + if pctx.Err() != nil { + return // cancelled by a worker failure; that error wins + } + if uc.completed.Load() { + producerErr = err + } else { + producerErr = &fallbackToFilesAPI{reason: fmt.Sprintf("failed to obtain upload URLs at part %d: %v", next, err)} + } + cancel() + return + } + for _, u := range urls { + offset := int64(u.PartNumber-1) * partSize + body := cloudstorage.SectionBody(ra, offset, min(partSize, fileSize-offset)) + select { + case tasks <- uploadPart{partNumber: u.PartNumber, body: body, url: u.URL, headers: octetStreamHeaders(u.Headers)}: + case <-pctx.Done(): + return + } + } + next += len(urls) + } + }() + + etags, err := wait() + if producerErr != nil { + err = producerErr // a producer failure wins over a worker's observed cancellation + } + if err != nil { + if fb, ok := errors.AsType[*fallbackToFilesAPI](err); ok { + return fb // the fallback re-reads the source, so no buffer is needed + } + return err + } + return e.completeMultipart(ctx, uc.targetPath, token, etags) +} + +// uploadResults collects part ETags from workers and records the first error. +type uploadResults struct { + mu sync.Mutex + etags map[int]string + firstErr error +} + +func (r *uploadResults) put(partNumber int, etag string) { + r.mu.Lock() + defer r.mu.Unlock() + r.etags[partNumber] = etag +} + +// setErr records the first error reported by any worker; later errors are +// dropped so the original cause survives. +func (r *uploadResults) setErr(err error) { + r.mu.Lock() + defer r.mu.Unlock() + if r.firstErr == nil { + r.firstErr = err + } +} + +func (r *uploadResults) err() error { + r.mu.Lock() + defer r.mu.Unlock() + return r.firstErr +} + +func (r *uploadResults) snapshot() map[int]string { + r.mu.Lock() + defer r.mu.Unlock() + return maps.Clone(r.etags) +} diff --git a/libs/tmp/files/v2/ext_resumable.go b/libs/tmp/files/v2/ext_resumable.go new file mode 100644 index 00000000000..76a1a1d22d2 --- /dev/null +++ b/libs/tmp/files/v2/ext_resumable.go @@ -0,0 +1,187 @@ +package files + +// Resumable upload for GCP: the single-stream, offset-resuming protocol used +// when the server hands back a resumable session instead of multipart. Chunk +// transfers go through uc.cloud.Attempt (raw responses, since 308 is the normal +// "continue" signal); the loop runs its own resume-aware retry. + +import ( + "context" + "fmt" + "io" + "net/http" + "regexp" + "slices" + "strconv" + + "github.com/databricks/sdk-go/core/apierr" + "github.com/databricks/cli/libs/tmp/files/v2/internal/cloudstorage" +) + +var resumableRangeRe = regexp.MustCompile(`bytes=0-(\d+)`) + +// performResumableUpload uploads a stream using the GCP resumable protocol, +// aborting the session if any chunk fails. +func (e *engine) performResumableUpload(ctx context.Context, uc *uploadContext, token string, r io.Reader, preRead []byte) error { + purl, err := e.createResumableURL(ctx, uc.targetPath, token) + if err != nil { + return &fallbackToFilesAPI{buffer: preRead, reason: fmt.Sprintf("failed to obtain resumable upload URL: %v", err)} + } + if err := e.resumableUploadLoop(ctx, uc, purl, r, preRead); err != nil { + if aerr := e.abortResumableUpload(ctx, uc, purl); aerr != nil { + e.c.logger.WarnContext(ctx, "failed to abort resumable upload", "error", aerr) + } + return err + } + return nil +} + +func (e *engine) resumableUploadLoop(ctx context.Context, uc *uploadContext, purl presignedURL, r io.Reader, preRead []byte) error { + // Buffer one part plus a read-ahead byte: a resumable chunk cannot be empty, + // so reading one byte past the part tells us whether this is the last chunk + // (which must be sent with the real total size, not "*"). + minBuffer := uc.partSize + uc.tun.readAheadBytes + buffer := slices.Clone(preRead) + chunkOffset := int64(0) + retryCount := 0 + noProgress := 0 + + for { + lastChunk := false + if need := minBuffer - int64(len(buffer)); need > 0 { + tmp := make([]byte, need) + n, rerr := io.ReadFull(r, tmp) + if rerr != nil && rerr != io.EOF && rerr != io.ErrUnexpectedEOF { + return rerr + } + buffer = append(buffer, tmp[:n]...) + if int64(n) < need { + lastChunk = true + } + } + + var chunkLen int64 + var totalSize string + if lastChunk { + chunkLen = int64(len(buffer)) + totalSize = strconv.FormatInt(chunkOffset+chunkLen, 10) + } else { + chunkLen = uc.partSize + totalSize = "*" + } + // An empty stream cannot be sent as a resumable upload: a chunk must carry + // at least one byte, and the range would be the malformed "bytes 0--1/0". + // Fall back to a single-shot PUT, which creates the empty object. Only + // reachable at offset 0, since the read-ahead byte folds a non-empty + // stream's tail into the preceding chunk. + if chunkLen == 0 { + return &fallbackToFilesAPI{reason: "empty input stream"} + } + chunkLastByte := chunkOffset + chunkLen - 1 + + headers := octetStreamHeaders(purl.Headers) + headers["Content-Range"] = fmt.Sprintf("bytes %d-%d/%s", chunkOffset, chunkLastByte, totalSize) + + if err := uc.limiter.Acquire(ctx); err != nil { + return err + } + resp, rerr := uc.cloud.Attempt(ctx, http.MethodPut, purl.URL, headers, cloudstorage.BytesBody(buffer[:chunkLen])) + uc.limiter.Release() + switch { + case rerr != nil: + // On a transient transport error, query the server for the confirmed + // offset and continue from there; otherwise surface the error. + if retryCount >= uc.tun.maxRetries || !cloudstorage.IsRetriable(rerr) { + return rerr + } + retryCount++ + status, qerr := e.resumableStatusQuery(ctx, uc, purl) + if qerr != nil || status == nil { + return rerr + } + resp = status + case cloudstorage.IsRetryableStatus(resp.StatusCode): + if retryCount < uc.tun.maxRetries { + retryCount++ + if status, qerr := e.resumableStatusQuery(ctx, uc, purl); qerr == nil && status != nil { + resp = status + } + } + default: + retryCount = 0 + } + + switch { + case resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusCreated: + if totalSize == "*" { + return fmt.Errorf("resumable upload reported complete with status %d before end of stream", resp.StatusCode) + } + uc.progress.add(chunkLen) + return nil + case resp.StatusCode == http.StatusPermanentRedirect: // 308: chunk accepted, more to come + confirmed, ok, perr := extractRangeOffset(resp.Header.Get("Range")) + if perr != nil { + return perr + } + var nextOffset int64 + if ok { + if confirmed < chunkOffset-1 || confirmed > chunkLastByte { + return fmt.Errorf("resumable upload confirmed offset %d outside chunk [%d, %d]", confirmed, chunkOffset, chunkLastByte) + } + nextOffset = confirmed + 1 + } else { + if chunkOffset > 0 { + return fmt.Errorf("resumable upload returned no confirmed offset at chunk offset %d", chunkOffset) + } + nextOffset = chunkOffset + } + // A 308 confirming no new bytes (nextOffset == chunkOffset) leaves the + // chunk unsent, and the first switch's default arm resets retryCount, so a + // server stuck at a fixed offset would loop forever (fs cp sets no context + // deadline). Bound it with a counter that survives that reset. + if nextOffset == chunkOffset { + noProgress++ + if noProgress > uc.tun.maxRetries { + return fmt.Errorf("resumable upload made no progress at offset %d after %d attempts", chunkOffset, noProgress) + } + } else { + // Forward progress recovered from a transient error via the status + // query: reset both budgets so they bound consecutive stalls, not the + // total blip count over a long, steadily-advancing upload. + noProgress = 0 + retryCount = 0 + } + buffer = buffer[nextOffset-chunkOffset:] + uc.progress.add(nextOffset - chunkOffset) + chunkOffset = nextOffset + case resp.StatusCode == http.StatusPreconditionFailed && (uc.overwrite == nil || !*uc.overwrite): + return ErrAlreadyExists + default: + if apiErr := apierr.FromHTTPError(resp.StatusCode, resp.Header, resp.Body); apiErr != nil { + return apiErr + } + return fmt.Errorf("resumable chunk upload failed with status %d", resp.StatusCode) + } + } +} + +func (e *engine) resumableStatusQuery(ctx context.Context, uc *uploadContext, purl presignedURL) (*cloudstorage.Response, error) { + return uc.cloud.Attempt(ctx, http.MethodPut, purl.URL, map[string]string{"Content-Range": "bytes */*"}, nil) +} + +// extractRangeOffset parses a resumable upload "Range: bytes=0-N" response +// header and returns N. An empty header means the server has confirmed nothing. +func extractRangeOffset(rangeHeader string) (offset int64, ok bool, err error) { + if rangeHeader == "" { + return 0, false, nil + } + m := resumableRangeRe.FindStringSubmatch(rangeHeader) + if m == nil { + return 0, false, fmt.Errorf("cannot parse Range header %q", rangeHeader) + } + n, perr := strconv.ParseInt(m[1], 10, 64) + if perr != nil { + return 0, false, fmt.Errorf("cannot parse Range header %q: %w", rangeHeader, perr) + } + return n, true, nil +} diff --git a/libs/tmp/files/v2/ext_resumable_test.go b/libs/tmp/files/v2/ext_resumable_test.go new file mode 100644 index 00000000000..bdf579ceb0d --- /dev/null +++ b/libs/tmp/files/v2/ext_resumable_test.go @@ -0,0 +1,21 @@ +package files + +import "testing" + +func TestExtractRangeOffset(t *testing.T) { + n, ok, err := extractRangeOffset("bytes=0-99") + noErr(t, err) + if !ok || n != 99 { + t.Errorf("extractRangeOffset(bytes=0-99) = (%d, %v), want (99, true)", n, ok) + } + + _, ok, err = extractRangeOffset("") + noErr(t, err) + if ok { + t.Error("empty header should report nothing confirmed") + } + + if _, _, err := extractRangeOffset("garbage"); err == nil { + t.Error("malformed header should error") + } +} diff --git a/libs/tmp/files/v2/ext_straggler.go b/libs/tmp/files/v2/ext_straggler.go new file mode 100644 index 00000000000..5b60d19c4b7 --- /dev/null +++ b/libs/tmp/files/v2/ext_straggler.go @@ -0,0 +1,81 @@ +package files + +// Straggler control for multipart part uploads. A wedged cloud connection +// trickles bytes for minutes while its peers stay fast; it trips neither the +// idle timeout (bytes keep moving) nor the response timeout (the response phase +// is fine), so a single slow part gates the whole upload (completeMultipart +// needs every ETag). doUploadOnePart cancels an attempt that runs far longer +// than the recent part-duration tail and re-issues it on a fresh connection. +// +// The trigger keys off the recent p95, not the median: at high concurrency the +// opening burst of connections makes many parts legitimately slow (tens of +// seconds) even though the median stays low, so a median-relative trigger would +// both miss the real outliers and falsely cancel healthy burst parts. A p95 +// multiple tracks that legitimate tail and fires only well above it, while still +// adapting to a uniformly slow network (where p95 rises with everything else). + +import ( + "errors" + "slices" + "sync" + "time" +) + +// errSlowAttempt is the cancel cause for a part attempt that outlived its soft +// deadline. It never escapes doUploadOnePart; it only distinguishes the guard's +// own cancellation from the caller cancelling ctx. +var errSlowAttempt = errors.New("part attempt exceeded soft deadline") + +// slowAttemptGuard tracks the duration of recently completed part attempts so an +// attempt that far exceeds the recent tail can be cancelled and re-issued. Safe +// for concurrent use by the upload workers. Its policy comes from tun (see the +// slow* fields on tunables). +type slowAttemptGuard struct { + tun tunables + mu sync.Mutex + samples []time.Duration +} + +// record adds a completed attempt's duration to the rolling window. +func (g *slowAttemptGuard) record(d time.Duration) { + g.mu.Lock() + defer g.mu.Unlock() + g.samples = append(g.samples, d) + if len(g.samples) > g.tun.slowWindow { + g.samples = g.samples[len(g.samples)-g.tun.slowWindow:] + } +} + +// deadline returns the soft deadline for the next attempt: the cold-start +// deadline until enough attempts have completed for the p95 to be meaningful, +// then slowFactor x the recent p95 (floored at slowMinDeadline). The p95 ignores +// the rare wedged attempts in the window (they sit above it), so the deadline +// does not drift up toward the stragglers it is meant to catch. +func (g *slowAttemptGuard) deadline() time.Duration { + g.mu.Lock() + defer g.mu.Unlock() + if len(g.samples) < g.tun.slowWarmup { + return g.tun.slowColdDeadline + } + s := slices.Clone(g.samples) + slices.Sort(s) + p95 := s[min(len(s)*95/100, len(s)-1)] + return max(time.Duration(g.tun.slowFactor)*p95, g.tun.slowMinDeadline) +} + +// attemptDeadline returns the current soft deadline for a part attempt. It is +// called repeatedly while an attempt is in flight (see sendPart), so the value +// tracks the warming-up guard. The synchronous first part gets its own tight, +// contention-free deadline; once a part has been re-issued slowMaxReissue times +// the guard disarms (0) so the part rides out the normal timeouts instead of +// being cancelled again. +func (uc *uploadContext) attemptDeadline(isFirstPart bool, slowRetries int) time.Duration { + switch { + case slowRetries >= uc.tun.slowMaxReissue: + return 0 + case isFirstPart: + return uc.tun.slowFirstPartDeadline + default: + return uc.slowGuard.deadline() + } +} diff --git a/libs/tmp/files/v2/ext_straggler_test.go b/libs/tmp/files/v2/ext_straggler_test.go new file mode 100644 index 00000000000..ff74e289c0f --- /dev/null +++ b/libs/tmp/files/v2/ext_straggler_test.go @@ -0,0 +1,101 @@ +package files + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/databricks/cli/libs/tmp/files/v2/internal/cloudstorage" +) + +func TestSlowAttemptGuardDeadline(t *testing.T) { + tun := defaultTunables() + tun.slowWarmup = 10 + tun.slowWindow = 200 + tun.slowFactor = 3 + tun.slowMinDeadline = 1 * time.Second + + g := &slowAttemptGuard{tun: tun} + if d := g.deadline(); d != tun.slowColdDeadline { + t.Fatalf("before warmup: deadline = %v, want cold-start %v", d, tun.slowColdDeadline) + } + + // 90 fast parts and 10 slow ones: p95 falls in the slow tail (20s), so the + // deadline is 3*20s = 60s, well above the fast bulk but below a real straggler. + for range 90 { + g.record(1 * time.Second) + } + for range 10 { + g.record(20 * time.Second) + } + if d := g.deadline(); d != 60*time.Second { + t.Fatalf("deadline = %v, want 60s (3 x p95 of 20s)", d) + } + + // A fast network (tiny p95) is floored so the deadline never clips variance. + g2 := &slowAttemptGuard{tun: tun} + for range 20 { + g2.record(100 * time.Millisecond) + } + if d := g2.deadline(); d != tun.slowMinDeadline { + t.Fatalf("floored deadline = %v, want %v", d, tun.slowMinDeadline) + } +} + +// TestSendPartSoftDeadlineCancels verifies a wedged attempt (the server never +// responds) is cancelled at the soft deadline and reported as slow, promptly. +func TestSendPartSoftDeadlineCancels(t *testing.T) { + tun := defaultTunables() + tun.slowCheckInterval = 5 * time.Millisecond + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-r.Context().Done(): // the client cancelled at the soft deadline + case <-time.After(3 * time.Second): + w.WriteHeader(http.StatusOK) + } + })) + t.Cleanup(srv.Close) + + e := &engine{c: &Client{}, tun: tun} + uc := &uploadContext{tun: tun, cloud: cloudstorage.New(srv.Client()), limiter: NewLimiter(0)} + + start := time.Now() + _, slow, err := e.sendPart(t.Context(), uc, srv.URL, nil, cloudstorage.BytesBody([]byte("payload")), + func() time.Duration { return 50 * time.Millisecond }) + elapsed := time.Since(start) + + if !slow { + t.Fatalf("slow = false, want true (err=%v, elapsed=%v)", err, elapsed) + } + if elapsed > time.Second { + t.Fatalf("sendPart returned after %v, want prompt cancel near the 50ms deadline", elapsed) + } +} + +// TestSendPartFastNotSlow verifies a fast attempt under a generous deadline is +// not flagged slow and returns the response. +func TestSendPartFastNotSlow(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ETag", "etag-1") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + + tun := defaultTunables() + e := &engine{c: &Client{}, tun: tun} + uc := &uploadContext{tun: tun, cloud: cloudstorage.New(srv.Client()), limiter: NewLimiter(0)} + + resp, slow, err := e.sendPart(t.Context(), uc, srv.URL, nil, cloudstorage.BytesBody([]byte("payload")), + func() time.Duration { return 5 * time.Second }) + if err != nil { + t.Fatalf("sendPart: %v", err) + } + if slow { + t.Fatal("slow = true, want false for a fast attempt") + } + if resp.Header.Get("ETag") != "etag-1" { + t.Errorf("ETag = %q, want etag-1", resp.Header.Get("ETag")) + } +} diff --git a/libs/tmp/files/v2/ext_tunables.go b/libs/tmp/files/v2/ext_tunables.go new file mode 100644 index 00000000000..8bbe75fa130 --- /dev/null +++ b/libs/tmp/files/v2/ext_tunables.go @@ -0,0 +1,171 @@ +package files + +// The large-file upload engine's policy knobs and the engine value that carries +// them. tunables is injected at construction (defaultTunables in production, +// custom values in tests) rather than read from package globals, so a test never +// mutates shared state and uploads with different policies cannot interfere. + +import ( + "time" + + "github.com/databricks/sdk-go/core/ops" +) + +// tunables holds every policy value the engine reads. It is copied by value into +// the engine, the per-upload context, and the straggler guard, so nothing shared +// is mutated after construction. +type tunables struct { + // minStreamSize is the threshold below which a known-size upload is sent in a + // single PUT instead of being split into parts. + minStreamSize int64 + + // defaultPartSize is the part size used when no part size is requested. A + // fixed, modestly sized part (rather than one scaled to keep the part count + // low) bounds the memory a buffering upload holds (about parallelism*partSize + // for a non-seekable stream) and spreads stragglers across more parts. It is + // only grown when a file is large enough to exceed maxUploadParts. + defaultPartSize int64 + + // maxPartSize is the largest part size a cloud provider accepts. + maxPartSize int64 + + // batchURLCount is the number of presigned URLs requested per batch when the + // content length is unknown. + batchURLCount int + + // fileParallelism is the default worker count when no parallelism is requested + // and the source is randomly readable (a local file or other io.ReaderAt). + // Each part streams from a positioned read, so resident memory stays bounded + // to the transport's per-connection buffers regardless of the worker count + // (not parallelism*partSize); a high default saturates a fast uplink, and the + // driver caps the workers at the part count for small files. + fileParallelism int + + // streamMemoryBudget bounds the bytes a buffered upload (a non-seekable + // stream) holds in memory, where each worker keeps one partSize buffer for its + // in-flight part. The default stream worker count is this budget divided by + // the part size, so a larger part size lowers the worker count instead of + // growing memory. At the 16 MiB default part size this is 16 workers. + streamMemoryBudget int64 + + // readAheadBytes is how far past a chunk the resumable upload reads to detect + // end-of-stream (a resumable chunk cannot be empty, so the last chunk must be + // marked with the real total size). + readAheadBytes int64 + + // maxRetries caps how many times a single part is retried after a presigned + // URL expires (multipart) or a chunk must be resumed (resumable). + maxRetries int + + // maxUploadParts is the provider's maximum number of parts per multipart + // upload (S3's 10,000 is the common floor). The default part size is grown for + // files large enough to otherwise exceed it. + maxUploadParts int64 + + // responseHeaderTimeout bounds the response phase of a cloud transfer so a + // connection that accepts the upload but never replies cannot hang a goroutine + // forever. It is applied to the transfer client built in resolveCloudClient. + responseHeaderTimeout time.Duration + + // cleanupTimeout bounds a best-effort abort, which runs on a context detached + // from the upload's (see cleanupContext). + cleanupTimeout time.Duration + + // slowFactor cancels a part attempt that exceeds this multiple of the recent + // p95 part duration. It is the main straggler knob: lower catches stragglers + // sooner but re-issues more healthy parts. + slowFactor int + + // slowWarmup is the number of completed parts required before the guard + // switches from the cold-start deadline to the p95-relative one. + slowWarmup int + + // slowWindow bounds the rolling sample set so the p95 tracks recent conditions + // rather than the whole upload. + slowWindow int + + // slowMinDeadline floors the soft deadline so a fast network (tiny p95) cannot + // produce a deadline that clips normal variance. + slowMinDeadline time.Duration + + // slowColdDeadline is the absolute soft deadline for a part in the opening + // wave, before slowWarmup parts have completed and the p95-relative deadline is + // armed. Because the deadline is re-evaluated in flight (see sendPart), this + // only bites in the first seconds of an upload or for a file too small to warm + // the guard; once warmed, an in-flight early part is caught at the tighter p95 + // deadline instead. It is set above the legitimate opening-burst latency (tens + // of seconds at high concurrency). + slowColdDeadline time.Duration + + // slowFirstPartDeadline is the soft deadline for the synchronous first part, + // which runs alone before any worker (so it has no samples and, unlike the + // opening wave, no burst contention). A healthy first PUT is a few seconds, so + // a tight deadline re-issues a wedged first connection quickly without blocking + // the whole upload at 0% on the cold-start floor. + slowFirstPartDeadline time.Duration + + // slowCheckInterval is how often an in-flight attempt is re-checked against the + // current deadline, so warmup and a falling p95 take effect while the attempt + // is still running. + slowCheckInterval time.Duration + + // slowMaxReissue caps soft re-issues per part; afterward the part rides out the + // normal timeouts rather than being cancelled again, so it is never failed + // solely for staying slow. + slowMaxReissue int + + // cpBackoff is the exponential-backoff-with-jitter policy for control-plane + // retries. + cpBackoff ops.BackoffPolicy + + // cpRetryTimeout bounds the total time spent retrying a single control-plane + // call. Without it, a persistently failing endpoint would be retried until the + // caller's context expires (and forever if it has no deadline). + cpRetryTimeout time.Duration + + // urlExpiry is how long requested presigned URLs are valid. + urlExpiry time.Duration +} + +// defaultTunables returns the production policy values. +func defaultTunables() tunables { + return tunables{ + minStreamSize: 50 * 1024 * 1024, + defaultPartSize: 16 << 20, // 16 MiB + maxPartSize: 4 << 30, + batchURLCount: 1, + fileParallelism: 128, + streamMemoryBudget: 256 << 20, + readAheadBytes: 1, + maxRetries: 3, + maxUploadParts: 10_000, + responseHeaderTimeout: 60 * time.Second, + cleanupTimeout: 30 * time.Second, + slowFactor: 3, + slowWarmup: 32, + slowWindow: 128, + slowMinDeadline: 5 * time.Second, + slowColdDeadline: 30 * time.Second, + slowFirstPartDeadline: 10 * time.Second, + slowCheckInterval: 1 * time.Second, + slowMaxReissue: 2, + cpBackoff: ops.BackoffPolicy{}, + cpRetryTimeout: 5 * time.Minute, + urlExpiry: time.Hour, + } +} + +// engine performs a single or shared-limiter set of uploads for a Client. It +// pairs the Client (for its authenticated transport, host, logger, and workspace +// routing) with the tunables that govern part sizing, retries, and straggler +// control. All internal upload logic hangs off *engine so tests can drive it +// with custom tunables without touching package state. +type engine struct { + c *Client + tun tunables +} + +// newEngine builds an engine over c with the production tunables. +func newEngine(c *Client) *engine { + return &engine{c: c, tun: defaultTunables()} +} diff --git a/libs/tmp/files/v2/ext_upload.go b/libs/tmp/files/v2/ext_upload.go new file mode 100644 index 00000000000..d2715a7d8dd --- /dev/null +++ b/libs/tmp/files/v2/ext_upload.go @@ -0,0 +1,568 @@ +package files + +// Large-file upload engine for Databricks Unity Catalog Volumes. It chooses a +// single-shot, multipart (AWS/Azure), or resumable (GCP) protocol based on the +// stream size and the protocol the workspace's Files API selects. +// +// The engine is pure orchestration over two transports: this client's own +// authenticated Files API calls (the control plane; see upload_control.go), +// which handle the single-shot upload and the multipart/resumable coordination, +// and the unauthenticated cloud-storage client (the cloudstorage subpackage), +// which transfers parts and chunks directly to object storage over the presigned +// URLs the control plane mints. + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "math" + "net/http" + "os" + "sync" + "sync/atomic" + + "github.com/databricks/cli/libs/tmp/files/v2/internal/cloudstorage" +) + +// errUploadURLExpired is returned when a part's presigned URL keeps expiring +// across re-mint retries. An already-exists condition is surfaced as the +// [ErrAlreadyExists] sentinel; a malformed server response is an internal error. +var errUploadURLExpired = errors.New("upload URL expired after retries") + +// fallbackToFilesAPI signals that a multipart or resumable upload failed in a +// way that warrants retrying through a single-shot Files API PUT (for example, +// a cloud firewall rejecting the first part). It carries the bytes already +// buffered but not yet uploaded so they can be replayed. +type fallbackToFilesAPI struct { + buffer []byte + reason string +} + +func (e *fallbackToFilesAPI) Error() string { + return "falling back to single-shot upload: " + e.reason +} + +// UploadResult holds the result of an upload. It is currently empty and exists +// for forward compatibility. +type UploadResult struct{} + +// UploadOption configures an Upload or UploadFrom call. +type UploadOption func(*uploadConfig) + +type uploadConfig struct { + overwrite *bool + partSize int64 + parallelism *int + progress ProgressFunc + transferClient *http.Client + limiter Limiter +} + +// Progress reports the state of an in-flight upload. Fields may be added in +// future releases, so callers must not depend on the struct being comparable or +// on its exact size; always refer to fields by name. +type Progress struct { + // Transferred is the cumulative number of bytes confirmed uploaded so far. + Transferred int64 + // Total is the total size in bytes, or -1 if it is not known in advance (a + // non-seekable stream). + Total int64 +} + +// ProgressFunc is invoked as an upload makes progress. It is called from +// internal goroutines but never concurrently with itself, so it needs no +// locking of its own; it must return promptly. +type ProgressFunc func(Progress) + +// WithOverwrite controls whether an existing file is overwritten. When this +// option is not supplied the parameter is omitted and the server applies its +// default. +func WithOverwrite(overwrite bool) UploadOption { + return func(c *uploadConfig) { c.overwrite = &overwrite } +} + +// WithPartSize sets the multipart part size in bytes. It must not exceed the +// cloud provider maximum. When not supplied an appropriate size is chosen from +// the content length. +func WithPartSize(partSize int64) UploadOption { + return func(c *uploadConfig) { c.partSize = partSize } +} + +// WithParallelism sets the number of concurrent upload workers used for a large +// file. A value of 1 uploads the parts sequentially on a single goroutine; +// higher values upload that many parts at once. It must be at least 1. When not +// set, a default is used. +func WithParallelism(n int) UploadOption { + return func(c *uploadConfig) { c.parallelism = &n } +} + +// WithProgress registers a callback invoked as the upload progresses, reporting +// the cumulative bytes uploaded and the total size (-1 if unknown). It is useful +// for rendering an upload progress bar. +func WithProgress(fn ProgressFunc) UploadOption { + return func(c *uploadConfig) { c.progress = fn } +} + +// WithTransferClient overrides the HTTP client used to transfer file contents +// during a large-file upload, replacing the one the call would otherwise create +// internally. Use it to route the transfer through a custom transport, proxy, or +// CA. +// +// Because these transfers use self-authenticating presigned URLs, the client must +// not attach Databricks credentials; the storage provider rejects extra +// authentication. Do not set http.Client.Timeout, a whole-request deadline that +// would abort a legitimately long part transfer; bound the upload with the context +// passed to Upload instead. +func WithTransferClient(client *http.Client) UploadOption { + return func(c *uploadConfig) { c.transferClient = client } +} + +func resolveUploadConfig(opts []UploadOption, tun tunables) (*uploadConfig, error) { + cfg := &uploadConfig{} + for _, opt := range opts { + opt(cfg) + } + if cfg.parallelism != nil && *cfg.parallelism < 1 { + return nil, fmt.Errorf("parallelism must be at least 1, got %d", *cfg.parallelism) + } + if cfg.partSize < 0 { + return nil, fmt.Errorf("part size must be non-negative, got %d", cfg.partSize) + } + if cfg.partSize > tun.maxPartSize { + return nil, fmt.Errorf("part size %d exceeds maximum %d", cfg.partSize, tun.maxPartSize) + } + return cfg, nil +} + +// resolvedParallelism returns the worker count: the caller's explicit value when +// set, otherwise a default that depends on whether the source is buffered. A +// randomly-readable source streams parts from positioned reads and uses a high, +// bandwidth-oriented default; a buffered (non-seekable) source holds +// parallelism*partSize in memory, so its default is sized to a memory budget, +// capped at the file default so a tiny part size cannot translate the budget into +// an unbounded number of goroutines and sockets. +func (cfg *uploadConfig) resolvedParallelism(buffered bool, partSize int64, tun tunables) int { + if cfg.parallelism != nil { + return *cfg.parallelism + } + if !buffered { + return tun.fileParallelism + } + workers := int(tun.streamMemoryBudget / partSize) + return min(max(workers, 1), tun.fileParallelism) +} + +// newTransferClient returns the default HTTP client for cloud-leg transfers, +// sized for n concurrent transfers: it keeps up to n idle connections to the +// storage host so they are reused rather than re-dialed (Go's default of 2 per +// host would otherwise 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 the +// context instead. Callers who need to supply their own client use +// WithTransferClient. +func newTransferClient(n int, tun tunables) *http.Client { + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.ResponseHeaderTimeout = tun.responseHeaderTimeout + transport.MaxIdleConnsPerHost = n + transport.MaxIdleConns = max(transport.MaxIdleConns, n) + return &http.Client{Transport: transport} +} + +// resolveCloudClient returns the caller-supplied transfer client when set, and +// otherwise a fresh one sized to the upload parallelism. +func (cfg *uploadConfig) resolveCloudClient(parallelism int, tun tunables) *http.Client { + if cfg.transferClient != nil { + return cfg.transferClient + } + return newTransferClient(parallelism, tun) +} + +// uploadContext is the resolved per-upload state. Its fields are set once at +// construction; the slowGuard and completed flag hold the only mutable state, +// both safe for concurrent use by the upload workers. +type uploadContext struct { + tun tunables + targetPath string + overwrite *bool + partSize int64 + batchSize int + contentLength int64 // -1 when unknown (non-seekable stream) + sourcePath string + parallel bool + parallelism int + cloud *cloudstorage.Client + progress *progressReporter + slowGuard *slowAttemptGuard + limiter Limiter // bounds concurrent transfers; never nil (unlimited by default) + + // completed records whether any part has finished uploading. The single-shot + // fallback is taken only while no part has landed; once a part succeeds, a + // later cloud rejection is a real error rather than a blanket-block signal. + completed atomic.Bool +} + +// newUploadContext resolves the per-upload state shared by Upload and UploadFrom. +// contentLength is -1 for a non-seekable stream; sourcePath is "" unless the +// source is a local file (UploadFrom). buffered reports whether the upload path +// holds parts in memory (a non-seekable stream) rather than streaming them from +// positioned reads, which sets the default worker count. +func (e *engine) newUploadContext(cfg *uploadConfig, targetPath, sourcePath string, contentLength int64, buffered bool) *uploadContext { + partSize, batchSize := optimizeParams(contentLength, cfg.partSize, e.tun) + parallelism := cfg.resolvedParallelism(buffered, partSize, e.tun) + limiter := cfg.limiter + if limiter == nil { + limiter = unlimitedLimiter{} + } + return &uploadContext{ + tun: e.tun, + targetPath: targetPath, + overwrite: cfg.overwrite, + partSize: partSize, + batchSize: batchSize, + contentLength: contentLength, + sourcePath: sourcePath, + parallel: parallelism > 1, + parallelism: parallelism, + cloud: cloudstorage.New(cfg.resolveCloudClient(parallelism, e.tun), cloudstorage.WithLogger(e.c.logger)), + progress: newProgressReporter(cfg.progress, contentLength), + slowGuard: &slowAttemptGuard{tun: e.tun}, + limiter: limiter, + } +} + +// progressReporter accumulates confirmed-uploaded bytes and forwards the running +// total to a ProgressFunc. Its methods are safe for concurrent use (parallel +// uploads report from many goroutines) and serialize callbacks so the +// ProgressFunc never runs concurrently with itself. +type progressReporter struct { + fn ProgressFunc + total int64 + + mu sync.Mutex + transferred int64 +} + +// newProgressReporter returns nil when fn is nil so that reporter calls are +// no-ops without the callers having to check. +func newProgressReporter(fn ProgressFunc, total int64) *progressReporter { + if fn == nil { + return nil + } + return &progressReporter{fn: fn, total: total} +} + +// add records n newly uploaded bytes and reports the new cumulative total. It is +// a no-op on a nil reporter, so callers need not check. +func (p *progressReporter) add(n int64) { + if p == nil || n <= 0 { + return + } + p.mu.Lock() + defer p.mu.Unlock() + p.transferred += n + p.fn(Progress{Transferred: p.transferred, Total: p.total}) +} + +// progressReader reports bytes as they are read. It tracks single-shot uploads, +// where the whole body streams through one request rather than discrete parts. +type progressReader struct { + r io.Reader + p *progressReporter +} + +func (pr *progressReader) Read(b []byte) (int, error) { + n, err := pr.r.Read(b) + pr.p.add(int64(n)) + return n, err +} + +// progressSeekReader is a progressReader over a seekable source. It forwards +// Seek so the wrapper preserves the underlying io.Seeker, which the single-shot +// path probes to decide whether a failed attempt can be rewound and retried. +// Kept distinct from progressReader so a non-seekable source is never wrongly +// reported as seekable. +type progressSeekReader struct { + progressReader + s io.Seeker +} + +func (pr *progressSeekReader) Seek(offset int64, whence int) (int64, error) { + return pr.s.Seek(offset, whence) +} + +// withProgress wraps r so reads report to p, preserving r's io.Seeker capability +// when it has one. Without this, wrapping a replayable body (bytes.Reader, +// *os.File, io.SectionReader) would hide its Seeker and make an otherwise +// retriable single-shot upload non-retriable. +func withProgress(r io.Reader, p *progressReporter) io.Reader { + pr := progressReader{r: r, p: p} + if s, ok := r.(io.Seeker); ok { + return &progressSeekReader{progressReader: pr, s: s} + } + return &pr +} + +// Upload uploads contents to the remote file at filePath, which must be an +// absolute path such as "/Volumes/catalog/schema/volume/file". It chooses a +// single-shot, multipart, or resumable upload based on the stream size (when +// the reader is seekable) and the protocol the workspace's Files API selects. +// +// When contents implements io.ReaderAt (for example an *os.File or a +// *bytes.Reader), parts are read from it with concurrent positioned reads rather +// than buffered in memory. A reader without random access has each in-flight part +// buffered as it is read. +// +// ctx bounds the entire operation, including all retries; pass a context with a +// deadline to cap the total upload time. +func (c *Client) Upload(ctx context.Context, filePath string, contents io.Reader, opts ...UploadOption) (*UploadResult, error) { + return newEngine(c).upload(ctx, filePath, contents, opts...) +} + +// UploadFrom uploads the local file at sourcePath to the remote file at +// filePath. The local file is opened as needed; its size is always known, so it +// never takes the non-seekable path. +// +// Parts are read from the file on demand rather than buffered, so the upload +// covers the file as of its size when the call begins: bytes appended afterward +// are not included, and truncating the file mid-upload fails the upload (it does +// not produce a partial object). Do not modify the file while an upload is in +// progress. +// +// ctx bounds the entire operation, including all retries; pass a context with a +// deadline to cap the total upload time. +func (c *Client) UploadFrom(ctx context.Context, filePath, sourcePath string, opts ...UploadOption) (*UploadResult, error) { + return newEngine(c).uploadFrom(ctx, filePath, sourcePath, opts...) +} + +func (e *engine) upload(ctx context.Context, filePath string, contents io.Reader, opts ...UploadOption) (*UploadResult, error) { + cfg, err := resolveUploadConfig(opts, e.tun) + if err != nil { + return nil, err + } + + // A seekable reader lets us learn the size up front; otherwise it is unknown. + // Seeking to the end moves the reader, so a failed rewind afterward is + // unrecoverable: the reader is left at EOF and a later read would send zero + // bytes. Fail loud in that case rather than silently truncating the object. + contentLength := int64(-1) + if seeker, ok := contents.(io.Seeker); ok { + if end, serr := seeker.Seek(0, io.SeekEnd); serr == nil { + if _, serr := seeker.Seek(0, io.SeekStart); serr != nil { + return nil, fmt.Errorf("cannot rewind reader after sizing: %w", serr) + } + contentLength = end + } + } + + // A randomly-readable source large enough for the parallel path streams its + // parts from positioned reads; any other source buffers each in-flight part. + _, isReaderAt := contents.(io.ReaderAt) + buffered := !isReaderAt || contentLength < e.tun.minStreamSize + uc := e.newUploadContext(cfg, filePath, "", contentLength, buffered) + + switch { + case uc.parallel && contentLength < 0: + // Non-seekable stream of unknown size: parts are buffered as they are read. + return &UploadResult{}, e.parallelUploadFromStream(ctx, uc, contents) + case uc.parallel && contentLength >= e.tun.minStreamSize: + // Known-size large upload. A reader supporting concurrent positioned reads + // (io.ReaderAt) is streamed in sections without buffering; anything else + // buffers parts. + if ra, ok := contents.(io.ReaderAt); ok { + return &UploadResult{}, e.parallelUploadFromReaderAt(ctx, uc, ra) + } + return &UploadResult{}, e.parallelUploadFromStream(ctx, uc, contents) + case contentLength >= 0: + return &UploadResult{}, e.uploadSingleThreadKnownSize(ctx, uc, contents) + default: + return &UploadResult{}, e.singleThreadMultipart(ctx, uc, contents) + } +} + +func (e *engine) uploadFrom(ctx context.Context, filePath, sourcePath string, opts ...UploadOption) (*UploadResult, error) { + cfg, err := resolveUploadConfig(opts, e.tun) + if err != nil { + return nil, err + } + info, err := os.Stat(sourcePath) + if err != nil { + return nil, err + } + size := info.Size() + + // A local file is randomly readable, so the parallel path streams parts from + // positioned reads and never buffers them: buffered is false. + uc := e.newUploadContext(cfg, filePath, sourcePath, size, false) + + if uc.parallel && size >= e.tun.minStreamSize { + return &UploadResult{}, e.parallelUploadFromFile(ctx, uc) + } + f, err := os.Open(sourcePath) + if err != nil { + return nil, err + } + defer f.Close() + return &UploadResult{}, e.uploadSingleThreadKnownSize(ctx, uc, f) +} + +// optimizeParams picks a part size and a presigned-URL batch size. contentLength +// is -1 when unknown; override is 0 when no part size was requested. The override +// is already range-checked in resolveUploadConfig. +func optimizeParams(contentLength, override int64, tun tunables) (partSize int64, batchSize int) { + switch { + case override > 0: + partSize = override + default: + partSize = tun.defaultPartSize + // Grow the part size only if the fixed default would exceed the + // provider's maximum part count for this (known) content length. + if contentLength >= 0 { + if minPart := (contentLength + tun.maxUploadParts - 1) / tun.maxUploadParts; minPart > partSize { + partSize = min(minPart, tun.maxPartSize) + } + } + } + + batchSize = tun.batchURLCount + if contentLength >= 0 && partSize > 0 { + numParts := (contentLength + partSize - 1) / partSize + // The square root of the part count is a heuristic that balances the + // number of URL-minting round trips against batch size. + batchSize = max(int(math.Ceil(math.Sqrt(float64(numParts)))), 1) + } + return partSize, batchSize +} + +func (e *engine) uploadSingleThreadKnownSize(ctx context.Context, uc *uploadContext, r io.Reader) error { + if uc.contentLength < e.tun.minStreamSize { + return e.singleShotUpload(ctx, uc, r) + } + return e.singleThreadMultipart(ctx, uc, r) +} + +// singleThreadMultipart pre-reads up to the single-shot threshold to handle +// small (or non-seekable but actually small) streams without a multipart +// round trip, then initiates the upload. +func (e *engine) singleThreadMultipart(ctx context.Context, uc *uploadContext, r io.Reader) error { + preRead, err := readUpTo(r, e.tun.minStreamSize) + if err != nil { + return err + } + if int64(len(preRead)) < e.tun.minStreamSize { + return e.singleShotUpload(ctx, uc, bytes.NewReader(preRead)) + } + return e.initiateAndUpload(ctx, uc, false, + func(token string) error { return e.performMultipartUpload(ctx, uc, token, r, preRead) }, + func(token string) error { return e.performResumableUpload(ctx, uc, token, r, preRead) }, + func(buffered []byte) error { + return e.singleShotUpload(ctx, uc, io.MultiReader(bytes.NewReader(buffered), r)) + }, + ) +} + +func (e *engine) parallelUploadFromStream(ctx context.Context, uc *uploadContext, r io.Reader) error { + return e.initiateAndUpload(ctx, uc, true, + func(token string) error { return e.parallelMultipartFromStream(ctx, uc, token, r) }, + func(token string) error { return e.performResumableUpload(ctx, uc, token, r, nil) }, + func(buffered []byte) error { + return e.singleShotUpload(ctx, uc, io.MultiReader(bytes.NewReader(buffered), r)) + }, + ) +} + +func (e *engine) parallelUploadFromFile(ctx context.Context, uc *uploadContext) error { + // Each branch opens its own handle: only one runs per upload, and the multipart + // driver keeps its handle open until all its workers finish (wait()). + withFile := func(use func(*os.File) error) error { + f, err := os.Open(uc.sourcePath) + if err != nil { + return err + } + defer f.Close() + return use(f) + } + return e.initiateAndUpload(ctx, uc, true, + func(token string) error { + return withFile(func(f *os.File) error { return e.parallelMultipartFromReaderAt(ctx, uc, token, f) }) + }, + func(token string) error { + return withFile(func(f *os.File) error { return e.performResumableUpload(ctx, uc, token, f, nil) }) + }, + func(buffered []byte) error { + return withFile(func(f *os.File) error { return e.singleShotUpload(ctx, uc, f) }) + }, + ) +} + +// parallelUploadFromReaderAt uploads a known-size, randomly-readable source +// (an *os.File or *bytes.Reader passed to Upload) by streaming sections of it +// concurrently, without buffering whole parts. The resumable and fallback paths +// re-read it through fresh section readers. +func (e *engine) parallelUploadFromReaderAt(ctx context.Context, uc *uploadContext, ra io.ReaderAt) error { + return e.initiateAndUpload(ctx, uc, true, + func(token string) error { return e.parallelMultipartFromReaderAt(ctx, uc, token, ra) }, + func(token string) error { + return e.performResumableUpload(ctx, uc, token, io.NewSectionReader(ra, 0, uc.contentLength), nil) + }, + func(buffered []byte) error { + return e.singleShotUpload(ctx, uc, io.NewSectionReader(ra, 0, uc.contentLength)) + }, + ) +} + +// initiateAndUpload initiates the server-coordinated upload, dispatches to the +// multipart or resumable path the server selected (multipart on AWS/Azure, +// resumable on GCP), and applies the shared abort-and-fallback handling. +func (e *engine) initiateAndUpload( + ctx context.Context, + uc *uploadContext, + parallel bool, + performMultipart func(token string) error, + performResumable func(token string) error, + fallbackSingleShot func(buffered []byte) error, +) error { + resp, err := e.initiateUpload(ctx, uc.targetPath, uc.overwrite) + if err != nil { + return err + } + + switch { + case resp.MultipartUpload != nil: + token := resp.MultipartUpload.SessionToken + if token == "" { + return fmt.Errorf("%w: missing multipart session token", errUnexpectedServerResponse) + } + uerr := performMultipart(token) + if uerr == nil { + return nil + } + // Abort is best-effort in both the fallback and the hard-error case. + e.abortMultipartBestEffort(ctx, uc, token) + if fb, ok := errors.AsType[*fallbackToFilesAPI](uerr); ok { + e.c.logger.InfoContext(ctx, "using single-shot fallback", "reason", fb.reason) + return fallbackSingleShot(fb.buffer) + } + return uerr + + case resp.ResumableUpload != nil: + token := resp.ResumableUpload.SessionToken + if token == "" { + return fmt.Errorf("%w: missing resumable session token", errUnexpectedServerResponse) + } + if parallel { + e.c.logger.WarnContext(ctx, "GCP does not support parallel resumable uploads; using single-threaded upload") + } + // The resumable path aborts itself on error, so no abort here. + uerr := performResumable(token) + if fb, ok := errors.AsType[*fallbackToFilesAPI](uerr); ok { + e.c.logger.InfoContext(ctx, "using single-shot fallback", "reason", fb.reason) + return fallbackSingleShot(fb.buffer) + } + return uerr + + default: + return fmt.Errorf("%w: neither multipart nor resumable upload offered", errUnexpectedServerResponse) + } +} diff --git a/libs/tmp/files/v2/ext_upload_control.go b/libs/tmp/files/v2/ext_upload_control.go new file mode 100644 index 00000000000..8beab493414 --- /dev/null +++ b/libs/tmp/files/v2/ext_upload_control.go @@ -0,0 +1,367 @@ +package files + +// Files API control-plane calls for large uploads: initiate a server-coordinated +// session, mint presigned URLs for parts (multipart) and the resumable session +// (GCP), complete a multipart upload, and mint abort URLs for cleanup. It also +// owns the single-shot octet-stream PUT used for small files and as the +// multipart/resumable fallback. +// +// These calls run over the client's own authenticated transport (the same one +// the generated methods use); the parts/chunks then transfer directly to cloud +// storage over the URLs minted here (see the cloudstorage subpackage), carrying +// no Databricks credentials. + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "maps" + "net/http" + "net/url" + "slices" + "strconv" + "time" + + "github.com/databricks/sdk-go/core/apierr" + "github.com/databricks/sdk-go/core/apierr/codes" + "github.com/databricks/sdk-go/core/apiretry" + "github.com/databricks/sdk-go/core/ops" +) + +// Sentinel errors surfaced by the large-file upload API. +var ( + // errUnexpectedServerResponse is returned when the Files API control plane + // returns a well-formed but semantically incomplete response (for example, + // no presigned URLs). It is an internal signal, not a caller-facing sentinel. + errUnexpectedServerResponse = errors.New("unexpected server response") + + // ErrAlreadyExists is returned when an upload targets an existing path with + // overwrite disabled. It is surfaced as a sentinel so callers detect + // "already exists" with errors.Is regardless of which protocol produced it. + ErrAlreadyExists = errors.New("the file being created already exists") +) + +// controlPlaneRetryStatusCodes are the HTTP statuses retried for an +// authenticated control-plane call. +var controlPlaneRetryStatusCodes = []int{ + http.StatusRequestTimeout, // 408 + http.StatusTooManyRequests, // 429 + http.StatusBadGateway, // 502 + http.StatusServiceUnavailable, // 503 + http.StatusGatewayTimeout, // 504 +} + +// --- Result types --- + +// initiateResult reports which large-file protocol the server selected. Exactly +// one of multipartUpload (AWS/Azure) and resumableUpload (GCP) is non-nil on a +// well-formed response; the pointers distinguish "not offered" from "offered +// with an empty token". +type initiateResult struct { + MultipartUpload *uploadSession `json:"multipart_upload"` + ResumableUpload *uploadSession `json:"resumable_upload"` +} + +type uploadSession struct { + SessionToken string `json:"session_token"` +} + +// presignedURL is a resolved cloud-storage URL with its associated request +// headers, as minted by the control plane and transferred over by the caller. +type presignedURL struct { + URL string + PartNumber int + Headers map[string]string +} + +// --- Wire types for the Files API multipart/resumable coordination --- + +type nameValue struct { + Name string `json:"name"` + Value string `json:"value"` +} + +type createPartURLsRequest struct { + Path string `json:"path"` + SessionToken string `json:"session_token"` + StartPartNumber int `json:"start_part_number"` + Count int `json:"count"` + ExpireTime string `json:"expire_time"` +} + +type createPartURLsResponse struct { + UploadPartURLs []struct { + URL string `json:"url"` + PartNumber int `json:"part_number"` + Headers []nameValue `json:"headers"` + } `json:"upload_part_urls"` +} + +type resumableURLRequest struct { + Path string `json:"path"` + SessionToken string `json:"session_token"` +} + +// urlWithHeaders is a presigned cloud-storage URL and its required request +// headers, as returned by the resumable and abort URL endpoints. +type urlWithHeaders struct { + URL string `json:"url"` + Headers []nameValue `json:"headers"` +} + +type resumableURLResponse struct { + ResumableUploadURL *urlWithHeaders `json:"resumable_upload_url"` +} + +type abortURLRequest struct { + Path string `json:"path"` + SessionToken string `json:"session_token"` + ExpireTime string `json:"expire_time"` +} + +type abortURLResponse struct { + AbortUploadURL *urlWithHeaders `json:"abort_upload_url"` +} + +type completePart struct { + PartNumber int `json:"part_number"` + ETag string `json:"etag"` +} + +type completeRequest struct { + Parts []completePart `json:"parts"` +} + +// uploadSingleShot sends the whole body in one PUT. Used for small files and as +// the multipart/resumable fallback. overwrite is tri-state: a non-nil value is +// sent explicitly; nil lets the server apply its default. A 409 ALREADY_EXISTS +// is mapped to ErrAlreadyExists. +// +// The generated UploadFile is unsuitable here: it JSON-marshals the reader +// rather than streaming the raw octet stream the Files API expects. A seekable +// body is rewound before each retry; a non-seekable body cannot be replayed, so +// it is attempted once. +func (e *engine) uploadSingleShot(ctx context.Context, path string, overwrite *bool, body io.Reader) error { + query := url.Values{} + if overwrite != nil { + query.Set("overwrite", strconv.FormatBool(*overwrite)) + } + urlStr, err := e.controlPlaneURL(filesAPIPath(path), query) + if err != nil { + return err + } + + seeker, seekable := body.(io.Seeker) + var start int64 + if seekable { + if start, err = seeker.Seek(0, io.SeekCurrent); err != nil { + seekable = false + } + } + + call := func(ctx context.Context) error { + if seekable { + if _, err := seeker.Seek(start, io.SeekStart); err != nil { + return fmt.Errorf("rewinding upload body: %w", err) + } + } + req, err := http.NewRequestWithContext(ctx, http.MethodPut, urlStr, body) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/octet-stream") + e.setWorkspaceHeader(req) + _, _, err = executeHTTPCall(httpCallOptions{req: req, client: e.c.httpClient, logger: e.c.logger}) + return err + } + + // Only retry when the body can be rewound; a consumed non-seekable stream + // cannot be replayed, so it is attempted exactly once. + opts := []ops.Option{ops.WithTimeout(e.tun.cpRetryTimeout)} + if seekable { + opts = append(opts, ops.WithRetrier(e.newControlPlaneRetrier)) + } + return asAlreadyExists(ops.Execute(ctx, call, opts...)) +} + +// initiateUpload opens a server-coordinated upload session for path. The result +// reports which protocol the workspace selected (multipart on AWS/Azure, +// resumable on GCP). +func (e *engine) initiateUpload(ctx context.Context, path string, overwrite *bool) (*initiateResult, error) { + query := url.Values{"action": {"initiate-upload"}} + if overwrite != nil { + query.Set("overwrite", strconv.FormatBool(*overwrite)) + } + var out initiateResult + if err := e.controlPlaneJSON(ctx, http.MethodPost, filesAPIPath(path), query, nil, &out); err != nil { + return nil, asAlreadyExists(err) + } + return &out, nil +} + +// createPartURLs mints count presigned URLs for multipart parts starting at +// startPart. +func (e *engine) createPartURLs(ctx context.Context, path, token string, startPart, count int) ([]presignedURL, error) { + body := createPartURLsRequest{ + Path: path, + SessionToken: token, + StartPartNumber: startPart, + Count: count, + ExpireTime: e.expireTime(), + } + var out createPartURLsResponse + if err := e.controlPlaneJSON(ctx, http.MethodPost, "/api/2.0/fs/create-upload-part-urls", nil, body, &out); err != nil { + return nil, err + } + if len(out.UploadPartURLs) == 0 { + return nil, fmt.Errorf("%w: no upload part URLs returned", errUnexpectedServerResponse) + } + result := make([]presignedURL, 0, len(out.UploadPartURLs)) + for _, p := range out.UploadPartURLs { + result = append(result, presignedURL{URL: p.URL, PartNumber: p.PartNumber, Headers: headerMap(p.Headers)}) + } + return result, nil +} + +// createResumableURL mints the single presigned URL for a GCP resumable session. +func (e *engine) createResumableURL(ctx context.Context, path, token string) (presignedURL, error) { + body := resumableURLRequest{Path: path, SessionToken: token} + var out resumableURLResponse + if err := e.controlPlaneJSON(ctx, http.MethodPost, "/api/2.0/fs/create-resumable-upload-url", nil, body, &out); err != nil { + return presignedURL{}, err + } + if out.ResumableUploadURL == nil || out.ResumableUploadURL.URL == "" { + return presignedURL{}, fmt.Errorf("%w: no resumable upload URL returned", errUnexpectedServerResponse) + } + return presignedURL{URL: out.ResumableUploadURL.URL, Headers: headerMap(out.ResumableUploadURL.Headers)}, nil +} + +// completeMultipart finalizes a multipart upload with the part ETags, which it +// sorts by part number before sending. +func (e *engine) completeMultipart(ctx context.Context, path, token string, etags map[int]string) error { + nums := slices.Sorted(maps.Keys(etags)) + parts := make([]completePart, 0, len(nums)) + for _, n := range nums { + parts = append(parts, completePart{PartNumber: n, ETag: etags[n]}) + } + query := url.Values{"action": {"complete-upload"}, "upload_type": {"multipart"}, "session_token": {token}} + return asAlreadyExists(e.controlPlaneJSON(ctx, http.MethodPost, filesAPIPath(path), query, completeRequest{Parts: parts}, nil)) +} + +// createAbortURL mints a presigned URL for aborting an upload session. The +// caller issues the (unauthenticated) cloud DELETE against it. +func (e *engine) createAbortURL(ctx context.Context, path, token string) (presignedURL, error) { + body := abortURLRequest{Path: path, SessionToken: token, ExpireTime: e.expireTime()} + var out abortURLResponse + if err := e.controlPlaneJSON(ctx, http.MethodPost, "/api/2.0/fs/create-abort-upload-url", nil, body, &out); err != nil { + return presignedURL{}, err + } + if out.AbortUploadURL == nil || out.AbortUploadURL.URL == "" { + return presignedURL{}, fmt.Errorf("%w: no abort upload URL returned", errUnexpectedServerResponse) + } + return presignedURL{URL: out.AbortUploadURL.URL, Headers: headerMap(out.AbortUploadURL.Headers)}, nil +} + +// controlPlaneJSON performs an authenticated JSON request against the Files API +// control plane, retrying transient failures via core/ops. reqBody and out may +// be nil. A fresh request is built on each attempt so retries re-apply +// credentials (the auth transport handles token refresh) and rewind the body. +func (e *engine) controlPlaneJSON(ctx context.Context, method, path string, query url.Values, reqBody, out any) error { + urlStr, err := e.controlPlaneURL(path, query) + if err != nil { + return err + } + var bodyBytes []byte + if reqBody != nil { + if bodyBytes, err = json.Marshal(reqBody); err != nil { + return err + } + } + + call := func(ctx context.Context) error { + var rdr io.Reader + if bodyBytes != nil { + rdr = bytes.NewReader(bodyBytes) + } + req, err := http.NewRequestWithContext(ctx, method, urlStr, rdr) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + e.setWorkspaceHeader(req) + respBody, _, err := executeHTTPCall(httpCallOptions{req: req, client: e.c.httpClient, logger: e.c.logger}) + if err != nil { + return err + } + if out != nil && len(respBody) > 0 { + return json.Unmarshal(respBody, out) + } + return nil + } + return ops.Execute(ctx, call, ops.WithRetrier(e.newControlPlaneRetrier), ops.WithTimeout(e.tun.cpRetryTimeout)) +} + +func (e *engine) newControlPlaneRetrier() ops.Retrier { + return apiretry.NewRetrier(e.tun.cpBackoff, apiretry.RetrierConfig{StatusCodes: controlPlaneRetryStatusCodes}) +} + +// setWorkspaceHeader applies the workspace routing header when the client is +// workspace-scoped. The Files API control plane routes a request to the right +// workspace by this header; the auth transport supplies only the credentials. +func (e *engine) setWorkspaceHeader(req *http.Request) { + if e.c.workspaceID != "" { + req.Header.Set("X-Databricks-Workspace-Id", e.c.workspaceID) + } +} + +// controlPlaneURL joins the client host with a path and optional query, +// mirroring the generated client's URL construction. +func (e *engine) controlPlaneURL(path string, query url.Values) (string, error) { + baseURL, err := url.Parse(e.c.host) + if err != nil { + return "", err + } + baseURL.Path = path + if len(query) > 0 { + baseURL.RawQuery = query.Encode() + } + return baseURL.String(), nil +} + +// asAlreadyExists maps an "already exists" API error to the ErrAlreadyExists +// sentinel, so callers detect it with errors.Is regardless of which path +// surfaced it. The Files API returns HTTP 409 with error_code ALREADY_EXISTS +// when an upload targets an existing path with overwrite=false; for multipart +// this surfaces from complete-upload (after the parts are sent). Every path runs +// over core/apierr, so a single code check suffices. Other errors pass through +// unchanged. +func asAlreadyExists(err error) error { + if apierr.Code(err) == codes.AlreadyExists { + return ErrAlreadyExists + } + return err +} + +// filesAPIPath builds the Files API URL path for an absolute volume path. It +// matches the single-shot UploadFile path construction so the control-plane +// initiate/complete calls escape the path identically. +func filesAPIPath(absPath string) string { + return fmt.Sprintf("/api/2.0/fs/files%v", absPath) +} + +// expireTime is the presigned-URL expiry timestamp sent to the control plane, +// urlExpiry into the future. +func (e *engine) expireTime() string { + return time.Now().UTC().Add(e.tun.urlExpiry).Format("2006-01-02T15:04:05Z") +} + +func headerMap(headers []nameValue) map[string]string { + out := make(map[string]string, len(headers)) + for _, h := range headers { + out[h.Name] = h.Value + } + return out +} diff --git a/libs/tmp/files/v2/ext_upload_control_test.go b/libs/tmp/files/v2/ext_upload_control_test.go new file mode 100644 index 00000000000..4500f73c54f --- /dev/null +++ b/libs/tmp/files/v2/ext_upload_control_test.go @@ -0,0 +1,206 @@ +package files + +import ( + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// newControlEngine builds an engine whose control plane points at an httptest +// server running h. +func newControlEngine(t *testing.T, h http.HandlerFunc) *engine { + t.Helper() + srv := httptest.NewServer(h) + t.Cleanup(srv.Close) + return newEngine(buildUploadClient(t, srv.URL, srv.Client(), "")) +} + +func TestInitiate(t *testing.T) { + c := newControlEngine(t, func(w http.ResponseWriter, r *http.Request) { + if !strings.HasPrefix(r.URL.Path, "/api/2.0/fs/files/") { + t.Errorf("path = %q", r.URL.Path) + } + if got := r.URL.Query().Get("action"); got != "initiate-upload" { + t.Errorf("action = %q", got) + } + if got := r.URL.Query().Get("overwrite"); got != "true" { + t.Errorf("overwrite = %q, want true", got) + } + _, _ = io.WriteString(w, `{"multipart_upload":{"session_token":"tok"}}`) + }) + + ow := true + res, err := c.initiateUpload(t.Context(), "/Volumes/c/s/v/f.bin", &ow) + if err != nil { + t.Fatal(err) + } + if res.MultipartUpload == nil || res.MultipartUpload.SessionToken != "tok" { + t.Errorf("multipart = %+v", res.MultipartUpload) + } + if res.ResumableUpload != nil { + t.Error("resumable should be nil") + } +} + +func TestCreatePartURLs(t *testing.T) { + c := newControlEngine(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, `{"upload_part_urls":[{"url":"https://cloud.test/part/1","part_number":1,"headers":[{"name":"x-h","value":"v"}]}]}`) + }) + + urls, err := c.createPartURLs(t.Context(), "/Volumes/c/s/v/f.bin", "tok", 1, 1) + if err != nil { + t.Fatal(err) + } + if len(urls) != 1 { + t.Fatalf("got %d URLs, want 1", len(urls)) + } + if urls[0].URL != "https://cloud.test/part/1" || urls[0].PartNumber != 1 || urls[0].Headers["x-h"] != "v" { + t.Errorf("url = %+v", urls[0]) + } +} + +func TestCreatePartURLsEmpty(t *testing.T) { + c := newControlEngine(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, `{"upload_part_urls":[]}`) + }) + if _, err := c.createPartURLs(t.Context(), "/Volumes/c/s/v/f.bin", "tok", 1, 1); !errors.Is(err, errUnexpectedServerResponse) { + t.Fatalf("err = %v, want errUnexpectedServerResponse", err) + } +} + +func TestCreateResumableURL(t *testing.T) { + c := newControlEngine(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, `{"resumable_upload_url":{"url":"https://cloud.test/resumable","headers":[{"name":"x-h","value":"v"}]}}`) + }) + + purl, err := c.createResumableURL(t.Context(), "/Volumes/c/s/v/f.bin", "tok") + if err != nil { + t.Fatal(err) + } + if purl.URL != "https://cloud.test/resumable" || purl.Headers["x-h"] != "v" { + t.Errorf("url = %+v", purl) + } +} + +func TestCreateResumableURLMissing(t *testing.T) { + c := newControlEngine(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, `{}`) + }) + if _, err := c.createResumableURL(t.Context(), "/Volumes/c/s/v/f.bin", "tok"); !errors.Is(err, errUnexpectedServerResponse) { + t.Fatalf("err = %v, want errUnexpectedServerResponse", err) + } +} + +func TestCreateAbortURL(t *testing.T) { + c := newControlEngine(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, `{"abort_upload_url":{"url":"https://cloud.test/abort"}}`) + }) + + purl, err := c.createAbortURL(t.Context(), "/Volumes/c/s/v/f.bin", "tok") + if err != nil { + t.Fatal(err) + } + if purl.URL != "https://cloud.test/abort" { + t.Errorf("url = %q", purl.URL) + } +} + +func TestCompleteMultipartSortsParts(t *testing.T) { + var got []int + c := newControlEngine(t, func(w http.ResponseWriter, r *http.Request) { + if action := r.URL.Query().Get("action"); action != "complete-upload" { + t.Errorf("action = %q", action) + } + var body struct { + Parts []struct { + PartNumber int `json:"part_number"` + ETag string `json:"etag"` + } `json:"parts"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + for _, p := range body.Parts { + got = append(got, p.PartNumber) + } + w.WriteHeader(http.StatusOK) + }) + + err := c.completeMultipart(t.Context(), "/Volumes/c/s/v/f.bin", "tok", map[int]string{3: "e3", 1: "e1", 2: "e2"}) + if err != nil { + t.Fatal(err) + } + if len(got) != 3 || got[0] != 1 || got[1] != 2 || got[2] != 3 { + t.Errorf("parts = %v, want sorted [1 2 3]", got) + } +} + +func TestControlPlaneAlreadyExists(t *testing.T) { + conflict := func(errorCode string) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusConflict) + _, _ = io.WriteString(w, `{"error_code":"`+errorCode+`","message":"the file being created already exists"}`) + } + } + + // complete-upload (the multipart finish) returning 409 ALREADY_EXISTS maps to + // the sentinel so callers can skip-if-exists, matching single-shot/resumable. + c := newControlEngine(t, conflict("ALREADY_EXISTS")) + if err := c.completeMultipart(t.Context(), "/Volumes/c/s/v/f.bin", "tok", map[int]string{1: "e1"}); !errors.Is(err, ErrAlreadyExists) { + t.Errorf("completeMultipart err = %v, want ErrAlreadyExists", err) + } + + // initiate-upload maps the same way. + c2 := newControlEngine(t, conflict("ALREADY_EXISTS")) + if _, err := c2.initiateUpload(t.Context(), "/Volumes/c/s/v/f.bin", nil); !errors.Is(err, ErrAlreadyExists) { + t.Errorf("initiateUpload err = %v, want ErrAlreadyExists", err) + } + + // A different 409 is not mistaken for the already-exists sentinel. + c3 := newControlEngine(t, conflict("RESOURCE_CONFLICT")) + if _, err := c3.initiateUpload(t.Context(), "/Volumes/c/s/v/f.bin", nil); errors.Is(err, ErrAlreadyExists) { + t.Errorf("a non-ALREADY_EXISTS 409 mapped to ErrAlreadyExists: %v", err) + } +} + +func TestSingleShotUploadOverwrite(t *testing.T) { + // The single-shot octet-stream PUT streams the raw body and sends the + // overwrite query only when explicitly set. + yes, no := true, false + cases := []struct { + name string + overwrite *bool + wantQuery string + }{ + {"unset", nil, ""}, + {"true", &yes, "true"}, + {"false", &no, "false"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var gotBody, gotOverwrite, gotContentType string + c := newControlEngine(t, func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + gotBody = string(b) + gotOverwrite = r.URL.Query().Get("overwrite") + gotContentType = r.Header.Get("Content-Type") + w.WriteHeader(http.StatusOK) + }) + err := c.uploadSingleShot(t.Context(), "/Volumes/c/s/v/f.bin", tc.overwrite, strings.NewReader("hello")) + if err != nil { + t.Fatal(err) + } + if gotBody != "hello" { + t.Errorf("body = %q, want the raw octet stream %q", gotBody, "hello") + } + if gotContentType != "application/octet-stream" { + t.Errorf("Content-Type = %q, want application/octet-stream", gotContentType) + } + if gotOverwrite != tc.wantQuery { + t.Errorf("overwrite query = %q, want %q", gotOverwrite, tc.wantQuery) + } + }) + } +} diff --git a/libs/tmp/files/v2/ext_upload_helpers.go b/libs/tmp/files/v2/ext_upload_helpers.go new file mode 100644 index 00000000000..99592a3b8b2 --- /dev/null +++ b/libs/tmp/files/v2/ext_upload_helpers.go @@ -0,0 +1,116 @@ +package files + +// Shared buffering and cloud-leg header helpers for the large-file uploader, +// plus the best-effort abort glue between the orchestration and the two +// transports. The cloud-transport pieces (part bodies, retry classification, +// presigned-URL-expiry detection) live in the cloudstorage subpackage. + +import ( + "context" + "io" + "maps" + "net/http" +) + +// --- Buffer helpers --- + +// fillBuffer reads from r until buf holds at least minSize bytes or the stream +// ends. A short read at end-of-stream is not an error. +func fillBuffer(buf []byte, minSize int64, r io.Reader) ([]byte, error) { + need := minSize - int64(len(buf)) + if need <= 0 { + return buf, nil + } + tmp := make([]byte, need) + n, err := io.ReadFull(r, tmp) + buf = append(buf, tmp[:n]...) + if err == nil || err == io.EOF || err == io.ErrUnexpectedEOF { + return buf, nil + } + return buf, err +} + +// readUpTo reads up to n bytes from r, returning fewer only at end-of-stream. +func readUpTo(r io.Reader, n int64) ([]byte, error) { + buf := make([]byte, n) + read, err := io.ReadFull(r, buf) + if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF { + return nil, err + } + return buf[:read], nil +} + +// --- Header helpers --- + +func mergeHeaders(base, override map[string]string) map[string]string { + out := make(map[string]string, len(base)+len(override)) + maps.Copy(out, base) + maps.Copy(out, override) + return out +} + +// octetStreamHeaders returns the headers for a cloud-storage request: the binary +// content type plus the presigned URL's own headers (which win on conflict). The +// returned map is freshly allocated, so callers may add request-specific headers +// such as Content-Range. +func octetStreamHeaders(presigned map[string]string) map[string]string { + return mergeHeaders(map[string]string{"Content-Type": "application/octet-stream"}, presigned) +} + +// --- Single-shot delegation and best-effort abort glue --- + +// singleShotUpload sends the whole body in one PUT through the control-plane +// single-shot path. Used for small files and as the multipart/resumable +// fallback. Progress is reported as the body streams. +func (e *engine) singleShotUpload(ctx context.Context, uc *uploadContext, body io.Reader) error { + if uc.progress != nil { + body = withProgress(body, uc.progress) + } + // A single-shot upload is one transfer; hold one slot for its duration so it + // draws from the same concurrency budget as multipart parts. + if err := uc.limiter.Acquire(ctx); err != nil { + return err + } + defer uc.limiter.Release() + return e.uploadSingleShot(ctx, uc.targetPath, uc.overwrite, body) +} + +// abortMultipartUpload mints a presigned abort URL on the control plane and +// issues the unauthenticated cloud DELETE against it. Both halves run on the +// caller's context, which abortMultipartBestEffort detaches and bounds. +func (e *engine) abortMultipartUpload(ctx context.Context, uc *uploadContext, token string) error { + purl, err := e.createAbortURL(ctx, uc.targetPath, token) + if err != nil { + return err + } + headers := octetStreamHeaders(purl.Headers) + _, err = uc.cloud.Send(ctx, http.MethodDelete, purl.URL, headers, nil) + return err +} + +func (e *engine) abortMultipartBestEffort(ctx context.Context, uc *uploadContext, token string) { + ctx, cancel := e.cleanupContext(ctx) + defer cancel() + if err := e.abortMultipartUpload(ctx, uc, token); err != nil { + // Best-effort cleanup; its failure is not user-actionable (e.g. some clouds + // do not support abort presigned URLs), so it stays at debug rather than + // surfacing as a warning on an otherwise-normal outcome like a skip. + e.c.logger.DebugContext(ctx, "failed to abort multipart upload", "error", err) + } +} + +// cleanupContext returns a context for a best-effort cleanup (aborting a partial +// upload) that is detached from the caller's cancellation and deadline but +// bounded by its own timeout. Cleanup most needs to run exactly when the upload +// context has already been cancelled or has expired; reusing it would make the +// abort fail instantly and leak the partial upload on the storage provider. +func (e *engine) cleanupContext(ctx context.Context) (context.Context, context.CancelFunc) { + return context.WithTimeout(context.WithoutCancel(ctx), e.tun.cleanupTimeout) +} + +func (e *engine) abortResumableUpload(ctx context.Context, uc *uploadContext, purl presignedURL) error { + ctx, cancel := e.cleanupContext(ctx) + defer cancel() + _, err := uc.cloud.Send(ctx, http.MethodDelete, purl.URL, purl.Headers, nil) + return err +} diff --git a/libs/tmp/files/v2/ext_upload_test.go b/libs/tmp/files/v2/ext_upload_test.go new file mode 100644 index 00000000000..dc07b0bc73d --- /dev/null +++ b/libs/tmp/files/v2/ext_upload_test.go @@ -0,0 +1,726 @@ +package files + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/databricks/cli/libs/tmp/auth" + "github.com/databricks/sdk-go/core/ops" + "github.com/databricks/cli/libs/tmp/options/client" +) + +// stubCredentials is a no-op credential for tests: NewClient requires +// credentials, but these tests exercise the control plane through a fake server +// that ignores auth headers. +type stubCredentials struct{} + +func (stubCredentials) Name() string { return "stub" } +func (stubCredentials) AuthHeaders(context.Context) ([]auth.Header, error) { return nil, nil } + +// buildUploadClient builds a v2 Client whose control plane is served by hc and +// addressed at host, with an optional workspace ID for the routing header. +func buildUploadClient(t *testing.T, host string, hc *http.Client, workspaceID string) *Client { + t.Helper() + opts := []client.Option{ + client.WithHost(host), + client.WithHTTPClient(hc), + client.WithCredentials(stubCredentials{}), + client.WithLogger(slog.New(slog.NewTextHandler(io.Discard, nil))), + // Keep tests hermetic: without this, an unset workspace ID would be + // filled from the developer's local profile. + client.WithoutProfileResolution(), + } + if workspaceID != "" { + opts = append(opts, client.WithWorkspaceID(workspaceID)) + } + c, err := NewClient(context.Background(), opts...) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + return c +} + +// fakeServer emulates both the Files API control plane and the cloud storage +// provider: the presigned URLs it mints point back at itself, so a single +// httptest server backs the whole upload. It reassembles the uploaded object so +// tests can assert byte-for-byte equality with the input. +type fakeServer struct { + srv *httptest.Server + + mu sync.Mutex + mode string // "multipart" or "resumable" + parts map[int][]byte + completed []completePart + singleBody []byte + resumable []byte + overwriteQ []string // overwrite query values seen on initiate + single-shot + partHeaders []http.Header // headers seen on cloud part PUTs + partURLReqs []int // Count of each create-upload-part-urls request + + // partHook injects faults. It returns (statusCode, etag, body); a zero + // statusCode means "use the default 200 + store". + partHook func(n, attempt int) (int, string, []byte) + partTry map[int]int + + // singleShotStatus, when non-zero, is returned by the single-shot PUT. + singleShotStatus int + singleShotBody string + + // singleShotHook injects per-attempt faults on the single-shot PUT. It + // returns the status for the given 1-based attempt; a zero status means "use + // the default 200 + store". singleTry counts single-shot PUTs seen. + singleShotHook func(attempt int) int + singleTry int + + // stuckResumable makes a resumable chunk PUT confirm no new bytes, modeling a + // server/proxy wedged at a fixed offset. + stuckResumable bool + + // resumableChunkHook injects per-chunk faults on the resumable data PUT. For a + // non-zero returned status the chunk's bytes are still committed (so the + // client's status query reveals forward progress), but the response is that + // status -- modeling a spurious transient failure on a chunk that did land. + // resumableTry counts data-chunk PUTs (not status queries). + resumableChunkHook func(attempt int) int + resumableTry int +} + +func newFakeServer(t *testing.T, mode string) *fakeServer { + f := &fakeServer{ + mode: mode, + parts: map[int][]byte{}, + partTry: map[int]int{}, + } + f.srv = httptest.NewServer(f) + t.Cleanup(f.srv.Close) + return f +} + +func (f *fakeServer) base() string { return f.srv.URL } + +func (f *fakeServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + path := r.URL.Path + switch { + case r.Method == http.MethodPost && strings.HasSuffix(path, "/create-upload-part-urls"): + f.handleCreatePartURLs(w, r) + case r.Method == http.MethodPost && strings.HasSuffix(path, "/create-resumable-upload-url"): + writeJSON(w, resumableURLResponse{ResumableUploadURL: &urlWithHeaders{URL: f.base() + "/cloud/resumable"}}) + case r.Method == http.MethodPost && strings.HasSuffix(path, "/create-abort-upload-url"): + writeJSON(w, abortURLResponse{AbortUploadURL: &urlWithHeaders{URL: f.base() + "/cloud/abort"}}) + case r.Method == http.MethodPut && strings.HasPrefix(path, "/cloud/part/"): + f.handlePartPut(w, r) + case r.Method == http.MethodPut && path == "/cloud/resumable": + f.handleResumablePut(w, r) + case r.Method == http.MethodDelete && (path == "/cloud/abort" || path == "/cloud/resumable"): + w.WriteHeader(http.StatusOK) + case strings.Contains(path, "/api/2.0/fs/files/"): + f.handleFiles(w, r) + default: + http.Error(w, "unexpected request: "+r.Method+" "+path, http.StatusInternalServerError) + } +} + +func (f *fakeServer) handleFiles(w http.ResponseWriter, r *http.Request) { + action := r.URL.Query().Get("action") + switch { + case r.Method == http.MethodPost && action == "initiate-upload": + f.mu.Lock() + f.overwriteQ = append(f.overwriteQ, r.URL.Query().Get("overwrite")) + f.mu.Unlock() + if f.mode == "resumable" { + writeJSON(w, initiateResult{ResumableUpload: &uploadSession{SessionToken: "tok"}}) + } else { + writeJSON(w, initiateResult{MultipartUpload: &uploadSession{SessionToken: "tok"}}) + } + case r.Method == http.MethodPost && action == "complete-upload": + var req completeRequest + _ = json.NewDecoder(r.Body).Decode(&req) + f.mu.Lock() + f.completed = req.Parts + f.mu.Unlock() + w.WriteHeader(http.StatusOK) + case r.Method == http.MethodPut: // single-shot + body, _ := io.ReadAll(r.Body) + f.mu.Lock() + f.singleBody = body + f.overwriteQ = append(f.overwriteQ, r.URL.Query().Get("overwrite")) + f.singleTry++ + status, respBody, hook, attempt := f.singleShotStatus, f.singleShotBody, f.singleShotHook, f.singleTry + f.mu.Unlock() + if hook != nil { + if s := hook(attempt); s != 0 { + w.WriteHeader(s) + return + } + w.WriteHeader(http.StatusOK) + return + } + if status != 0 { + w.WriteHeader(status) + _, _ = io.WriteString(w, respBody) + return + } + w.WriteHeader(http.StatusOK) + default: + http.Error(w, "unexpected files request", http.StatusInternalServerError) + } +} + +func (f *fakeServer) handleCreatePartURLs(w http.ResponseWriter, r *http.Request) { + var req createPartURLsRequest + _ = json.NewDecoder(r.Body).Decode(&req) + f.mu.Lock() + f.partURLReqs = append(f.partURLReqs, req.Count) + f.mu.Unlock() + var out createPartURLsResponse + for i := range req.Count { + n := req.StartPartNumber + i + out.UploadPartURLs = append(out.UploadPartURLs, struct { + URL string `json:"url"` + PartNumber int `json:"part_number"` + Headers []nameValue `json:"headers"` + }{URL: f.base() + "/cloud/part/" + strconv.Itoa(n), PartNumber: n}) + } + writeJSON(w, out) +} + +func (f *fakeServer) handlePartPut(w http.ResponseWriter, r *http.Request) { + n, _ := strconv.Atoi(strings.TrimPrefix(r.URL.Path, "/cloud/part/")) + body, _ := io.ReadAll(r.Body) + + f.mu.Lock() + f.partHeaders = append(f.partHeaders, r.Header.Clone()) + f.partTry[n]++ + attempt := f.partTry[n] + hook := f.partHook + f.mu.Unlock() + + if hook != nil { + if status, etag, respBody := hook(n, attempt); status != 0 { + if etag != "" { + w.Header().Set("ETag", etag) + } + w.WriteHeader(status) + _, _ = w.Write(respBody) + return + } + } + + f.mu.Lock() + f.parts[n] = body + f.mu.Unlock() + w.Header().Set("ETag", "etag-"+strconv.Itoa(n)) + w.WriteHeader(http.StatusOK) +} + +func (f *fakeServer) handleResumablePut(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + cr := r.Header.Get("Content-Range") + // Status query: "bytes */*" -> report current confirmed offset. + if cr == "bytes */*" { + f.mu.Lock() + n := len(f.resumable) + f.mu.Unlock() + if n == 0 { + w.WriteHeader(http.StatusPermanentRedirect) + return + } + w.Header().Set("Range", "bytes=0-"+strconv.Itoa(n-1)) + w.WriteHeader(http.StatusPermanentRedirect) + return + } + + // A wedged server: accept a non-final chunk but confirm no new bytes (no Range + // header), so the client sees a zero-progress 308 and must bound its retries. + if f.stuckResumable && strings.HasSuffix(cr, "/*") { + w.WriteHeader(http.StatusPermanentRedirect) + return + } + + f.mu.Lock() + f.resumable = append(f.resumable, body...) + total := len(f.resumable) + f.resumableTry++ + hook, attempt := f.resumableChunkHook, f.resumableTry + f.mu.Unlock() + + // The chunk's bytes are committed above; a hooked transient status models a + // blip on a chunk that actually landed, so the client's follow-up status + // query sees forward progress. + if hook != nil { + if s := hook(attempt); s != 0 { + w.WriteHeader(s) + return + } + } + + // Content-Range: bytes {start}-{end}/{total|*} + final := !strings.HasSuffix(cr, "/*") + if final { + w.WriteHeader(http.StatusOK) + return + } + w.Header().Set("Range", "bytes=0-"+strconv.Itoa(total-1)) + w.WriteHeader(http.StatusPermanentRedirect) +} + +// assembled returns the bytes the server reconstructed for the upload. +func (f *fakeServer) assembled() []byte { + f.mu.Lock() + defer f.mu.Unlock() + if f.singleBody != nil { + return f.singleBody + } + if f.mode == "resumable" { + return f.resumable + } + var out []byte + for _, p := range f.completed { + out = append(out, f.parts[p.PartNumber]...) + } + return out +} + +func writeJSON(w http.ResponseWriter, v any) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(v) +} + +func newTestClient(t *testing.T, f *fakeServer) *Client { + t.Helper() + return buildUploadClient(t, f.base(), f.srv.Client(), "") +} + +func TestSetWorkspaceHeader(t *testing.T) { + // A workspace-scoped client stamps the routing header the Files API control + // plane needs; a client without a workspace ID sets nothing. + e := newEngine(buildUploadClient(t, "https://h.test", http.DefaultClient, "1234")) + req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, "https://h.test/api/2.0/fs/files/x", nil) + noErr(t, err) + e.setWorkspaceHeader(req) + if got := req.Header.Get("X-Databricks-Workspace-Id"); got != "1234" { + t.Errorf("X-Databricks-Workspace-Id = %q, want 1234", got) + } + + e2 := newEngine(buildUploadClient(t, "https://h.test", http.DefaultClient, "")) + req2, err := http.NewRequestWithContext(t.Context(), http.MethodPost, "https://h.test/api/2.0/fs/files/x", nil) + noErr(t, err) + e2.setWorkspaceHeader(req2) + if got := req2.Header.Get("X-Databricks-Workspace-Id"); got != "" { + t.Errorf("X-Databricks-Workspace-Id = %q, want empty", got) + } +} + +// failingSeeker reports a size (Seek to the end succeeds) but cannot rewind +// (Seek to the start fails), modeling the unrecoverable probe case. +type failingSeeker struct{ io.ReadSeeker } + +func (f failingSeeker) Seek(offset int64, whence int) (int64, error) { + if whence == io.SeekStart { + return 0, errors.New("seek start failed") + } + return f.ReadSeeker.Seek(offset, whence) +} + +func TestUploadRewindFailureIsLoud(t *testing.T) { + f := newFakeServer(t, "multipart") + e := shrunkEngine(t, f) + // A reader the engine can size but not rewind must fail the upload rather + // than send a truncated object from the EOF position. + r := failingSeeker{bytes.NewReader(make([]byte, 1024))} + if _, err := e.upload(t.Context(), "/Volumes/c/s/v/f.bin", r); err == nil { + t.Fatal("expected an error when the reader cannot be rewound after sizing") + } +} + +// testEngine returns an engine over a client pointed at f, with production +// tunables. +func testEngine(t *testing.T, f *fakeServer) *engine { + t.Helper() + return newEngine(newTestClient(t, f)) +} + +// shrunkEngine returns an engine whose size thresholds are tiny, so multipart +// boundaries are hit with small inputs. Tunables are injected per engine, so +// tests never mutate shared state and can run in parallel. +func shrunkEngine(t *testing.T, f *fakeServer) *engine { + t.Helper() + e := testEngine(t, f) + e.tun.minStreamSize = 1024 + e.tun.defaultPartSize = 1024 + return e +} + +// --- assertion helpers (standard library testing; no third-party framework) --- + +func noErr(t *testing.T, err error) { + t.Helper() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func eqBytes(t *testing.T, got, want []byte) { + t.Helper() + if !bytes.Equal(got, want) { + t.Errorf("uploaded bytes differ: got %d bytes, want %d bytes", len(got), len(want)) + } +} + +func data(n int) []byte { + b := make([]byte, n) + for i := range b { + b[i] = byte(i % 251) + } + return b +} + +func writeTemp(t *testing.T, b []byte) string { + t.Helper() + p := t.TempDir() + "/src.bin" + noErr(t, os.WriteFile(p, b, 0o600)) + return p +} + +func TestUploadSingleShot(t *testing.T) { + f := newFakeServer(t, "multipart") + e := shrunkEngine(t, f) + + in := data(500) // below the 1024 single-shot threshold + _, err := e.upload(t.Context(), "/Volumes/c/s/v/small.bin", bytes.NewReader(in)) + noErr(t, err) + eqBytes(t, f.assembled(), in) +} + +func TestUploadMultipartParallelMatchesSequential(t *testing.T) { + in := data(20 * 1024) + + run := func(parallelism int) []byte { + f := newFakeServer(t, "multipart") + e := shrunkEngine(t, f) + _, err := e.uploadFrom(t.Context(), "/Volumes/c/s/v/big.bin", writeTemp(t, in), + WithParallelism(parallelism)) + noErr(t, err) + // Parts must be completed in sorted order with non-empty ETags. + if len(f.completed) == 0 { + t.Fatal("expected completed parts") + } + for i, p := range f.completed { + if p.PartNumber != i+1 { + t.Errorf("part %d: got part number %d, want %d", i, p.PartNumber, i+1) + } + if p.ETag == "" { + t.Errorf("part %d: empty ETag", p.PartNumber) + } + } + return f.assembled() + } + + eqBytes(t, run(1), in) + eqBytes(t, run(8), in) +} + +func TestUploadFromReaderAtKnownSize(t *testing.T) { + f := newFakeServer(t, "multipart") + e := shrunkEngine(t, f) + + in := data(5000) + _, err := e.upload(t.Context(), "/Volumes/c/s/v/r.bin", bytes.NewReader(in), WithParallelism(4)) + noErr(t, err) + eqBytes(t, f.assembled(), in) +} + +func TestUploadNonSeekableStream(t *testing.T) { + f := newFakeServer(t, "multipart") + e := shrunkEngine(t, f) + + in := data(5000) + // io.Reader wrapper that is neither Seeker nor ReaderAt. + r := io.MultiReader(bytes.NewReader(in)) + _, err := e.upload(t.Context(), "/Volumes/c/s/v/stream.bin", r, WithParallelism(4)) + noErr(t, err) + eqBytes(t, f.assembled(), in) +} + +func TestCloudPartsCarryNoDatabricksCredentials(t *testing.T) { + f := newFakeServer(t, "multipart") + e := shrunkEngine(t, f) + + _, err := e.uploadFrom(t.Context(), "/Volumes/c/s/v/big.bin", writeTemp(t, data(5000)), WithParallelism(4)) + noErr(t, err) + + if len(f.partHeaders) == 0 { + t.Fatal("expected cloud part requests") + } + for _, h := range f.partHeaders { + if got := h.Get("Authorization"); got != "" { + t.Errorf("cloud part PUT must not carry Databricks auth, got Authorization=%q", got) + } + if got := h.Get("X-Databricks-Workspace-Id"); got != "" { + t.Errorf("cloud part PUT must not carry workspace routing, got X-Databricks-Workspace-Id=%q", got) + } + } +} + +func TestUploadOverwrite(t *testing.T) { + cases := []struct { + name string + opts []UploadOption + want string + }{ + {"unset", nil, ""}, + {"true", []UploadOption{WithOverwrite(true)}, "true"}, + {"false", []UploadOption{WithOverwrite(false)}, "false"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + f := newFakeServer(t, "multipart") + e := shrunkEngine(t, f) + _, err := e.upload(t.Context(), "/Volumes/c/s/v/x.bin", bytes.NewReader(data(100)), tc.opts...) + noErr(t, err) + if len(f.overwriteQ) == 0 { + t.Fatal("expected an overwrite query value") + } + if f.overwriteQ[0] != tc.want { + t.Errorf("overwrite query = %q, want %q", f.overwriteQ[0], tc.want) + } + }) + } +} + +func TestFileUploadMintsURLsInBatches(t *testing.T) { + f := newFakeServer(t, "multipart") + e := shrunkEngine(t, f) + + in := data(16 * 1024) // 16 parts at the 1024 part size + _, err := e.uploadFrom(t.Context(), "/Volumes/c/s/v/batch.bin", writeTemp(t, in), WithParallelism(8)) + noErr(t, err) + eqBytes(t, f.assembled(), in) + + // The file path pre-mints in batches, so at least one create-upload-part-urls + // call requests more than one URL; per-part minting would never exceed 1. + maxCount := 0 + for _, cnt := range f.partURLReqs { + maxCount = max(maxCount, cnt) + } + if maxCount <= 1 { + t.Errorf("expected batched URL minting (a request with count > 1), got max count %d across %d calls", maxCount, len(f.partURLReqs)) + } +} + +func TestCloudRejectsAllPartsFallsBackToSingleShot(t *testing.T) { + f := newFakeServer(t, "multipart") + // A cloud firewall forbids every part PUT on every attempt (a blanket block of + // the storage host). Parts upload concurrently, so no single part is the canary; + // because none lands, the upload falls back to single-shot. + f.partHook = func(n, attempt int) (int, string, []byte) { + return http.StatusForbidden, "", []byte("forbidden") + } + e := shrunkEngine(t, f) + + in := data(5000) + _, err := e.uploadFrom(t.Context(), "/Volumes/c/s/v/fb.bin", writeTemp(t, in), WithParallelism(4)) + noErr(t, err) + // The fallback single-shot PUT carries the whole object. + eqBytes(t, f.assembled(), in) +} + +func TestUploadAlreadyExists(t *testing.T) { + f := newFakeServer(t, "multipart") + f.singleShotStatus = http.StatusConflict + f.singleShotBody = `{"error_code":"ALREADY_EXISTS","message":"file already exists"}` + e := shrunkEngine(t, f) + + _, err := e.upload(t.Context(), "/Volumes/c/s/v/exists.bin", bytes.NewReader(data(100)), WithOverwrite(false)) + if !errors.Is(err, ErrAlreadyExists) { + t.Fatalf("got error %v, want ErrAlreadyExists", err) + } +} + +// TestWithProgressPreservesSeekable is the unit-level guard for the retry bug: +// wrapping a body for progress reporting must not hide its io.Seeker (which the +// single-shot path probes to decide retriability), yet must not fabricate one +// for a non-seekable source. +func TestWithProgressPreservesSeekable(t *testing.T) { + p := newProgressReporter(func(Progress) {}, 10) + + // A seekable source stays seekable through the wrapper, and Seek is forwarded. + seekable := withProgress(bytes.NewReader([]byte("0123456789")), p) + s, ok := seekable.(io.Seeker) + if !ok { + t.Fatal("withProgress dropped io.Seeker for a seekable source") + } + if _, err := s.Seek(0, io.SeekStart); err != nil { + t.Errorf("forwarded Seek returned error: %v", err) + } + + // A non-seekable source must not be reported as seekable. + if _, ok := withProgress(io.MultiReader(bytes.NewReader(nil)), p).(io.Seeker); ok { + t.Error("withProgress fabricated io.Seeker for a non-seekable source") + } +} + +// TestSingleShotRetriesWithProgress guards both MAJOR review findings: with +// WithProgress set, a seekable single-shot body must still be retried on a +// transient failure (the progress wrapper previously hid the io.Seeker, which +// silently disabled retry). +func TestSingleShotRetriesWithProgress(t *testing.T) { + f := newFakeServer(t, "multipart") + // First single-shot PUT fails with a retryable 503; the second succeeds. + f.singleShotHook = func(attempt int) int { + if attempt == 1 { + return http.StatusServiceUnavailable + } + return 0 + } + e := shrunkEngine(t, f) + // Shrink the control-plane backoff so the retry does not sleep a full second. + e.tun.cpBackoff = ops.BackoffPolicy{Initial: time.Millisecond, Maximum: time.Millisecond, Factor: 1} + + var reported atomic.Int64 + in := data(500) // below the 1024 single-shot threshold + _, err := e.upload(t.Context(), "/Volumes/c/s/v/p.bin", bytes.NewReader(in), + WithProgress(func(p Progress) { reported.Store(p.Transferred) })) + noErr(t, err) + eqBytes(t, f.assembled(), in) + + if f.singleTry != 2 { + t.Errorf("single-shot attempts = %d, want 2 (one retry after the 503)", f.singleTry) + } + if reported.Load() == 0 { + t.Error("progress callback never fired") + } +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +func TestUploadWithTransferClient(t *testing.T) { + f := newFakeServer(t, "multipart") + e := shrunkEngine(t, f) + + var used atomic.Int32 + custom := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) { + used.Add(1) + return http.DefaultTransport.RoundTrip(r) + })} + + in := data(5000) + _, err := e.uploadFrom(t.Context(), "/Volumes/c/s/v/tc.bin", writeTemp(t, in), + WithParallelism(2), WithTransferClient(custom)) + noErr(t, err) + eqBytes(t, f.assembled(), in) + if used.Load() == 0 { + t.Error("the cloud transfer must go through the supplied client") + } +} + +func TestUploadResumableHappyPath(t *testing.T) { + f := newFakeServer(t, "resumable") + e := shrunkEngine(t, f) + + in := data(5000) + // Parallelism > 1 still uses single-threaded resumable; bytes must match. + _, err := e.uploadFrom(t.Context(), "/Volumes/c/s/v/gcp.bin", writeTemp(t, in), WithParallelism(4)) + noErr(t, err) + eqBytes(t, f.assembled(), in) +} + +func TestUploadResumableNoProgressBounded(t *testing.T) { + f := newFakeServer(t, "resumable") + f.stuckResumable = true + e := shrunkEngine(t, f) + + // The server confirms no new bytes on every chunk; the resumable loop must + // error rather than spin forever (fs cp sets no context deadline). + _, err := e.upload(t.Context(), "/Volumes/c/s/v/gcp.bin", bytes.NewReader(data(5000))) + if err == nil { + t.Fatal("expected an error when the resumable server makes no progress") + } +} + +// TestUploadResumableTransientRetryBudgetIsPerChunk guards against the retry +// budget being a global cap rather than per-progress: every chunk PUT returns a +// transient 503 (but commits its bytes, so the status query shows forward +// progress). With more chunks than maxRetries, a global budget would fail the +// upload after maxRetries total blips even though it keeps advancing; the budget +// must reset on each confirmed advance so a steadily-progressing flaky upload +// completes. +func TestUploadResumableTransientRetryBudgetIsPerChunk(t *testing.T) { + f := newFakeServer(t, "resumable") + e := shrunkEngine(t, f) // 1024-byte parts + // Every data-chunk PUT reports a retryable 503 after committing; with a + // 5000-byte input at 1024-byte chunks that is ~5 chunks, well over maxRetries. + f.resumableChunkHook = func(attempt int) int { return http.StatusServiceUnavailable } + + in := data(5000) + _, err := e.upload(t.Context(), "/Volumes/c/s/v/flaky.bin", bytes.NewReader(in)) + noErr(t, err) + eqBytes(t, f.assembled(), in) + if f.resumableTry <= e.tun.maxRetries { + t.Errorf("test did not exceed the retry budget: %d chunk PUTs, want > maxRetries (%d)", f.resumableTry, e.tun.maxRetries) + } +} + +func TestUploadPresignedURLExpiryRetried(t *testing.T) { + f := newFakeServer(t, "multipart") + // Every part's first presigned URL is reported expired; the second succeeds. + azure := `AuthenticationFailedSignature not valid in the specified time frame` + f.partHook = func(n, attempt int) (int, string, []byte) { + if attempt == 1 { + return http.StatusForbidden, "", []byte(azure) + } + return 0, "", nil + } + e := shrunkEngine(t, f) + + in := data(5000) + _, err := e.uploadFrom(t.Context(), "/Volumes/c/s/v/exp.bin", writeTemp(t, in), WithParallelism(4)) + noErr(t, err) + eqBytes(t, f.assembled(), in) +} + +// Retry-timing coverage (retryable-status retried/exhausted, with shrunk backoff) +// lives in the cloudstorage package, where Send's retry policy is defined. + +func TestResolvedParallelism(t *testing.T) { + n := func(v int) *int { return &v } + const mib int64 = 1 << 20 + tun := defaultTunables() + tests := []struct { + name string + cfg uploadConfig + buffered bool + partSize int64 + want int + }{ + {"explicit override wins for a file", uploadConfig{parallelism: n(7)}, false, 16 * mib, 7}, + {"explicit override wins for a stream", uploadConfig{parallelism: n(3)}, true, 16 * mib, 3}, + {"file default is bandwidth-oriented", uploadConfig{}, false, 16 * mib, tun.fileParallelism}, + {"stream default fits the memory budget", uploadConfig{}, true, 16 * mib, int(tun.streamMemoryBudget / (16 * mib))}, + {"larger stream part size lowers the worker count", uploadConfig{}, true, 64 * mib, int(tun.streamMemoryBudget / (64 * mib))}, + {"tiny stream part size is capped at the file default", uploadConfig{}, true, 1, tun.fileParallelism}, + {"huge stream part size floors at one worker", uploadConfig{}, true, tun.streamMemoryBudget * 2, 1}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := tc.cfg.resolvedParallelism(tc.buffered, tc.partSize, tun); got != tc.want { + t.Errorf("resolvedParallelism(%v, %d) = %d, want %d", tc.buffered, tc.partSize, got, tc.want) + } + }) + } +} diff --git a/libs/tmp/files/v2/genhelper.go b/libs/tmp/files/v2/genhelper.go new file mode 100644 index 00000000000..c757bcd6c05 --- /dev/null +++ b/libs/tmp/files/v2/genhelper.go @@ -0,0 +1,250 @@ +// Code generated by Databricks SDK Generator. DO NOT EDIT. + +package files + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "net/url" + "strings" + + "github.com/databricks/cli/libs/tmp/auth" + "github.com/databricks/cli/libs/tmp/options/call" + "github.com/databricks/cli/libs/tmp/options/internaloptions" + "github.com/databricks/sdk-go/core/apierr" + "github.com/databricks/sdk-go/core/ops" +) + +// unmarshalResponse decodes a JSON response body into v, treating an empty body +// as a zero-value response. Several Files API endpoints (HEAD metadata calls and +// 204 PUT/DELETE responses) return no body, which json.Unmarshal rejects with +// "unexpected end of JSON input". +// +// NOTE: local patch to the vendored files/v2 copy; the upstream generator emits +// a bare json.Unmarshal that crashes on these empty bodies. See libs/tmp/README. +func unmarshalResponse(respBody []byte, v any) error { + if len(respBody) == 0 { + return nil + } + return json.Unmarshal(respBody, v) +} + +type httpCallOptions struct { + req *http.Request + client *http.Client + logger *slog.Logger +} + +type httpRequestOptions struct { + Method string + URL string + Credentials auth.Credentials + Headers http.Header + Body io.Reader +} + +func newHTTPRequest(ctx context.Context, opts httpRequestOptions) (*http.Request, error) { + req, err := http.NewRequestWithContext(ctx, opts.Method, opts.URL, opts.Body) + if err != nil { + closeBody(opts.Body) + return nil, err + } + + req.Header = opts.Headers.Clone() + if opts.Credentials != nil { + headers, err := opts.Credentials.AuthHeaders(ctx) + if err != nil { + closeBody(opts.Body) + return nil, err + } + for _, h := range headers { + req.Header.Add(h.Key, h.Value) + } + } + + return req, nil +} + +func closeBody(body io.Reader) { + if body == nil { + return + } + if rc, ok := body.(io.ReadCloser); ok { + _ = rc.Close() + } +} + +// lazyRequest implements slog.LogValuer for lazy evaluation of request +// logging. +type lazyRequest struct{ req *http.Request } + +func (r lazyRequest) LogValue() slog.Value { + return slog.GroupValue( + slog.String("method", r.req.Method), + slog.String("url", r.req.URL.String()), + ) +} + +// lazyResponse implements slog.LogValuer for lazy evaluation of response +// logging. +type lazyResponse struct { + resp *http.Response + body []byte +} + +func (r lazyResponse) LogValue() slog.Value { + return slog.GroupValue( + slog.Int("status", r.resp.StatusCode), + slog.String("body", string(r.body)), + ) +} + +// executeHTTPCall executes an HTTP call and returns the response body and +// headers, or an error. The headers are returned so callers can populate +// header-mapped response fields alongside the JSON body. This function takes +// care of logging the request and response. +func executeHTTPCall(opts httpCallOptions) ([]byte, http.Header, error) { + opts.logger.Debug("HTTP request", "request", lazyRequest{opts.req}) + resp, err := opts.client.Do(opts.req) + if err != nil { + return nil, nil, err + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + opts.logger.Debug("HTTP response", "response", lazyResponse{resp, nil}) + return nil, nil, err + } + opts.logger.Debug("HTTP response", "response", lazyResponse{resp, body}) + if err := apierr.FromHTTPError(resp.StatusCode, resp.Header, body); err != nil { + return nil, nil, err + } + return body, resp.Header, nil +} + +// executeCall resolves call.Option values to ops.Option values and invokes +// ops.Execute. +func executeCall(ctx context.Context, op func(context.Context) error, opts []call.Option) error { + cfg := internaloptions.CallOptions{} + for _, opt := range opts { + if err := opt(&cfg); err != nil { + return err + } + } + var opsOpts []ops.Option + if cfg.Retrier != nil { + opsOpts = append(opsOpts, ops.WithRetrier(cfg.Retrier)) + } + if cfg.RateLimiter != nil { + opsOpts = append(opsOpts, ops.WithLimiter(cfg.RateLimiter)) + } + if cfg.Timeout != 0 { + opsOpts = append(opsOpts, ops.WithTimeout(cfg.Timeout)) + } + return ops.Execute(ctx, op, opsOpts...) +} +func addQueryValue(params url.Values, key string, value any) error { + data, err := json.Marshal(value) + if err != nil { + return err + } + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.UseNumber() + var decoded any + if err := decoder.Decode(&decoded); err != nil { + return err + } + flattenQueryValue(params, key, decoded) + return nil +} + +func flattenQueryValue(params url.Values, key string, value any) { + switch value := value.(type) { + case nil: + // JSON null represents an absent query value. + case map[string]any: + for childKey, child := range value { + flattenQueryValue(params, key+"."+childKey, child) + } + case []any: + for _, item := range value { + params.Add(key, fmt.Sprintf("%v", item)) + } + default: + params.Add(key, fmt.Sprintf("%v", value)) + } +} + +// pathBuilder assembles a request path from static literals and parameter +// values. It tracks the decoded path and its percent-escaped wire form in +// lockstep so build() can assign both url.URL.Path and url.URL.RawPath; because +// RawPath is a valid escaping of Path, url.URL.String() emits it verbatim +// instead of re-escaping (which would double-encode "%"). +type pathBuilder struct { + path strings.Builder + raw strings.Builder +} + +// literal appends a static path segment, identical on the decoded and escaped +// paths. +func (b *pathBuilder) literal(s string) { + b.path.WriteString(s) + b.raw.WriteString(s) +} + +// singleSegment appends a single-segment path parameter: the value occupies one +// path segment, so everything is escaped, including "/". The value is formatted +// with %v so strings, enums, and numbers all work. +func (b *pathBuilder) singleSegment(v any) { + s := fmt.Sprintf("%v", v) + b.path.WriteString(s) + b.raw.WriteString(url.PathEscape(s)) +} + +// multiSegments appends a multi-segment path parameter: the value spans several +// path segments, so each segment is escaped but the "/" separators are kept. +// The value is formatted with %v so strings, enums, and numbers all work. +func (b *pathBuilder) multiSegments(v any) { + s := fmt.Sprintf("%v", v) + b.path.WriteString(s) + segments := strings.Split(s, "/") + for i, seg := range segments { + segments[i] = url.PathEscape(seg) + } + b.raw.WriteString(strings.Join(segments, "/")) +} + +func (b *pathBuilder) build() (path, rawPath string) { + return b.path.String(), b.raw.String() +} + +// executeStreamingHTTPCall executes an HTTP call whose response body is a raw +// byte stream. On success it returns the response with its body still open, so +// the caller can stream and close it. On a non-2xx status it reads and closes +// the body to build the API error; the returned response is nil in that case. +func executeStreamingHTTPCall(opts httpCallOptions) (*http.Response, error) { + opts.logger.Debug("HTTP request", "request", lazyRequest{opts.req}) + resp, err := opts.client.Do(opts.req) + if err != nil { + return nil, err + } + if resp.StatusCode >= 400 { + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + opts.logger.Debug("HTTP response", "response", lazyResponse{resp, nil}) + return nil, err + } + opts.logger.Debug("HTTP response", "response", lazyResponse{resp, body}) + if apiErr := apierr.FromHTTPError(resp.StatusCode, resp.Header, body); apiErr != nil { + return nil, apiErr + } + } + opts.logger.Debug("HTTP response", "response", lazyResponse{resp, []byte("")}) + return resp, nil +} diff --git a/libs/tmp/files/v2/internal/cloudstorage/body.go b/libs/tmp/files/v2/internal/cloudstorage/body.go new file mode 100644 index 00000000000..828bdf638df --- /dev/null +++ b/libs/tmp/files/v2/internal/cloudstorage/body.go @@ -0,0 +1,55 @@ +package cloudstorage + +import ( + "bytes" + "io" +) + +// Body supplies a request's bytes and can re-supply them for each retry attempt, +// so a request is retried without the caller buffering it again. +type Body interface { + // Reader returns a reader over the body positioned at the start. It is called + // once per attempt. + Reader() (io.Reader, error) + // Size reports the body length in bytes, for Content-Length and the + // short-read guard. + Size() int64 +} + +// BytesBody returns a Body backed by an in-memory buffer. +func BytesBody(data []byte) Body { + return bytesBody{data} +} + +type bytesBody struct { + data []byte +} + +func (b bytesBody) Reader() (io.Reader, error) { + return bytes.NewReader(b.data), nil +} + +func (b bytesBody) Size() int64 { + return int64(len(b.data)) +} + +// SectionBody returns a Body that reads length bytes starting at offset from ra +// on each attempt. ra may be shared across concurrent bodies: positioned reads +// via io.ReaderAt are safe to run in parallel. +func SectionBody(ra io.ReaderAt, offset, length int64) Body { + return sectionBody{ra: ra, offset: offset, length: length} +} + +type sectionBody struct { + ra io.ReaderAt + offset int64 + length int64 +} + +func (s sectionBody) Reader() (io.Reader, error) { + return io.NewSectionReader(s.ra, s.offset, s.length), nil +} + +func (s sectionBody) Size() int64 { + return s.length +} diff --git a/libs/tmp/files/v2/internal/cloudstorage/body_test.go b/libs/tmp/files/v2/internal/cloudstorage/body_test.go new file mode 100644 index 00000000000..f67d052a5aa --- /dev/null +++ b/libs/tmp/files/v2/internal/cloudstorage/body_test.go @@ -0,0 +1,57 @@ +package cloudstorage + +import ( + "bytes" + "io" + "strings" + "testing" +) + +func TestBytesBody(t *testing.T) { + data := []byte("hello world") + b := BytesBody(data) + + if b.Size() != int64(len(data)) { + t.Errorf("Size = %d, want %d", b.Size(), len(data)) + } + + // Reader is called once per attempt and must re-supply the full body from the + // start each time, so a retry needs no rewind. + for i := range 2 { + r, err := b.Reader() + if err != nil { + t.Fatalf("attempt %d: Reader returned error: %v", i, err) + } + got, err := io.ReadAll(r) + if err != nil { + t.Fatalf("attempt %d: ReadAll returned error: %v", i, err) + } + if !bytes.Equal(got, data) { + t.Errorf("attempt %d: read %q, want %q", i, got, data) + } + } +} + +func TestSectionBody(t *testing.T) { + src := strings.NewReader("0123456789") + b := SectionBody(src, 3, 4) + + if b.Size() != 4 { + t.Errorf("Size = %d, want 4", b.Size()) + } + + // Each Reader reads only the [offset, offset+length) section, from the start. + for i := range 2 { + r, err := b.Reader() + if err != nil { + t.Fatalf("attempt %d: Reader returned error: %v", i, err) + } + got, err := io.ReadAll(r) + if err != nil { + t.Fatalf("attempt %d: ReadAll returned error: %v", i, err) + } + if string(got) != "3456" { + t.Errorf("attempt %d: read %q, want %q", i, got, "3456") + } + } +} diff --git a/libs/tmp/files/v2/internal/cloudstorage/cloudstorage.go b/libs/tmp/files/v2/internal/cloudstorage/cloudstorage.go new file mode 100644 index 00000000000..5574f54ac7c --- /dev/null +++ b/libs/tmp/files/v2/internal/cloudstorage/cloudstorage.go @@ -0,0 +1,320 @@ +// Package cloudstorage issues unauthenticated, idempotent HTTP requests to cloud +// object storage (S3, Azure Blob, GCS) using short-lived presigned URLs. The +// URLs are self-authenticating, so the client deliberately carries no Databricks +// credentials. It owns the inactivity-stall and exact-length guards and a capped +// retry-with-backoff policy; callers supply the bytes via a [Body] and interpret +// the returned [Response]. +package cloudstorage + +import ( + "context" + "encoding/xml" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptrace" + "strings" + "time" + + "github.com/databricks/sdk-go/core/apierr" + "github.com/databricks/sdk-go/core/ops" +) + +// defaultIdleTimeout cuts a transfer that stops making progress: if the request +// body is not read for this long, no bytes are moving and the attempt is cancelled. +var defaultIdleTimeout = 60 * time.Second + +var ( + errStalled = errors.New("transfer stalled (no progress within the idle timeout)") + errShortRead = errors.New("body produced fewer bytes than expected (was the source modified mid-transfer?)") +) + +// Response is the captured result of a presigned request. Unlike a Databricks +// API call, cloud requests expose the raw status and headers (ETag, Range, 308) +// and are never mapped to a Databricks API error. +type Response struct { + StatusCode int + Header http.Header + Body []byte +} + +// Client issues requests to cloud object storage over presigned URLs. It holds +// only an *http.Client: no Databricks credentials, no config. +type Client struct { + http *http.Client + retrier func() ops.Retrier + idleTimeout time.Duration + logger *slog.Logger +} + +// Option configures a [Client]. +type Option func(*Client) + +// WithLogger sets the logger used for debug-level transfer diagnostics. When +// unset, transfer diagnostics are discarded. +func WithLogger(logger *slog.Logger) Option { + return func(c *Client) { + if logger != nil { + c.logger = logger + } + } +} + +// New returns a Client that sends over the given HTTP client. The caller owns +// the client's transport, timeouts, and connection-pool sizing; it must not +// attach Databricks credentials, and should not set a whole-request +// http.Client.Timeout (it would abort a legitimately long transfer; bound the +// operation with the request context instead). +func New(httpClient *http.Client, opts ...Option) *Client { + c := &Client{ + http: httpClient, + retrier: newRetrier, + idleTimeout: defaultIdleTimeout, + logger: slog.New(slog.DiscardHandler), + } + for _, opt := range opts { + opt(c) + } + return c +} + +// Attempt performs a single request and returns the raw response. body may be +// nil for a bodyless request (a status query or an abort). +func (c *Client) Attempt(ctx context.Context, method, url string, headers map[string]string, body Body) (*Response, error) { + var rdr io.Reader + var bodyLen int64 + if body != nil { + var err error + if rdr, err = body.Reader(); err != nil { + return nil, err + } + bodyLen = body.Size() + } + + // Cancel the attempt on two conditions, each with a distinct cause: the + // transfer stalls (a dead connection that accepts no more bytes but never + // errors), or the source comes up short (truncated mid-transfer). Both are + // told apart from the caller cancelling ctx; the stall is retried, the short + // read is not. + attemptCtx, cancel := context.WithCancelCause(ctx) + defer cancel(nil) + + req, err := http.NewRequestWithContext(attemptCtx, method, url, rdr) + if err != nil { + return nil, err + } + for k, v := range headers { + req.Header.Set(k, v) + } + if bodyLen > 0 && req.Body != nil { + // A non-bytes source (a file section) is not length-detected by NewRequest, + // so set Content-Length explicitly. Enforce the exact length and guard + // against a stall; the body is read as bytes move to the socket, so a gap + // longer than IdleTimeout means the transfer has stalled. The response phase + // is bounded by the transport's ResponseHeaderTimeout. + req.ContentLength = bodyLen + req.Body = stallGuardedBody(exactBody(req.Body, bodyLen, cancel), c.idleTimeout, cancel) + } + + // Trace connection reuse so a slow request can be attributed to connection + // behavior (e.g. HTTP/2 multiplexing onto one connection vs a fresh one). + var connReused, connWasIdle bool + req = req.WithContext(httptrace.WithClientTrace(req.Context(), &httptrace.ClientTrace{ + GotConn: func(info httptrace.GotConnInfo) { + connReused, connWasIdle = info.Reused, info.WasIdle + }, + })) + + start := time.Now() + resp, err := c.http.Do(req) + if err != nil { + // Surface our own cancellations as distinct errors; a caller ctx + // cancellation stays as-is. errShortRead carries its detail message. + if ctx.Err() == nil { + switch cause := context.Cause(attemptCtx); { + case errors.Is(cause, errShortRead): + err = cause + case errors.Is(cause, errStalled): + err = errStalled + } + } + c.logger.DebugContext(ctx, "request failed", + "method", method, "host", req.URL.Host, "req_bytes", bodyLen, + "duration_ms", time.Since(start).Milliseconds(), "error", err) + return nil, err + } + defer resp.Body.Close() + respBody, err := io.ReadAll(resp.Body) + duration := time.Since(start) + if err != nil { + return nil, err + } + c.logger.DebugContext(ctx, "response", + "method", method, "host", req.URL.Host, "status", resp.StatusCode, "proto", resp.Proto, + "conn_reused", connReused, "conn_was_idle", connWasIdle, "req_bytes", bodyLen, + "duration_ms", duration.Milliseconds()) + return &Response{StatusCode: resp.StatusCode, Header: resp.Header, Body: respBody}, nil +} + +// Send performs Attempt with capped retries: transient network errors, the +// retryable statuses in retryStatusCodes (honoring Retry-After), and stalled +// attempts are retried with exponential backoff. A fresh request and body reader +// are built per attempt, so no explicit rewind is needed. Exhausting retries on +// a retryable status returns that status as an error. +func (c *Client) Send(ctx context.Context, method, url string, headers map[string]string, body Body) (*Response, error) { + var result *Response + op := func(ctx context.Context) error { + resp, err := c.Attempt(ctx, method, url, headers, body) + if err != nil { + return err // network error or inactivity cut; the retrier decides + } + // Surface a failure status as an APIError and let the retrier decide: it + // retries the retryable statuses (honoring Retry-After) and transient + // transport failures, and returns everything else. A non-retryable status + // (403 URL expired, 412 already exists, ...) therefore reaches the caller as + // this error, which carries the status, headers, and body for inspection. + // FromHTTPError returns nil for 2xx and 3xx, so both pass through as a + // successful response (the http.Client already follows real redirects). + if apiErr := apierr.FromHTTPError(resp.StatusCode, resp.Header, resp.Body); apiErr != nil { + return apiErr + } + result = resp + return nil + } + if err := ops.Execute(ctx, op, ops.WithRetrier(c.retrier)); err != nil { + return nil, err + } + return result, nil +} + +// IsURLExpired reports whether err (as returned by Send) is one of the known +// cloud-provider "presigned URL expired" errors: a 403 with the AWS or Azure +// signature-expiry body. +func IsURLExpired(err error) bool { + aerr, ok := errors.AsType[*apierr.APIError](err) + if !ok || aerr.HTTPStatusCode() != http.StatusForbidden { + return false + } + var e struct { + XMLName xml.Name `xml:"Error"` + Code string `xml:"Code"` + Message string `xml:"Message"` + AuthDetail string `xml:"AuthenticationErrorDetail"` + } + if uerr := xml.Unmarshal(aerr.HTTPBody(), &e); uerr != nil { + return false + } + switch e.Code { + case "AuthenticationFailed": // Azure + return strings.Contains(e.AuthDetail, "Signature not valid in the specified time frame") + case "AccessDenied": // AWS + return e.Message == "Request has expired" + default: + return false + } +} + +// stallReader wraps a request body so a transfer that stops making progress is +// cancelled. The transport reads the body as it writes bytes to the socket, so a +// gap longer than idle between reads means no bytes are moving -- a stalled or +// dead connection. A transfer making any progress keeps resetting the timer and +// so cannot be cut, unlike a whole-attempt or throughput-floor deadline. +type stallReader struct { + r io.Reader + idle time.Duration + cancel context.CancelCauseFunc + timer *time.Timer +} + +func newStallReader(r io.Reader, idle time.Duration, cancel context.CancelCauseFunc) *stallReader { + return &stallReader{r: r, idle: idle, cancel: cancel} +} + +func (s *stallReader) Read(p []byte) (int, error) { + // Arm on the first read (after the connection is established, so dial time is + // not charged against the idle budget) and reset on each later read. + if s.timer == nil { + s.timer = time.AfterFunc(s.idle, func() { s.cancel(errStalled) }) + } else { + s.timer.Reset(s.idle) + } + n, err := s.r.Read(p) + if err != nil { + // EOF (transfer done) or a read error: stop arming so the response phase is + // not subject to the stall timer. + s.stop() + } + return n, err +} + +// stop disarms the idle timer. Safe to call before the first Read (timer nil) +// and more than once. +func (s *stallReader) stop() { + if s.timer != nil { + s.timer.Stop() + } +} + +// stallGuardedBody wraps a request body's reads with stall detection while +// preserving its Close. +func stallGuardedBody(body io.ReadCloser, idle time.Duration, cancel context.CancelCauseFunc) io.ReadCloser { + return &stallGuardedReadCloser{Closer: body, sr: newStallReader(body, idle, cancel)} +} + +type stallGuardedReadCloser struct { + io.Closer + sr *stallReader +} + +func (b *stallGuardedReadCloser) Read(p []byte) (int, error) { + return b.sr.Read(p) +} + +// Close disarms the stall timer. The transport closes the request body once it +// has finished writing it, which on a fixed-length request happens without a +// trailing EOF read (the writer stops at Content-Length). Without stopping here, +// the timer would stay armed into the response phase and could cancel a +// successful upload as errStalled while waiting on a slow response. +func (b *stallGuardedReadCloser) Close() error { + b.sr.stop() + return b.Closer.Close() +} + +// exactReader enforces that a body delivers exactly the declared length. If the +// source ends early -- e.g. the file was truncated during the transfer -- it +// cancels the attempt with errShortRead rather than letting a short body be sent +// (silent corruption) or surfacing net/http's opaque ContentLength-mismatch +// error (which would then be pointlessly retried). +type exactReader struct { + r io.Reader + left int64 + cancel context.CancelCauseFunc +} + +func (e *exactReader) Read(p []byte) (int, error) { + n, err := e.r.Read(p) + e.left -= int64(n) + if err == io.EOF && e.left > 0 { + cause := fmt.Errorf("%w: stream ended %d bytes short of the expected length", errShortRead, e.left) + e.cancel(cause) + return n, cause + } + return n, err +} + +// exactBody wraps a request body's reads with the exact-length check while +// preserving its Close. +func exactBody(body io.ReadCloser, n int64, cancel context.CancelCauseFunc) io.ReadCloser { + return &exactReadCloser{Closer: body, er: &exactReader{r: body, left: n, cancel: cancel}} +} + +type exactReadCloser struct { + io.Closer + er *exactReader +} + +func (b *exactReadCloser) Read(p []byte) (int, error) { + return b.er.Read(p) +} diff --git a/libs/tmp/files/v2/internal/cloudstorage/cloudstorage_test.go b/libs/tmp/files/v2/internal/cloudstorage/cloudstorage_test.go new file mode 100644 index 00000000000..d92ae2c22a4 --- /dev/null +++ b/libs/tmp/files/v2/internal/cloudstorage/cloudstorage_test.go @@ -0,0 +1,174 @@ +package cloudstorage + +import ( + "bytes" + "errors" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/databricks/sdk-go/core/apierr" + "github.com/databricks/sdk-go/core/ops" +) + +// fastClient returns a Client whose retrier backs off in milliseconds, so retry +// tests do not sleep. +func fastClient(srv *httptest.Server) *Client { + fast := ops.BackoffPolicy{Initial: time.Millisecond, Maximum: time.Millisecond, Factor: 1} + c := New(srv.Client()) + c.retrier = func() ops.Retrier { + return &retrier{backoff: fast, maxRetries: 3} + } + return c +} + +func TestSendRetriesRetryableStatus(t *testing.T) { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) == 1 { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + resp, err := fastClient(srv).Send(t.Context(), http.MethodPut, srv.URL, nil, BytesBody([]byte("x"))) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } + if calls.Load() != 2 { + t.Errorf("calls = %d, want 2 (one retry)", calls.Load()) + } +} + +func TestSendExhaustsRetryableStatus(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + })) + defer srv.Close() + + _, err := fastClient(srv).Send(t.Context(), http.MethodPut, srv.URL, nil, BytesBody([]byte("x"))) + aerr, ok := errors.AsType[*apierr.APIError](err) + if !ok || aerr.HTTPStatusCode() != http.StatusServiceUnavailable { + t.Fatalf("err = %v, want a 503 APIError after exhausting retries", err) + } +} + +func TestSendReturnsNonRetryableStatusAsError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte("nope")) + })) + defer srv.Close() + + resp, err := New(srv.Client()).Send(t.Context(), http.MethodPut, srv.URL, nil, BytesBody([]byte("x"))) + if resp != nil { + t.Errorf("expected nil response on error, got %v", resp) + } + aerr, ok := errors.AsType[*apierr.APIError](err) + if !ok || aerr.HTTPStatusCode() != http.StatusForbidden { + t.Fatalf("err = %v, want a 403 APIError", err) + } + if string(aerr.HTTPBody()) != "nope" { + t.Errorf("HTTPBody = %q, want %q", aerr.HTTPBody(), "nope") + } +} + +func TestAttemptReturnsRawResponse(t *testing.T) { + // Attempt never maps a status to an error: a 308 (resumable continue) and its + // Range header come back as a Response. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Range", "bytes=0-41") + w.WriteHeader(http.StatusPermanentRedirect) + })) + defer srv.Close() + + resp, err := New(srv.Client()).Attempt(t.Context(), http.MethodPut, srv.URL, nil, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.StatusCode != http.StatusPermanentRedirect { + t.Errorf("status = %d, want 308", resp.StatusCode) + } + if got := resp.Header.Get("Range"); got != "bytes=0-41" { + t.Errorf("Range = %q, want bytes=0-41", got) + } +} + +func TestIsURLExpired(t *testing.T) { + azure := `AuthenticationFailedSignature not valid in the specified time frame` + aws := `AccessDeniedRequest has expired` + other := `AccessDeniedAccess Denied` + + cases := []struct { + name string + err error + want bool + }{ + {"azure expired", apierr.FromHTTPError(http.StatusForbidden, nil, []byte(azure)), true}, + {"aws expired", apierr.FromHTTPError(http.StatusForbidden, nil, []byte(aws)), true}, + {"other 403", apierr.FromHTTPError(http.StatusForbidden, nil, []byte(other)), false}, + {"404", apierr.FromHTTPError(http.StatusNotFound, nil, []byte(azure)), false}, + {"nil", nil, false}, + {"non-apierr", errors.New("boom"), false}, + } + for _, tc := range cases { + if got := IsURLExpired(tc.err); got != tc.want { + t.Errorf("%s: IsURLExpired = %v, want %v", tc.name, got, tc.want) + } + } +} + +// TestStallGuardStopsTimerOnClose verifies the idle timer is disarmed when the +// transport closes the request body, not only on a terminal Read. On a +// fixed-length request the transport normally issues a trailing EOF read that +// stops the timer, but Close is the reliable signal that body transmission is +// done: if the timer stayed armed into the response phase, a slow response could +// be cancelled as errStalled even though the upload succeeded. After Close, the +// cancel cause must never be set. +func TestStallGuardStopsTimerOnClose(t *testing.T) { + var cancelled atomic.Bool + cancel := func(error) { cancelled.Store(true) } + + // A tiny idle window so a leaked timer would fire quickly. + const idle = 20 * time.Millisecond + payload := []byte("payload") + body := stallGuardedBody(io.NopCloser(bytes.NewReader(payload)), idle, cancel) + + // Simulate a transport that stops reading at Content-Length without a trailing + // EOF read: read exactly len(payload) bytes (each Read returns n>0, nil, which + // arms the timer), then Close. Only Close can disarm the timer here. + buf := make([]byte, len(payload)) + if _, err := io.ReadFull(body, buf); err != nil { + t.Fatalf("ReadFull: %v", err) + } + if err := body.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + // Wait past the idle window; a timer left armed by Close would fire here. + time.Sleep(3 * idle) + if cancelled.Load() { + t.Fatal("stall timer fired after the body was fully read and closed") + } +} + +func TestExactReaderShortRead(t *testing.T) { + var caught error + er := &exactReader{r: bytes.NewReader(make([]byte, 10)), left: 100, cancel: func(c error) { caught = c }} + + _, err := io.ReadAll(er) + if !errors.Is(err, errShortRead) { + t.Fatalf("read error = %v, want errShortRead", err) + } + if !errors.Is(caught, errShortRead) { + t.Errorf("attempt cancelled with %v, want errShortRead", caught) + } +} diff --git a/libs/tmp/files/v2/internal/cloudstorage/retries.go b/libs/tmp/files/v2/internal/cloudstorage/retries.go new file mode 100644 index 00000000000..1f6ce8c2c89 --- /dev/null +++ b/libs/tmp/files/v2/internal/cloudstorage/retries.go @@ -0,0 +1,76 @@ +package cloudstorage + +import ( + "errors" + "net/http" + "slices" + "time" + + "github.com/databricks/sdk-go/core/apierr" + "github.com/databricks/sdk-go/core/apiretry" + "github.com/databricks/sdk-go/core/ops" +) + +// maxRetries caps how many times a single request is retried before giving up. +var maxRetries = 4 + +// retryStatusCodes are the HTTP statuses retried for an idempotent request. +var retryStatusCodes = []int{ + http.StatusRequestTimeout, // 408 + http.StatusTooManyRequests, // 429 + http.StatusBadGateway, // 502 + http.StatusServiceUnavailable, // 503 + http.StatusGatewayTimeout, // 504 +} + +// IsRetryableStatus reports whether an HTTP status code should be retried. It is +// exported for callers (such as a resumable upload) that run their own +// resume-aware retry loop over Attempt rather than using Send. +func IsRetryableStatus(code int) bool { + return slices.Contains(retryStatusCodes, code) +} + +func newRetrier() ops.Retrier { + return &retrier{maxRetries: maxRetries} +} + +type retrier struct { + backoff ops.BackoffPolicy + maxRetries int + attempts int +} + +func (r *retrier) IsRetriable(err error) (time.Duration, bool) { + r.attempts++ + if r.attempts > r.maxRetries { + return 0, false + } + if !IsRetriable(err) { + return 0, false + } + // Honor a Retry-After hint (carried on a throttling response's APIError) + // when it exceeds the next backoff delay. + return max(r.backoff.Delay(), apiretry.RetryDurationHint(err)), true +} + +// IsRetriable reports whether an error from Attempt is worth retrying: a stalled +// attempt is, a short read (deterministic) is not, a retryable HTTP status (as +// an APIError) is, and any other failure is retried only if it is a transient +// network error (which excludes context cancellation and deadline). It is +// exported for callers (such as a resumable upload) that run their own +// resume-aware retry loop over Attempt. +func IsRetriable(err error) bool { + if err == nil { + return false + } + if errors.Is(err, errShortRead) { + return false + } + if errors.Is(err, errStalled) { + return true + } + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { + return IsRetryableStatus(apiErr.HTTPStatusCode()) + } + return apiretry.IsTransientNetworkError(err) +} diff --git a/libs/tmp/files/v2/internal/cloudstorage/retries_test.go b/libs/tmp/files/v2/internal/cloudstorage/retries_test.go new file mode 100644 index 00000000000..51f403bf799 --- /dev/null +++ b/libs/tmp/files/v2/internal/cloudstorage/retries_test.go @@ -0,0 +1,43 @@ +package cloudstorage + +import ( + "context" + "errors" + "net/http" + "syscall" + "testing" + + "github.com/databricks/sdk-go/core/apierr" +) + +func TestIsRetriable(t *testing.T) { + cases := []struct { + name string + err error + want bool + }{ + {"nil", nil, false}, + {"stall", errStalled, true}, + {"short read", errShortRead, false}, + {"canceled", context.Canceled, false}, + {"deadline", context.DeadlineExceeded, false}, + {"transient network", syscall.ECONNRESET, true}, + {"retryable status", apierr.FromHTTPError(http.StatusServiceUnavailable, nil, nil), true}, + {"non-retryable status", apierr.FromHTTPError(http.StatusForbidden, nil, nil), false}, + {"generic", errors.New("connection reset"), false}, + } + for _, tc := range cases { + if got := IsRetriable(tc.err); got != tc.want { + t.Errorf("%s: IsRetriable = %v, want %v", tc.name, got, tc.want) + } + } +} + +func TestIsRetryableStatus(t *testing.T) { + if !IsRetryableStatus(http.StatusServiceUnavailable) { + t.Error("503 should be retryable") + } + if IsRetryableStatus(http.StatusNotFound) { + t.Error("404 should not be retryable") + } +} diff --git a/libs/tmp/files/v2/model.go b/libs/tmp/files/v2/model.go new file mode 100644 index 00000000000..4875cde1e4e --- /dev/null +++ b/libs/tmp/files/v2/model.go @@ -0,0 +1,190 @@ +// Code generated by Databricks SDK Generator. DO NOT EDIT. + +package files + +import ( + "io" +) + +type AddBlockRequest struct { + Handle *int64 + Data []byte +} + +type AddBlockResponse struct { +} + +type CloseRequest struct { + Handle *int64 +} + +type CloseResponse struct { +} + +// Create a directory. +type CreateDirectoryRequest struct { + DirectoryPath *string +} + +type CreateDirectoryResponse struct { +} + +type CreateRequest struct { + Path *string + Overwrite *bool +} + +type CreateResponse struct { + Handle *int64 +} + +// Delete a directory. +type DeleteDirectoryRequest struct { + DirectoryPath *string +} + +type DeleteDirectoryResponse struct { +} + +// Delete a file. +type DeleteFileRequest struct { + FilePath *string +} + +type DeleteFileResponse struct { +} + +type DeleteRequest struct { + Path *string + Recursive *bool +} + +type DeleteResponse struct { +} + +type DirectoryEntry struct { + FileSize *int + IsDirectory *bool + LastModified *int + Name *string + Path *string +} + +// Download a file. +type DownloadFileRequest struct { + FilePath *string + Range *string + IfUnmodifiedSince *string +} + +type DownloadFileResponse struct { + ContentLength *int64 + ContentType *string + Contents io.ReadCloser + LastModified *string +} + +// Stores the attributes of a file or directory.. +type FileInfo struct { + Path *string + IsDir *bool + FileSize *int64 + ModificationTime *int64 +} + +// Get directory metadata. +type GetDirectoryMetadataRequest struct { + DirectoryPath *string +} + +type GetDirectoryMetadataResponse struct { +} + +// Get file metadata. +type GetFileMetadataRequest struct { + FilePath *string + Range *string + IfUnmodifiedSince *string +} + +type GetFileMetadataResponse struct { + ContentLength *int64 + ContentType *string + LastModified *string +} + +type GetStatusRequest struct { + Path *string +} + +type GetStatusResponse struct { + Path *string + IsDir *bool + FileSize *int64 + ModificationTime *int64 +} + +// List directory contents. +type ListDirectoryContentsRequest struct { + DirectoryPath *string + PageSize *int64 + PageToken *string +} + +type ListDirectoryResponse struct { + Contents []DirectoryEntry + NextPageToken *string +} + +type ListStatusRequest struct { + Path *string +} + +type ListStatusResponse struct { + Files []FileInfo +} + +type MkDirsRequest struct { + Path *string +} + +type MkDirsResponse struct { +} + +type MoveRequest struct { + SourcePath *string + DestinationPath *string +} + +type MoveResponse struct { +} + +type PutRequest struct { + Path *string + Contents []byte + Overwrite *bool +} + +type PutResponse struct { +} + +type ReadRequest struct { + Path *string + Offset *int64 + Length *int64 +} + +type ReadResponse struct { + BytesRead *int64 + Data []byte +} + +// Upload a file. +type UploadFileRequest struct { + FilePath *string + Contents io.ReadCloser + Overwrite *bool +} + +type UploadFileResponse struct { +} diff --git a/libs/tmp/files/v2/wire.go b/libs/tmp/files/v2/wire.go new file mode 100644 index 00000000000..fcea4363e38 --- /dev/null +++ b/libs/tmp/files/v2/wire.go @@ -0,0 +1,854 @@ +// Code generated by Databricks SDK Generator. DO NOT EDIT. + +package files + +type addBlockRequestWire struct { + Handle *int64 `json:"handle,omitempty"` + Data []byte `json:"data,omitempty"` +} + +func addBlockRequestToWire(v *AddBlockRequest) *addBlockRequestWire { + if v == nil { + return nil + } + return &addBlockRequestWire{ + Handle: v.Handle, + Data: v.Data, + } +} + +func addBlockRequestFromWire(w *addBlockRequestWire) *AddBlockRequest { + if w == nil { + return nil + } + return &AddBlockRequest{ + Handle: w.Handle, + Data: w.Data, + } +} + +type addBlockResponseWire struct { +} + +func addBlockResponseToWire(v *AddBlockResponse) *addBlockResponseWire { + if v == nil { + return nil + } + return &addBlockResponseWire{} +} + +func addBlockResponseFromWire(w *addBlockResponseWire) *AddBlockResponse { + if w == nil { + return nil + } + return &AddBlockResponse{} +} + +type closeRequestWire struct { + Handle *int64 `json:"handle,omitempty"` +} + +func closeRequestToWire(v *CloseRequest) *closeRequestWire { + if v == nil { + return nil + } + return &closeRequestWire{ + Handle: v.Handle, + } +} + +func closeRequestFromWire(w *closeRequestWire) *CloseRequest { + if w == nil { + return nil + } + return &CloseRequest{ + Handle: w.Handle, + } +} + +type closeResponseWire struct { +} + +func closeResponseToWire(v *CloseResponse) *closeResponseWire { + if v == nil { + return nil + } + return &closeResponseWire{} +} + +func closeResponseFromWire(w *closeResponseWire) *CloseResponse { + if w == nil { + return nil + } + return &CloseResponse{} +} + +type createDirectoryRequestWire struct { + DirectoryPath *string `json:"directory_path,omitempty"` +} + +func createDirectoryRequestToWire(v *CreateDirectoryRequest) *createDirectoryRequestWire { + if v == nil { + return nil + } + return &createDirectoryRequestWire{ + DirectoryPath: v.DirectoryPath, + } +} + +func createDirectoryRequestFromWire(w *createDirectoryRequestWire) *CreateDirectoryRequest { + if w == nil { + return nil + } + return &CreateDirectoryRequest{ + DirectoryPath: w.DirectoryPath, + } +} + +type createDirectoryResponseWire struct { +} + +func createDirectoryResponseToWire(v *CreateDirectoryResponse) *createDirectoryResponseWire { + if v == nil { + return nil + } + return &createDirectoryResponseWire{} +} + +func createDirectoryResponseFromWire(w *createDirectoryResponseWire) *CreateDirectoryResponse { + if w == nil { + return nil + } + return &CreateDirectoryResponse{} +} + +type createRequestWire struct { + Path *string `json:"path,omitempty"` + Overwrite *bool `json:"overwrite,omitempty"` +} + +func createRequestToWire(v *CreateRequest) *createRequestWire { + if v == nil { + return nil + } + return &createRequestWire{ + Path: v.Path, + Overwrite: v.Overwrite, + } +} + +func createRequestFromWire(w *createRequestWire) *CreateRequest { + if w == nil { + return nil + } + return &CreateRequest{ + Path: w.Path, + Overwrite: w.Overwrite, + } +} + +type createResponseWire struct { + Handle *int64 `json:"handle,omitempty"` +} + +func createResponseToWire(v *CreateResponse) *createResponseWire { + if v == nil { + return nil + } + return &createResponseWire{ + Handle: v.Handle, + } +} + +func createResponseFromWire(w *createResponseWire) *CreateResponse { + if w == nil { + return nil + } + return &CreateResponse{ + Handle: w.Handle, + } +} + +type deleteDirectoryRequestWire struct { + DirectoryPath *string `json:"directory_path,omitempty"` +} + +func deleteDirectoryRequestToWire(v *DeleteDirectoryRequest) *deleteDirectoryRequestWire { + if v == nil { + return nil + } + return &deleteDirectoryRequestWire{ + DirectoryPath: v.DirectoryPath, + } +} + +func deleteDirectoryRequestFromWire(w *deleteDirectoryRequestWire) *DeleteDirectoryRequest { + if w == nil { + return nil + } + return &DeleteDirectoryRequest{ + DirectoryPath: w.DirectoryPath, + } +} + +type deleteDirectoryResponseWire struct { +} + +func deleteDirectoryResponseToWire(v *DeleteDirectoryResponse) *deleteDirectoryResponseWire { + if v == nil { + return nil + } + return &deleteDirectoryResponseWire{} +} + +func deleteDirectoryResponseFromWire(w *deleteDirectoryResponseWire) *DeleteDirectoryResponse { + if w == nil { + return nil + } + return &DeleteDirectoryResponse{} +} + +type deleteFileRequestWire struct { + FilePath *string `json:"file_path,omitempty"` +} + +func deleteFileRequestToWire(v *DeleteFileRequest) *deleteFileRequestWire { + if v == nil { + return nil + } + return &deleteFileRequestWire{ + FilePath: v.FilePath, + } +} + +func deleteFileRequestFromWire(w *deleteFileRequestWire) *DeleteFileRequest { + if w == nil { + return nil + } + return &DeleteFileRequest{ + FilePath: w.FilePath, + } +} + +type deleteFileResponseWire struct { +} + +func deleteFileResponseToWire(v *DeleteFileResponse) *deleteFileResponseWire { + if v == nil { + return nil + } + return &deleteFileResponseWire{} +} + +func deleteFileResponseFromWire(w *deleteFileResponseWire) *DeleteFileResponse { + if w == nil { + return nil + } + return &DeleteFileResponse{} +} + +type deleteRequestWire struct { + Path *string `json:"path,omitempty"` + Recursive *bool `json:"recursive,omitempty"` +} + +func deleteRequestToWire(v *DeleteRequest) *deleteRequestWire { + if v == nil { + return nil + } + return &deleteRequestWire{ + Path: v.Path, + Recursive: v.Recursive, + } +} + +func deleteRequestFromWire(w *deleteRequestWire) *DeleteRequest { + if w == nil { + return nil + } + return &DeleteRequest{ + Path: w.Path, + Recursive: w.Recursive, + } +} + +type deleteResponseWire struct { +} + +func deleteResponseToWire(v *DeleteResponse) *deleteResponseWire { + if v == nil { + return nil + } + return &deleteResponseWire{} +} + +func deleteResponseFromWire(w *deleteResponseWire) *DeleteResponse { + if w == nil { + return nil + } + return &DeleteResponse{} +} + +type directoryEntryWire struct { + FileSize *int `json:"file_size,omitempty"` + IsDirectory *bool `json:"is_directory,omitempty"` + LastModified *int `json:"last_modified,omitempty"` + Name *string `json:"name,omitempty"` + Path *string `json:"path,omitempty"` +} + +func directoryEntryToWire(v *DirectoryEntry) *directoryEntryWire { + if v == nil { + return nil + } + return &directoryEntryWire{ + FileSize: v.FileSize, + IsDirectory: v.IsDirectory, + LastModified: v.LastModified, + Name: v.Name, + Path: v.Path, + } +} + +func directoryEntryFromWire(w *directoryEntryWire) *DirectoryEntry { + if w == nil { + return nil + } + return &DirectoryEntry{ + FileSize: w.FileSize, + IsDirectory: w.IsDirectory, + LastModified: w.LastModified, + Name: w.Name, + Path: w.Path, + } +} + +type downloadFileRequestWire struct { + FilePath *string `json:"file_path,omitempty"` +} + +func downloadFileRequestToWire(v *DownloadFileRequest) *downloadFileRequestWire { + if v == nil { + return nil + } + return &downloadFileRequestWire{ + FilePath: v.FilePath, + } +} + +func downloadFileRequestFromWire(w *downloadFileRequestWire) *DownloadFileRequest { + if w == nil { + return nil + } + return &DownloadFileRequest{ + FilePath: w.FilePath, + } +} + +type downloadFileResponseWire struct { +} + +func downloadFileResponseToWire(v *DownloadFileResponse) *downloadFileResponseWire { + if v == nil { + return nil + } + return &downloadFileResponseWire{} +} + +func downloadFileResponseFromWire(w *downloadFileResponseWire) *DownloadFileResponse { + if w == nil { + return nil + } + return &DownloadFileResponse{} +} + +type fileInfoWire struct { + Path *string `json:"path,omitempty"` + IsDir *bool `json:"is_dir,omitempty"` + FileSize *int64 `json:"file_size,omitempty"` + ModificationTime *int64 `json:"modification_time,omitempty"` +} + +func fileInfoToWire(v *FileInfo) *fileInfoWire { + if v == nil { + return nil + } + return &fileInfoWire{ + Path: v.Path, + IsDir: v.IsDir, + FileSize: v.FileSize, + ModificationTime: v.ModificationTime, + } +} + +func fileInfoFromWire(w *fileInfoWire) *FileInfo { + if w == nil { + return nil + } + return &FileInfo{ + Path: w.Path, + IsDir: w.IsDir, + FileSize: w.FileSize, + ModificationTime: w.ModificationTime, + } +} + +type getDirectoryMetadataRequestWire struct { + DirectoryPath *string `json:"directory_path,omitempty"` +} + +func getDirectoryMetadataRequestToWire(v *GetDirectoryMetadataRequest) *getDirectoryMetadataRequestWire { + if v == nil { + return nil + } + return &getDirectoryMetadataRequestWire{ + DirectoryPath: v.DirectoryPath, + } +} + +func getDirectoryMetadataRequestFromWire(w *getDirectoryMetadataRequestWire) *GetDirectoryMetadataRequest { + if w == nil { + return nil + } + return &GetDirectoryMetadataRequest{ + DirectoryPath: w.DirectoryPath, + } +} + +type getDirectoryMetadataResponseWire struct { +} + +func getDirectoryMetadataResponseToWire(v *GetDirectoryMetadataResponse) *getDirectoryMetadataResponseWire { + if v == nil { + return nil + } + return &getDirectoryMetadataResponseWire{} +} + +func getDirectoryMetadataResponseFromWire(w *getDirectoryMetadataResponseWire) *GetDirectoryMetadataResponse { + if w == nil { + return nil + } + return &GetDirectoryMetadataResponse{} +} + +type getFileMetadataRequestWire struct { + FilePath *string `json:"file_path,omitempty"` +} + +func getFileMetadataRequestToWire(v *GetFileMetadataRequest) *getFileMetadataRequestWire { + if v == nil { + return nil + } + return &getFileMetadataRequestWire{ + FilePath: v.FilePath, + } +} + +func getFileMetadataRequestFromWire(w *getFileMetadataRequestWire) *GetFileMetadataRequest { + if w == nil { + return nil + } + return &GetFileMetadataRequest{ + FilePath: w.FilePath, + } +} + +type getFileMetadataResponseWire struct { +} + +func getFileMetadataResponseToWire(v *GetFileMetadataResponse) *getFileMetadataResponseWire { + if v == nil { + return nil + } + return &getFileMetadataResponseWire{} +} + +func getFileMetadataResponseFromWire(w *getFileMetadataResponseWire) *GetFileMetadataResponse { + if w == nil { + return nil + } + return &GetFileMetadataResponse{} +} + +type getStatusRequestWire struct { + Path *string `json:"path,omitempty"` +} + +func getStatusRequestToWire(v *GetStatusRequest) *getStatusRequestWire { + if v == nil { + return nil + } + return &getStatusRequestWire{ + Path: v.Path, + } +} + +func getStatusRequestFromWire(w *getStatusRequestWire) *GetStatusRequest { + if w == nil { + return nil + } + return &GetStatusRequest{ + Path: w.Path, + } +} + +type getStatusResponseWire struct { + Path *string `json:"path,omitempty"` + IsDir *bool `json:"is_dir,omitempty"` + FileSize *int64 `json:"file_size,omitempty"` + ModificationTime *int64 `json:"modification_time,omitempty"` +} + +func getStatusResponseToWire(v *GetStatusResponse) *getStatusResponseWire { + if v == nil { + return nil + } + return &getStatusResponseWire{ + Path: v.Path, + IsDir: v.IsDir, + FileSize: v.FileSize, + ModificationTime: v.ModificationTime, + } +} + +func getStatusResponseFromWire(w *getStatusResponseWire) *GetStatusResponse { + if w == nil { + return nil + } + return &GetStatusResponse{ + Path: w.Path, + IsDir: w.IsDir, + FileSize: w.FileSize, + ModificationTime: w.ModificationTime, + } +} + +type listDirectoryContentsRequestWire struct { + DirectoryPath *string `json:"directory_path,omitempty"` + PageSize *int64 `json:"page_size,omitempty"` + PageToken *string `json:"page_token,omitempty"` +} + +func listDirectoryContentsRequestToWire(v *ListDirectoryContentsRequest) *listDirectoryContentsRequestWire { + if v == nil { + return nil + } + return &listDirectoryContentsRequestWire{ + DirectoryPath: v.DirectoryPath, + PageSize: v.PageSize, + PageToken: v.PageToken, + } +} + +func listDirectoryContentsRequestFromWire(w *listDirectoryContentsRequestWire) *ListDirectoryContentsRequest { + if w == nil { + return nil + } + return &ListDirectoryContentsRequest{ + DirectoryPath: w.DirectoryPath, + PageSize: w.PageSize, + PageToken: w.PageToken, + } +} + +type listDirectoryResponseWire struct { + Contents []directoryEntryWire `json:"contents,omitempty"` + NextPageToken *string `json:"next_page_token,omitempty"` +} + +func listDirectoryResponseToWire(v *ListDirectoryResponse) *listDirectoryResponseWire { + if v == nil { + return nil + } + return &listDirectoryResponseWire{ + Contents: convertSlice(v.Contents, directoryEntryToWire), + NextPageToken: v.NextPageToken, + } +} + +func listDirectoryResponseFromWire(w *listDirectoryResponseWire) *ListDirectoryResponse { + if w == nil { + return nil + } + return &ListDirectoryResponse{ + Contents: convertSlice(w.Contents, directoryEntryFromWire), + NextPageToken: w.NextPageToken, + } +} + +type listStatusRequestWire struct { + Path *string `json:"path,omitempty"` +} + +func listStatusRequestToWire(v *ListStatusRequest) *listStatusRequestWire { + if v == nil { + return nil + } + return &listStatusRequestWire{ + Path: v.Path, + } +} + +func listStatusRequestFromWire(w *listStatusRequestWire) *ListStatusRequest { + if w == nil { + return nil + } + return &ListStatusRequest{ + Path: w.Path, + } +} + +type listStatusResponseWire struct { + Files []fileInfoWire `json:"files,omitempty"` +} + +func listStatusResponseToWire(v *ListStatusResponse) *listStatusResponseWire { + if v == nil { + return nil + } + return &listStatusResponseWire{ + Files: convertSlice(v.Files, fileInfoToWire), + } +} + +func listStatusResponseFromWire(w *listStatusResponseWire) *ListStatusResponse { + if w == nil { + return nil + } + return &ListStatusResponse{ + Files: convertSlice(w.Files, fileInfoFromWire), + } +} + +type mkDirsRequestWire struct { + Path *string `json:"path,omitempty"` +} + +func mkDirsRequestToWire(v *MkDirsRequest) *mkDirsRequestWire { + if v == nil { + return nil + } + return &mkDirsRequestWire{ + Path: v.Path, + } +} + +func mkDirsRequestFromWire(w *mkDirsRequestWire) *MkDirsRequest { + if w == nil { + return nil + } + return &MkDirsRequest{ + Path: w.Path, + } +} + +type mkDirsResponseWire struct { +} + +func mkDirsResponseToWire(v *MkDirsResponse) *mkDirsResponseWire { + if v == nil { + return nil + } + return &mkDirsResponseWire{} +} + +func mkDirsResponseFromWire(w *mkDirsResponseWire) *MkDirsResponse { + if w == nil { + return nil + } + return &MkDirsResponse{} +} + +type moveRequestWire struct { + SourcePath *string `json:"source_path,omitempty"` + DestinationPath *string `json:"destination_path,omitempty"` +} + +func moveRequestToWire(v *MoveRequest) *moveRequestWire { + if v == nil { + return nil + } + return &moveRequestWire{ + SourcePath: v.SourcePath, + DestinationPath: v.DestinationPath, + } +} + +func moveRequestFromWire(w *moveRequestWire) *MoveRequest { + if w == nil { + return nil + } + return &MoveRequest{ + SourcePath: w.SourcePath, + DestinationPath: w.DestinationPath, + } +} + +type moveResponseWire struct { +} + +func moveResponseToWire(v *MoveResponse) *moveResponseWire { + if v == nil { + return nil + } + return &moveResponseWire{} +} + +func moveResponseFromWire(w *moveResponseWire) *MoveResponse { + if w == nil { + return nil + } + return &MoveResponse{} +} + +type putRequestWire struct { + Path *string `json:"path,omitempty"` + Contents []byte `json:"contents,omitempty"` + Overwrite *bool `json:"overwrite,omitempty"` +} + +func putRequestToWire(v *PutRequest) *putRequestWire { + if v == nil { + return nil + } + return &putRequestWire{ + Path: v.Path, + Contents: v.Contents, + Overwrite: v.Overwrite, + } +} + +func putRequestFromWire(w *putRequestWire) *PutRequest { + if w == nil { + return nil + } + return &PutRequest{ + Path: w.Path, + Contents: w.Contents, + Overwrite: w.Overwrite, + } +} + +type putResponseWire struct { +} + +func putResponseToWire(v *PutResponse) *putResponseWire { + if v == nil { + return nil + } + return &putResponseWire{} +} + +func putResponseFromWire(w *putResponseWire) *PutResponse { + if w == nil { + return nil + } + return &PutResponse{} +} + +type readRequestWire struct { + Path *string `json:"path,omitempty"` + Offset *int64 `json:"offset,omitempty"` + Length *int64 `json:"length,omitempty"` +} + +func readRequestToWire(v *ReadRequest) *readRequestWire { + if v == nil { + return nil + } + return &readRequestWire{ + Path: v.Path, + Offset: v.Offset, + Length: v.Length, + } +} + +func readRequestFromWire(w *readRequestWire) *ReadRequest { + if w == nil { + return nil + } + return &ReadRequest{ + Path: w.Path, + Offset: w.Offset, + Length: w.Length, + } +} + +type readResponseWire struct { + BytesRead *int64 `json:"bytes_read,omitempty"` + Data []byte `json:"data,omitempty"` +} + +func readResponseToWire(v *ReadResponse) *readResponseWire { + if v == nil { + return nil + } + return &readResponseWire{ + BytesRead: v.BytesRead, + Data: v.Data, + } +} + +func readResponseFromWire(w *readResponseWire) *ReadResponse { + if w == nil { + return nil + } + return &ReadResponse{ + BytesRead: w.BytesRead, + Data: w.Data, + } +} + +type uploadFileRequestWire struct { + FilePath *string `json:"file_path,omitempty"` + Overwrite *bool `json:"overwrite,omitempty"` +} + +func uploadFileRequestToWire(v *UploadFileRequest) *uploadFileRequestWire { + if v == nil { + return nil + } + return &uploadFileRequestWire{ + FilePath: v.FilePath, + Overwrite: v.Overwrite, + } +} + +func uploadFileRequestFromWire(w *uploadFileRequestWire) *UploadFileRequest { + if w == nil { + return nil + } + return &UploadFileRequest{ + FilePath: w.FilePath, + Overwrite: w.Overwrite, + } +} + +type uploadFileResponseWire struct { +} + +func uploadFileResponseToWire(v *UploadFileResponse) *uploadFileResponseWire { + if v == nil { + return nil + } + return &uploadFileResponseWire{} +} + +func uploadFileResponseFromWire(w *uploadFileResponseWire) *UploadFileResponse { + if w == nil { + return nil + } + return &UploadFileResponse{} +} + +func convertSlice[T, W any](s []T, conv func(*T) *W) []W { + if s == nil { + return nil + } + out := make([]W, len(s)) + for i := range s { + out[i] = *conv(&s[i]) + } + return out +} diff --git a/libs/tmp/options/call/call.go b/libs/tmp/options/call/call.go new file mode 100644 index 00000000000..7248237ecfb --- /dev/null +++ b/libs/tmp/options/call/call.go @@ -0,0 +1,53 @@ +// Package call defines the options used to configure individual calls +// against the Databricks API. +package call + +import ( + "time" + + "github.com/databricks/sdk-go/core/ops" + "github.com/databricks/cli/libs/tmp/options/internaloptions" +) + +// Option configures a single call against the Databricks API. +type Option func(*internaloptions.CallOptions) error + +// WithRetrier returns an Option that uses the given Retrier provider. If no +// retrier is provided, the call is not retried. +// +// The provider function must be thread-safe. +func WithRetrier(provider func() ops.Retrier) Option { + return func(c *internaloptions.CallOptions) error { + c.Retrier = provider + return nil + } +} + +// WithDisableRetry is a convenience option that disables retries. +func WithDisableRetry() Option { + return func(c *internaloptions.CallOptions) error { + c.Retrier = nil + return nil + } +} + +// WithTimeout returns an Option that sets the timeout for the call. If the +// context already has a deadline, it is updated to the minimum of the +// context's deadline and the timeout. A timeout of zero means no timeout. +// +// The timeout covers the entire execution, including retries. +func WithTimeout(t time.Duration) Option { + return func(c *internaloptions.CallOptions) error { + c.Timeout = t + return nil + } +} + +// WithLimiter returns an Option that uses the given Limiter. If no limiter is +// provided, the call is not rate limited. +func WithLimiter(l ops.Limiter) Option { + return func(c *internaloptions.CallOptions) error { + c.RateLimiter = l + return nil + } +} diff --git a/libs/tmp/options/call/call_test.go b/libs/tmp/options/call/call_test.go new file mode 100644 index 00000000000..59aa59b0ce9 --- /dev/null +++ b/libs/tmp/options/call/call_test.go @@ -0,0 +1,58 @@ +package call + +import ( + "context" + "testing" + "time" + + "github.com/databricks/sdk-go/core/ops" + "github.com/databricks/cli/libs/tmp/options/internaloptions" +) + +func TestOptionsApply_AllFields(t *testing.T) { + provider := func() ops.Retrier { return nil } + limiter := stubLimiter{} + + opts := []Option{ + WithRetrier(provider), + WithLimiter(limiter), + WithTimeout(3 * time.Second), + } + + cfg := internaloptions.CallOptions{} + for _, opt := range opts { + if err := opt(&cfg); err != nil { + t.Fatalf("Apply: %v", err) + } + } + + if cfg.Retrier == nil { + t.Error("expected Retrier provider") + } + if cfg.RateLimiter != ops.Limiter(limiter) { + t.Error("RateLimiter mismatch") + } + if cfg.Timeout != 3*time.Second { + t.Errorf("Timeout = %v", cfg.Timeout) + } +} + +func TestWithDisableRetry_OverridesEarlierRetrier(t *testing.T) { + cfg := internaloptions.CallOptions{} + if err := WithRetrier(func() ops.Retrier { return nil })(&cfg); err != nil { + t.Fatalf("Apply: %v", err) + } + if cfg.Retrier == nil { + t.Fatal("expected Retrier to be set after WithRetrier") + } + if err := WithDisableRetry()(&cfg); err != nil { + t.Fatalf("Apply: %v", err) + } + if cfg.Retrier != nil { + t.Fatal("expected Retrier to be cleared after WithDisableRetry") + } +} + +type stubLimiter struct{} + +func (stubLimiter) Wait(_ context.Context) error { return nil } diff --git a/libs/tmp/options/client/client.go b/libs/tmp/options/client/client.go new file mode 100644 index 00000000000..8d75a574bd0 --- /dev/null +++ b/libs/tmp/options/client/client.go @@ -0,0 +1,119 @@ +// Package client defines the options used to configure Databricks API +// clients. +// +// Databricks API clients can resolve their options from a variety of sources: +// +// - Profile from file or default. +// - Environment variables override. +// - Explicit value override. +// +// If no credentials are provided, the credentials are automatically resolved +// from the client configuration after the above chain of resolution. +// +// In practice, we recommend users to stick to a single resolution level. +// That is either rely purely on automatic resolution from the environment or +// explicitly set all options in code. Mixing the two approaches is not recommended. +package client + +import ( + "log/slog" + "net/http" + "time" + + "github.com/databricks/cli/libs/tmp/auth" + "github.com/databricks/cli/libs/tmp/options/internaloptions" +) + +// Option configures a Databricks API client. +type Option func(*internaloptions.ClientOptions) error + +// WithHost returns an Option that sets the host for the client. +func WithHost(h string) Option { + return func(c *internaloptions.ClientOptions) error { + c.Host = h + return nil + } +} + +// WithHTTPClient returns an Option that uses a specific HTTP client when +// making HTTP requests. +// +// Important: When set, this option ignores all other options. +func WithHTTPClient(hc *http.Client) Option { + return func(c *internaloptions.ClientOptions) error { + c.HTTPClient = hc + return nil + } +} + +// WithCredentials returns an Option that sets a specific credentials. +func WithCredentials(creds auth.Credentials) Option { + return func(c *internaloptions.ClientOptions) error { + c.Credentials = creds + return nil + } +} + +// WithTimeout returns an Option that sets the overall API call timeout to +// the given duration by default. +func WithTimeout(d time.Duration) Option { + return func(c *internaloptions.ClientOptions) error { + c.Timeout = d + return nil + } +} + +// WithLogger returns an Option that uses the provided logger. Log messages +// are only logged if the logger is enabled. +func WithLogger(l *slog.Logger) Option { + return func(c *internaloptions.ClientOptions) error { + c.Logger = l + return nil + } +} + +// WithAccountID returns an Option that sets the account ID for the client. +func WithAccountID(id string) Option { + return func(c *internaloptions.ClientOptions) error { + c.AccountID = id + return nil + } +} + +// WithWorkspaceID returns an Option that sets the workspace ID for the client. +func WithWorkspaceID(id string) Option { + return func(c *internaloptions.ClientOptions) error { + c.WorkspaceID = id + return nil + } +} + +// WithoutProfileResolution returns an Option that entirely disables profile +// resolution. This is useful when you want your client to only be explicitly +// configured in code. +func WithoutProfileResolution() Option { + return func(c *internaloptions.ClientOptions) error { + c.DisableProfileResolution = true + return nil + } +} + +// WithProfileFile returns an Option that sets the profile file to use for +// profile resolution. By default, the profile file is resolved from the +// environment variable $DATABRICKS_CONFIG_FILE. +func WithProfileFile(file string) Option { + return func(c *internaloptions.ClientOptions) error { + c.ProfileFile = file + return nil + } +} + +// WithProfile returns an Option that sets the profile name to use for +// profile resolution. By default, the profile name is resolved from the +// environment variable $DATABRICKS_CONFIG_NAME. +func WithProfile(name string) Option { + return func(c *internaloptions.ClientOptions) error { + c.ProfileName = name + return nil + } +} diff --git a/libs/tmp/options/client/client_test.go b/libs/tmp/options/client/client_test.go new file mode 100644 index 00000000000..e1bdc315916 --- /dev/null +++ b/libs/tmp/options/client/client_test.go @@ -0,0 +1,57 @@ +package client + +import ( + "context" + "log/slog" + "net/http" + "testing" + "time" + + "github.com/databricks/cli/libs/tmp/auth" + "github.com/databricks/cli/libs/tmp/options/internaloptions" +) + +func TestOptionsApply_AllFields(t *testing.T) { + httpClient := &http.Client{} + creds := stubCredentials{} + logger := slog.Default() + + opts := []Option{ + WithHost("https://example.cloud.databricks.com"), + WithHTTPClient(httpClient), + WithCredentials(creds), + WithTimeout(7 * time.Second), + WithLogger(logger), + } + + cfg := internaloptions.ClientOptions{} + for _, opt := range opts { + if err := opt(&cfg); err != nil { + t.Fatalf("Apply: %v", err) + } + } + + if cfg.Host != "https://example.cloud.databricks.com" { + t.Errorf("Host = %q", cfg.Host) + } + if cfg.HTTPClient != httpClient { + t.Error("HTTPClient mismatch") + } + if cfg.Credentials != auth.Credentials(creds) { + t.Error("Credentials mismatch") + } + if cfg.Timeout != 7*time.Second { + t.Errorf("Timeout = %v", cfg.Timeout) + } + if cfg.Logger != logger { + t.Error("Logger mismatch") + } +} + +type stubCredentials struct{} + +func (stubCredentials) Name() string { return "stub" } + +func (stubCredentials) AuthHeaders(_ context.Context) ([]auth.Header, error) { + return nil, nil +} diff --git a/libs/tmp/options/internal/version.go b/libs/tmp/options/internal/version.go new file mode 100644 index 00000000000..ab946dc5c93 --- /dev/null +++ b/libs/tmp/options/internal/version.go @@ -0,0 +1,5 @@ +package internal + +const ModuleName = "sdk-go-options" + +const Version = "0.0.0-dev" diff --git a/libs/tmp/options/internaloptions/internaloptions.go b/libs/tmp/options/internaloptions/internaloptions.go new file mode 100644 index 00000000000..1cf71bfc7a8 --- /dev/null +++ b/libs/tmp/options/internaloptions/internaloptions.go @@ -0,0 +1,113 @@ +// Package internaloptions contains the resolved options types produced +// by applying [github.com/databricks/cli/libs/tmp/options/client.Option] and +// [github.com/databricks/cli/libs/tmp/options/call.Option] values. +// +// IMPORTANT: This package is NOT part of the public API of the Databricks +// SDK. Its contents may change at any time without notice. Clients should +// not directly depend on functionalities in this package. +package internaloptions + +import ( + "errors" + "log/slog" + "net/http" + "time" + + "github.com/databricks/cli/libs/tmp/auth" + "github.com/databricks/sdk-go/core/ops" + "github.com/databricks/sdk-go/core/profiles" +) + +// ClientOptions is the resolved client configuration produced by applying +// client.Option values. +type ClientOptions struct { + // Profile resolution. + ProfileName string + ProfileFile string + DisableProfileResolution bool + + Host string + AccountID string + WorkspaceID string + + Credentials auth.Credentials + + Timeout time.Duration + HTTPClient *http.Client + Logger *slog.Logger +} + +// Resolve fills in defaults and validates the resolved client options. +// +// Resolve always populates the HTTPClient and Logger fields with default +// values if not provided. +func (c *ClientOptions) Resolve() error { + if c.Logger == nil { + c.Logger = slog.New(slog.DiscardHandler) + } + if c.HTTPClient == nil { + c.HTTPClient = &http.Client{ + Timeout: c.Timeout, + // Share the default transport so the connection pool is reused + // across all clients. + Transport: http.DefaultTransport, + } + } + + if err := c.resolve(); err != nil { + return err + } + return c.validate() +} + +// resolve fills unset options from the profile. Explicitly set options take +// precedence and are never overwritten. +// +// TODO: Apply environment-variable overrides, and resolve workspace/account ID +// and credentials from the profile when not provided. +func (c *ClientOptions) resolve() error { + if c.DisableProfileResolution { + return nil + } + + var opts []profiles.ResolveOption + if c.ProfileName != "" { + opts = append(opts, profiles.WithProfile(c.ProfileName)) + } else { + opts = append(opts, profiles.WithDefaultProfile()) + } + if c.ProfileFile != "" { + opts = append(opts, profiles.WithFile(c.ProfileFile)) + } + + p, err := profiles.Resolve(opts...) + if err != nil { + return err + } + + if c.Host == "" { + c.Host = p.Host + } + if c.AccountID == "" { + c.AccountID = p.AccountID + } + if c.WorkspaceID == "" { + c.WorkspaceID = p.WorkspaceID + } + return nil +} + +func (c *ClientOptions) validate() error { + if c.Credentials == nil { + return errors.New("credentials are required") + } + return nil +} + +// CallOptions is the resolved per-call configuration produced by applying +// call.Option values. +type CallOptions struct { + Retrier func() ops.Retrier + RateLimiter ops.Limiter + Timeout time.Duration +} diff --git a/libs/tmp/options/internaloptions/internaloptions_test.go b/libs/tmp/options/internaloptions/internaloptions_test.go new file mode 100644 index 00000000000..233d976b031 --- /dev/null +++ b/libs/tmp/options/internaloptions/internaloptions_test.go @@ -0,0 +1,41 @@ +package internaloptions + +import ( + "context" + "log/slog" + "testing" + + "github.com/databricks/cli/libs/tmp/auth" +) + +// stubCredentials satisfies the credentials requirement of Resolve for tests +// that exercise other resolution behavior (e.g. logger defaulting). +type stubCredentials struct{} + +func (stubCredentials) Name() string { return "stub" } + +func (stubCredentials) AuthHeaders(context.Context) ([]auth.Header, error) { return nil, nil } + +func TestClientOptionsResolve_DefaultLogger(t *testing.T) { + c := &ClientOptions{Credentials: stubCredentials{}} + if err := c.Resolve(); err != nil { + t.Fatalf("Resolve: %v", err) + } + if c.Logger == nil { + t.Fatal("expected default logger to be set") + } + if c.Logger.Enabled(context.Background(), slog.LevelError) { + t.Fatal("expected default logger to be disabled") + } +} + +func TestClientOptionsResolve_PreservesProvidedLogger(t *testing.T) { + provided := slog.Default() + c := &ClientOptions{Logger: provided, Credentials: stubCredentials{}} + if err := c.Resolve(); err != nil { + t.Fatalf("Resolve: %v", err) + } + if c.Logger != provided { + t.Fatal("expected provided logger to be preserved") + } +} diff --git a/tools/check_deadcode.py b/tools/check_deadcode.py index 5c37f1f4750..d3dd8dd43e0 100755 --- a/tools/check_deadcode.py +++ b/tools/check_deadcode.py @@ -42,6 +42,7 @@ EXCLUDED_DIRS = [ "libs/gorules/", # Lint rule definitions loaded by golangci-lint's ruleguard "bundle/internal/tf/schema/", # Generated from Terraform provider schema + "libs/tmp/", # Temporary verbatim copy of the unpublished sdk-go/files/v2; keep upstream API intact ] ALLOW_COMMENT = "//deadcode:allow"