From 39fe34792a67c7126c547ebec012db13bacc1d19 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 22 Jul 2026 04:47:02 +0000 Subject: [PATCH 1/4] 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" From 3a6d8ebd2c8fff26629e3ae8b2d9deaf1c4bab3e Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Fri, 24 Jul 2026 08:49:19 +0000 Subject: [PATCH 2/4] fs: depend on published sdk-go modules, drop libs/tmp vendored copy The sdk-go files/v2 client and its auth/options/core dependencies are now published to the Databricks Go proxy, so the temporary in-repo copy under libs/tmp is no longer needed. Replaces it with real module dependencies (files v0.0.0-dev.1, auth v0.0.0-dev, options v0.0.0-dev, core bumped to v0.0.1-dev), rewrites the filer and fs command imports from github.com/databricks/cli/libs/tmp/... back to github.com/databricks/sdk-go/..., and removes libs/tmp along with its carve-outs in .golangci.yaml, tools/check_deadcode.py, and .github/OWNERS. The two local fixes that had been applied to the vendored generated client (empty response bodies and HTTP method casing) are fixed upstream in the published version, so nothing is lost. This pull request and its description were written by Isaac. Co-authored-by: Isaac --- .github/OWNERS | 1 - .golangci.yaml | 10 - cmd/fs/upload_progress.go | 2 +- go.mod | 8 +- go.sum | 10 +- internal/build/notice_test.go | 3 + libs/filer/files_client.go | 14 +- libs/filer/files_client_auth.go | 10 +- libs/filer/files_client_test.go | 2 +- 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 - 48 files changed, 31 insertions(+), 8270 deletions(-) delete mode 100644 libs/tmp/README.md delete mode 100644 libs/tmp/auth/auth.go delete mode 100644 libs/tmp/auth/auth_test.go delete mode 100644 libs/tmp/auth/cache.go delete mode 100644 libs/tmp/auth/cache_test.go delete mode 100644 libs/tmp/auth/retrying.go delete mode 100644 libs/tmp/auth/retrying_test.go delete mode 100644 libs/tmp/files/internal/version.go delete mode 100644 libs/tmp/files/v2/client.go delete mode 100644 libs/tmp/files/v2/ext_limiter.go delete mode 100644 libs/tmp/files/v2/ext_limiter_test.go delete mode 100644 libs/tmp/files/v2/ext_multipart.go delete mode 100644 libs/tmp/files/v2/ext_resumable.go delete mode 100644 libs/tmp/files/v2/ext_resumable_test.go delete mode 100644 libs/tmp/files/v2/ext_straggler.go delete mode 100644 libs/tmp/files/v2/ext_straggler_test.go delete mode 100644 libs/tmp/files/v2/ext_tunables.go delete mode 100644 libs/tmp/files/v2/ext_upload.go delete mode 100644 libs/tmp/files/v2/ext_upload_control.go delete mode 100644 libs/tmp/files/v2/ext_upload_control_test.go delete mode 100644 libs/tmp/files/v2/ext_upload_helpers.go delete mode 100644 libs/tmp/files/v2/ext_upload_test.go delete mode 100644 libs/tmp/files/v2/genhelper.go delete mode 100644 libs/tmp/files/v2/internal/cloudstorage/body.go delete mode 100644 libs/tmp/files/v2/internal/cloudstorage/body_test.go delete mode 100644 libs/tmp/files/v2/internal/cloudstorage/cloudstorage.go delete mode 100644 libs/tmp/files/v2/internal/cloudstorage/cloudstorage_test.go delete mode 100644 libs/tmp/files/v2/internal/cloudstorage/retries.go delete mode 100644 libs/tmp/files/v2/internal/cloudstorage/retries_test.go delete mode 100644 libs/tmp/files/v2/model.go delete mode 100644 libs/tmp/files/v2/wire.go delete mode 100644 libs/tmp/options/call/call.go delete mode 100644 libs/tmp/options/call/call_test.go delete mode 100644 libs/tmp/options/client/client.go delete mode 100644 libs/tmp/options/client/client_test.go delete mode 100644 libs/tmp/options/internal/version.go delete mode 100644 libs/tmp/options/internaloptions/internaloptions.go delete mode 100644 libs/tmp/options/internaloptions/internaloptions_test.go diff --git a/.github/OWNERS b/.github/OWNERS index 865d29ff3f4..76bc081f0eb 100644 --- a/.github/OWNERS +++ b/.github/OWNERS @@ -39,7 +39,6 @@ /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 e2f03cf9db3..565f46a5ca8 100644 --- a/.golangci.yaml +++ b/.golangci.yaml @@ -163,11 +163,6 @@ 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 @@ -204,8 +199,3 @@ 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/cmd/fs/upload_progress.go b/cmd/fs/upload_progress.go index 47a5461b1eb..a186fb5fa76 100644 --- a/cmd/fs/upload_progress.go +++ b/cmd/fs/upload_progress.go @@ -10,7 +10,7 @@ import ( "github.com/charmbracelet/lipgloss" "github.com/databricks/cli/libs/cmdio" - files "github.com/databricks/cli/libs/tmp/files/v2" + files "github.com/databricks/sdk-go/files/v2" ) // renderThrottle caps how often the interactive bar is redrawn so a fast upload diff --git a/go.mod b/go.mod index f2924fac769..305c226651e 100644 --- a/go.mod +++ b/go.mod @@ -45,7 +45,13 @@ 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 github.com/databricks/sdk-go/core v0.0.1-dev // Apache-2.0 + +require ( + github.com/databricks/sdk-go/auth v0.0.0-dev // Apache-2.0 + github.com/databricks/sdk-go/files v0.0.0-dev.1 // Apache-2.0 + github.com/databricks/sdk-go/options v0.0.0-dev // Apache-2.0 +) require ( cloud.google.com/go/auth v0.18.1 // indirect diff --git a/go.sum b/go.sum index a1994be2135..6a56d025ddf 100644 --- a/go.sum +++ b/go.sum @@ -71,8 +71,14 @@ github.com/danieljoos/wincred v1.2.3 h1:v7dZC2x32Ut3nEfRH+vhoZGvN72+dQ/snVXo/vMF github.com/danieljoos/wincred v1.2.3/go.mod h1:6qqX0WNrS4RzPZ1tnroDzq9kY3fu1KwE7MRLQK4X0bs= github.com/databricks/databricks-sdk-go v0.160.0 h1:vwgT/11y2vMw41BxcKbUUqarg45lmoEdukk9yYJg5AM= github.com/databricks/databricks-sdk-go v0.160.0/go.mod h1:C5LNgGe6hGuRrTwoxFmuup3XtQQEaqtq0e+K8IFDIS4= -github.com/databricks/sdk-go/core v0.0.0 h1:1Zg54LjihKnFTfcVwtEsbv/nm446QVOtenS8dSPjQ9U= -github.com/databricks/sdk-go/core v0.0.0/go.mod h1:7Ckau34bOsaZhHJE6r7ORNa+/kO33jIs6DfrJL6UbsM= +github.com/databricks/sdk-go/auth v0.0.0-dev h1:SyRGZvExH5TG9ynZEqNn+m+xTY42VOXr2GOxgNOQexk= +github.com/databricks/sdk-go/auth v0.0.0-dev/go.mod h1:Tj09W13MScUaix94cR0He9msAwe40JAQKuS2jP/ofQA= +github.com/databricks/sdk-go/core v0.0.1-dev h1:sIA1hCEJyl8/012vOBUIkDjsycNCD8RwfP76SCM0QbQ= +github.com/databricks/sdk-go/core v0.0.1-dev/go.mod h1:7Ckau34bOsaZhHJE6r7ORNa+/kO33jIs6DfrJL6UbsM= +github.com/databricks/sdk-go/files v0.0.0-dev.1 h1:MAcqxXpTFuM8dV+Yeb+X2+2iQlXWV2MVYEsYKle09+0= +github.com/databricks/sdk-go/files v0.0.0-dev.1/go.mod h1:4xc94QVa2BnkqR3eVC4o73phlwPgB/N2EAPKlK/2/zQ= +github.com/databricks/sdk-go/options v0.0.0-dev h1:+3bgKy5OH6G8VqzAox0eh0Ca28eQnoi9qDi5j/deiRY= +github.com/databricks/sdk-go/options v0.0.0-dev/go.mod h1:+lMasXZ/AfRAUYFNC8KFuT2vCJoF+TYNZ7XE0t26zfk= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= diff --git a/internal/build/notice_test.go b/internal/build/notice_test.go index 4adbc2670d0..ea71f1f77a8 100644 --- a/internal/build/notice_test.go +++ b/internal/build/notice_test.go @@ -23,6 +23,9 @@ var moduleToGitHub = map[string]string{ var noticeExclude = map[string]bool{ "github.com/databricks/databricks-sdk-go": true, "github.com/databricks/sdk-go/core": true, + "github.com/databricks/sdk-go/auth": true, + "github.com/databricks/sdk-go/files": true, + "github.com/databricks/sdk-go/options": true, } // Additional entries required in the NOTICE file that are not direct go.mod diff --git a/libs/filer/files_client.go b/libs/filer/files_client.go index cca9b7a47c6..4351a3023b9 100644 --- a/libs/filer/files_client.go +++ b/libs/filer/files_client.go @@ -13,17 +13,16 @@ import ( "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/config" "github.com/databricks/sdk-go/core/apierr" + files "github.com/databricks/sdk-go/files/v2" + "github.com/databricks/sdk-go/options/client" "golang.org/x/sync/errgroup" ) // cloudResponseHeaderTimeout bounds the wait for a cloud storage response header -// on a large-file part transfer. It matches the files/v2 engine's own default; -// see the transfer-client note in libs/tmp/files/v2. +// on a large-file part transfer. It matches the files/v2 engine's own default. const cloudResponseHeaderTimeout = 60 * time.Second // httpStatus returns the HTTP status code of err if it is (or wraps) an @@ -204,10 +203,9 @@ func newTransferClient(n int) *http.Client { return &http.Client{Transport: transport} } -// 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. +// newFilesAPIClient builds the files/v2 client from the CLI's already resolved +// config, reusing its auth instead of re-reading a profile. cfg is passed by +// pointer because config.Config embeds a sync.Mutex and must not be copied. func newFilesAPIClient(ctx context.Context, cfg *config.Config) (*files.Client, error) { copts := []client.Option{ client.WithHost(cfg.Host), diff --git a/libs/filer/files_client_auth.go b/libs/filer/files_client_auth.go index 5d0d30afae9..1d28d1cacdf 100644 --- a/libs/filer/files_client_auth.go +++ b/libs/filer/files_client_auth.go @@ -4,12 +4,12 @@ import ( "context" "net/http" - tmpauth "github.com/databricks/cli/libs/tmp/auth" "github.com/databricks/databricks-sdk-go/config" + sdkauth "github.com/databricks/sdk-go/auth" ) // configCredentials adapts the CLI's resolved SDK config into the credentials -// interface expected by the vendored files/v2 client. It signs a throwaway +// interface expected by the files/v2 client. It signs a throwaway // request with config.Authenticate and hands the resulting headers to the // client, so the files/v2 client reuses the exact auth the rest of the CLI // already resolved instead of re-reading a profile. @@ -21,7 +21,7 @@ func (c configCredentials) Name() string { return "databricks-cli" } -func (c configCredentials) AuthHeaders(ctx context.Context) ([]tmpauth.Header, error) { +func (c configCredentials) AuthHeaders(ctx context.Context) ([]sdkauth.Header, error) { // The URL is irrelevant: config.Authenticate only reads it to pick the auth // scheme, and every files/v2 request targets the same workspace host. req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.cfg.Host, nil) @@ -32,10 +32,10 @@ func (c configCredentials) AuthHeaders(ctx context.Context) ([]tmpauth.Header, e return nil, err } - headers := make([]tmpauth.Header, 0, len(req.Header)) + headers := make([]sdkauth.Header, 0, len(req.Header)) for key, values := range req.Header { for _, value := range values { - headers = append(headers, tmpauth.Header{Key: key, Value: value}) + headers = append(headers, sdkauth.Header{Key: key, Value: value}) } } return headers, nil diff --git a/libs/filer/files_client_test.go b/libs/filer/files_client_test.go index 7008e6858a9..484715bcd96 100644 --- a/libs/filer/files_client_test.go +++ b/libs/filer/files_client_test.go @@ -10,8 +10,8 @@ import ( "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" + files "github.com/databricks/sdk-go/files/v2" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) diff --git a/libs/tmp/README.md b/libs/tmp/README.md deleted file mode 100644 index 70337684528..00000000000 --- a/libs/tmp/README.md +++ /dev/null @@ -1,48 +0,0 @@ -# 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 deleted file mode 100644 index fac53e7db42..00000000000 --- a/libs/tmp/auth/auth.go +++ /dev/null @@ -1,98 +0,0 @@ -// 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 deleted file mode 100644 index 1294e404e44..00000000000 --- a/libs/tmp/auth/auth_test.go +++ /dev/null @@ -1,16 +0,0 @@ -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 deleted file mode 100644 index 96e69c1d133..00000000000 --- a/libs/tmp/auth/cache.go +++ /dev/null @@ -1,245 +0,0 @@ -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 deleted file mode 100644 index e23d847a9a2..00000000000 --- a/libs/tmp/auth/cache_test.go +++ /dev/null @@ -1,608 +0,0 @@ -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 deleted file mode 100644 index 8944d3aa6fa..00000000000 --- a/libs/tmp/auth/retrying.go +++ /dev/null @@ -1,66 +0,0 @@ -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 deleted file mode 100644 index 93209270fc3..00000000000 --- a/libs/tmp/auth/retrying_test.go +++ /dev/null @@ -1,150 +0,0 @@ -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 deleted file mode 100644 index 4b332087817..00000000000 --- a/libs/tmp/files/internal/version.go +++ /dev/null @@ -1,5 +0,0 @@ -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 deleted file mode 100644 index 01ac14d5ca8..00000000000 --- a/libs/tmp/files/v2/client.go +++ /dev/null @@ -1,1303 +0,0 @@ -// 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 deleted file mode 100644 index e485e617f10..00000000000 --- a/libs/tmp/files/v2/ext_limiter.go +++ /dev/null @@ -1,55 +0,0 @@ -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 deleted file mode 100644 index e695ffe239e..00000000000 --- a/libs/tmp/files/v2/ext_limiter_test.go +++ /dev/null @@ -1,187 +0,0 @@ -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 deleted file mode 100644 index a882d84e715..00000000000 --- a/libs/tmp/files/v2/ext_multipart.go +++ /dev/null @@ -1,450 +0,0 @@ -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 deleted file mode 100644 index 76a1a1d22d2..00000000000 --- a/libs/tmp/files/v2/ext_resumable.go +++ /dev/null @@ -1,187 +0,0 @@ -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 deleted file mode 100644 index bdf579ceb0d..00000000000 --- a/libs/tmp/files/v2/ext_resumable_test.go +++ /dev/null @@ -1,21 +0,0 @@ -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 deleted file mode 100644 index 5b60d19c4b7..00000000000 --- a/libs/tmp/files/v2/ext_straggler.go +++ /dev/null @@ -1,81 +0,0 @@ -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 deleted file mode 100644 index ff74e289c0f..00000000000 --- a/libs/tmp/files/v2/ext_straggler_test.go +++ /dev/null @@ -1,101 +0,0 @@ -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 deleted file mode 100644 index 8bbe75fa130..00000000000 --- a/libs/tmp/files/v2/ext_tunables.go +++ /dev/null @@ -1,171 +0,0 @@ -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 deleted file mode 100644 index d2715a7d8dd..00000000000 --- a/libs/tmp/files/v2/ext_upload.go +++ /dev/null @@ -1,568 +0,0 @@ -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 deleted file mode 100644 index 8beab493414..00000000000 --- a/libs/tmp/files/v2/ext_upload_control.go +++ /dev/null @@ -1,367 +0,0 @@ -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 deleted file mode 100644 index 4500f73c54f..00000000000 --- a/libs/tmp/files/v2/ext_upload_control_test.go +++ /dev/null @@ -1,206 +0,0 @@ -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 deleted file mode 100644 index 99592a3b8b2..00000000000 --- a/libs/tmp/files/v2/ext_upload_helpers.go +++ /dev/null @@ -1,116 +0,0 @@ -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 deleted file mode 100644 index dc07b0bc73d..00000000000 --- a/libs/tmp/files/v2/ext_upload_test.go +++ /dev/null @@ -1,726 +0,0 @@ -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 deleted file mode 100644 index c757bcd6c05..00000000000 --- a/libs/tmp/files/v2/genhelper.go +++ /dev/null @@ -1,250 +0,0 @@ -// 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 deleted file mode 100644 index 828bdf638df..00000000000 --- a/libs/tmp/files/v2/internal/cloudstorage/body.go +++ /dev/null @@ -1,55 +0,0 @@ -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 deleted file mode 100644 index f67d052a5aa..00000000000 --- a/libs/tmp/files/v2/internal/cloudstorage/body_test.go +++ /dev/null @@ -1,57 +0,0 @@ -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 deleted file mode 100644 index 5574f54ac7c..00000000000 --- a/libs/tmp/files/v2/internal/cloudstorage/cloudstorage.go +++ /dev/null @@ -1,320 +0,0 @@ -// 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 deleted file mode 100644 index d92ae2c22a4..00000000000 --- a/libs/tmp/files/v2/internal/cloudstorage/cloudstorage_test.go +++ /dev/null @@ -1,174 +0,0 @@ -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 deleted file mode 100644 index 1f6ce8c2c89..00000000000 --- a/libs/tmp/files/v2/internal/cloudstorage/retries.go +++ /dev/null @@ -1,76 +0,0 @@ -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 deleted file mode 100644 index 51f403bf799..00000000000 --- a/libs/tmp/files/v2/internal/cloudstorage/retries_test.go +++ /dev/null @@ -1,43 +0,0 @@ -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 deleted file mode 100644 index 4875cde1e4e..00000000000 --- a/libs/tmp/files/v2/model.go +++ /dev/null @@ -1,190 +0,0 @@ -// 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 deleted file mode 100644 index fcea4363e38..00000000000 --- a/libs/tmp/files/v2/wire.go +++ /dev/null @@ -1,854 +0,0 @@ -// 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 deleted file mode 100644 index 7248237ecfb..00000000000 --- a/libs/tmp/options/call/call.go +++ /dev/null @@ -1,53 +0,0 @@ -// 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 deleted file mode 100644 index 59aa59b0ce9..00000000000 --- a/libs/tmp/options/call/call_test.go +++ /dev/null @@ -1,58 +0,0 @@ -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 deleted file mode 100644 index 8d75a574bd0..00000000000 --- a/libs/tmp/options/client/client.go +++ /dev/null @@ -1,119 +0,0 @@ -// 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 deleted file mode 100644 index e1bdc315916..00000000000 --- a/libs/tmp/options/client/client_test.go +++ /dev/null @@ -1,57 +0,0 @@ -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 deleted file mode 100644 index ab946dc5c93..00000000000 --- a/libs/tmp/options/internal/version.go +++ /dev/null @@ -1,5 +0,0 @@ -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 deleted file mode 100644 index 1cf71bfc7a8..00000000000 --- a/libs/tmp/options/internaloptions/internaloptions.go +++ /dev/null @@ -1,113 +0,0 @@ -// 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 deleted file mode 100644 index 233d976b031..00000000000 --- a/libs/tmp/options/internaloptions/internaloptions_test.go +++ /dev/null @@ -1,41 +0,0 @@ -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 d3dd8dd43e0..5c37f1f4750 100755 --- a/tools/check_deadcode.py +++ b/tools/check_deadcode.py @@ -42,7 +42,6 @@ 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" From cb8b260fc2be20f6a885048b666d5b466c701e31 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Fri, 24 Jul 2026 12:42:52 +0000 Subject: [PATCH 3/4] filer: name files/v2 credentials by the config's resolved auth type The configCredentials adapter reported a hardcoded "databricks-cli" as its credentials name, which the files/v2 client folds into the User-Agent. Report the config's actual resolved auth type (pat, oauth-m2m, azure-cli, ...) instead. The CLI resolves auth when it builds the workspace client, so AuthType is set before the filer is created. This pull request and its description were written by Isaac. Co-authored-by: Isaac --- libs/filer/files_client_auth.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/libs/filer/files_client_auth.go b/libs/filer/files_client_auth.go index 1d28d1cacdf..9b0a96be818 100644 --- a/libs/filer/files_client_auth.go +++ b/libs/filer/files_client_auth.go @@ -17,8 +17,11 @@ type configCredentials struct { cfg *config.Config } +// Name reports the resolved auth mechanism (e.g. "pat", "oauth-m2m"), which the +// files/v2 client folds into its User-Agent. The CLI resolves auth when it +// builds the workspace client, so AuthType is set before the filer is created. func (c configCredentials) Name() string { - return "databricks-cli" + return c.cfg.AuthType } func (c configCredentials) AuthHeaders(ctx context.Context) ([]sdkauth.Header, error) { From 24741c42f8cb77b460c53637d066c0ae55a25900 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Fri, 24 Jul 2026 13:21:20 +0000 Subject: [PATCH 4/4] filer: fall back to "unk" auth type when config auth is unresolved The files/v2 client validates the credential name at construction and rejects an empty value with an "invalid value for \"auth\"" error. configCredentials.Name returned cfg.AuthType, which is empty when auth has not been resolved (as in the bundle/libraries UC Volumes filer tests that build a bare config), breaking task test across all platforms. Co-authored-by: Isaac --- libs/filer/files_client_auth.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/libs/filer/files_client_auth.go b/libs/filer/files_client_auth.go index 9b0a96be818..6ed9098532f 100644 --- a/libs/filer/files_client_auth.go +++ b/libs/filer/files_client_auth.go @@ -19,8 +19,13 @@ type configCredentials struct { // Name reports the resolved auth mechanism (e.g. "pat", "oauth-m2m"), which the // files/v2 client folds into its User-Agent. The CLI resolves auth when it -// builds the workspace client, so AuthType is set before the filer is created. +// builds the workspace client, so AuthType is normally set before the filer is +// created; fall back to "unk" when it is empty because the files/v2 client +// rejects an empty auth name with an "invalid value" error at construction. func (c configCredentials) Name() string { + if c.cfg.AuthType == "" { + return "unk" + } return c.cfg.AuthType }