From 978544a89557f4963662327db797c2179c55072a Mon Sep 17 00:00:00 2001 From: simonfaltum Date: Tue, 9 Jun 2026 23:05:52 +0200 Subject: [PATCH 1/3] Check Close error when generate downloader writes files FlushToDisk closed downloaded files via defer, discarding the Close error and logging success before the file was closed. A failed flush could leave a truncated file while the command exits 0. Close the file explicitly and propagate the error, mirroring the pattern in libs/filer/local_client.go. Co-authored-by: Isaac --- bundle/generate/downloader.go | 7 ++++- bundle/generate/downloader_test.go | 45 ++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/bundle/generate/downloader.go b/bundle/generate/downloader.go index acfef1ff3db..672f20a52d8 100644 --- a/bundle/generate/downloader.go +++ b/bundle/generate/downloader.go @@ -321,9 +321,14 @@ func (n *Downloader) FlushToDisk(ctx context.Context, force bool) error { if err != nil { return err } - defer file.Close() _, err = io.Copy(file, reader) + // Close flushes buffered writes; if it fails the file may be truncated, + // so check it before reporting success. + cerr := file.Close() + if err == nil { + err = cerr + } if err != nil { return err } diff --git a/bundle/generate/downloader_test.go b/bundle/generate/downloader_test.go index 9363ce98d3c..d5c20168af7 100644 --- a/bundle/generate/downloader_test.go +++ b/bundle/generate/downloader_test.go @@ -8,6 +8,7 @@ import ( "path/filepath" "testing" + "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/experimental/mocks" "github.com/databricks/databricks-sdk-go/service/jobs" @@ -261,6 +262,50 @@ func TestDownloader_MarkTasksForDownload_NoNotebooks(t *testing.T) { assert.Empty(t, downloader.files) } +func TestDownloader_FlushToDisk(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + + contents := map[string]string{ + "/Users/user/project/notebook": "# Databricks notebook source\nprint(1)", + "/Users/user/project/utils.py": "def helper(): pass", + } + w := newTestWorkspaceClient(t, func(rw http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/2.0/workspace/export" { + t.Fatalf("unexpected request path: %s", r.URL.Path) + } + content, ok := contents[r.URL.Query().Get("path")] + if !ok { + t.Fatalf("unexpected export path: %s", r.URL.Query().Get("path")) + } + _, err := rw.Write([]byte(content)) + if err != nil { + t.Fatal(err) + } + }) + + sourceDir := t.TempDir() + downloader := NewDownloader(w, sourceDir, "config") + downloader.files[filepath.Join(sourceDir, "notebook.py")] = exportFile{ + path: "/Users/user/project/notebook", + format: workspace.ExportFormatSource, + } + downloader.files[filepath.Join(sourceDir, "utils.py")] = exportFile{ + path: "/Users/user/project/utils.py", + format: workspace.ExportFormatSource, + } + + err := downloader.FlushToDisk(ctx, false) + require.NoError(t, err) + + data, err := os.ReadFile(filepath.Join(sourceDir, "notebook.py")) + require.NoError(t, err) + assert.Equal(t, contents["/Users/user/project/notebook"], string(data)) + + data, err = os.ReadFile(filepath.Join(sourceDir, "utils.py")) + require.NoError(t, err) + assert.Equal(t, contents["/Users/user/project/utils.py"], string(data)) +} + func TestDownloader_CleanupOldFiles(t *testing.T) { ctx := t.Context() sourceDir := t.TempDir() From d34c4ee535c99877d77ed86f58399c517dac79e1 Mon Sep 17 00:00:00 2001 From: simonfaltum Date: Wed, 10 Jun 2026 07:31:06 +0200 Subject: [PATCH 2/3] Remove redundant comments Co-authored-by: Isaac --- bundle/generate/downloader.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/bundle/generate/downloader.go b/bundle/generate/downloader.go index 672f20a52d8..c86a1b44fcd 100644 --- a/bundle/generate/downloader.go +++ b/bundle/generate/downloader.go @@ -323,8 +323,7 @@ func (n *Downloader) FlushToDisk(ctx context.Context, force bool) error { } _, err = io.Copy(file, reader) - // Close flushes buffered writes; if it fails the file may be truncated, - // so check it before reporting success. + // Close flushes buffered writes; if it fails the file may be truncated. cerr := file.Close() if err == nil { err = cerr From 246ff61f56c58dc76870e4035b62801f084d8e00 Mon Sep 17 00:00:00 2001 From: simonfaltum Date: Wed, 10 Jun 2026 18:24:05 +0200 Subject: [PATCH 3/3] Extract writeAndClose and test the failure path A real os.File Close error cannot be triggered portably in a test, so extract the copy-and-close logic into a helper and cover the failure paths with a fake WriteCloser: a Close error is returned, and a copy error is not masked by a Close error. Co-authored-by: Isaac --- bundle/generate/downloader.go | 19 ++++++---- bundle/generate/downloader_test.go | 56 ++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 6 deletions(-) diff --git a/bundle/generate/downloader.go b/bundle/generate/downloader.go index c86a1b44fcd..bf12b375866 100644 --- a/bundle/generate/downloader.go +++ b/bundle/generate/downloader.go @@ -322,12 +322,7 @@ func (n *Downloader) FlushToDisk(ctx context.Context, force bool) error { return err } - _, err = io.Copy(file, reader) - // Close flushes buffered writes; if it fails the file may be truncated. - cerr := file.Close() - if err == nil { - err = cerr - } + err = writeAndClose(file, reader) if err != nil { return err } @@ -340,6 +335,18 @@ func (n *Downloader) FlushToDisk(ctx context.Context, force bool) error { return errs.Wait() } +// writeAndClose copies src into dst and closes dst. A copy error takes +// precedence; otherwise the Close error is returned because a failed Close +// can mean buffered writes were lost and the file is truncated. +func writeAndClose(dst io.WriteCloser, src io.Reader) error { + _, err := io.Copy(dst, src) + cerr := dst.Close() + if err == nil { + err = cerr + } + return err +} + func NewDownloader(w *databricks.WorkspaceClient, sourceDir, configDir string) *Downloader { return &Downloader{ files: make(map[string]exportFile), diff --git a/bundle/generate/downloader_test.go b/bundle/generate/downloader_test.go index d5c20168af7..5ce4889c5a2 100644 --- a/bundle/generate/downloader_test.go +++ b/bundle/generate/downloader_test.go @@ -1,12 +1,17 @@ package generate import ( + "bytes" "encoding/json" + "errors" + "io" "net/http" "net/http/httptest" "os" "path/filepath" + "strings" "testing" + "testing/iotest" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go" @@ -306,6 +311,57 @@ func TestDownloader_FlushToDisk(t *testing.T) { assert.Equal(t, contents["/Users/user/project/utils.py"], string(data)) } +type fakeWriteCloser struct { + bytes.Buffer + closeErr error +} + +func (f *fakeWriteCloser) Close() error { return f.closeErr } + +func TestWriteAndClose(t *testing.T) { + closeErr := errors.New("close failed") + readErr := errors.New("read failed") + + tests := []struct { + name string + src io.Reader + closeErr error + wantErr error + wantData string + }{ + { + name: "success", + src: strings.NewReader("data"), + wantData: "data", + }, + { + name: "close error is returned", + src: strings.NewReader("data"), + closeErr: closeErr, + wantErr: closeErr, + wantData: "data", + }, + { + name: "copy error takes precedence over close error", + src: iotest.ErrReader(readErr), + closeErr: closeErr, + wantErr: readErr, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dst := &fakeWriteCloser{closeErr: tt.closeErr} + err := writeAndClose(dst, tt.src) + if tt.wantErr == nil { + assert.NoError(t, err) + } else { + assert.ErrorIs(t, err, tt.wantErr) + } + assert.Equal(t, tt.wantData, dst.String()) + }) + } +} + func TestDownloader_CleanupOldFiles(t *testing.T) { ctx := t.Context() sourceDir := t.TempDir()