From 978544a89557f4963662327db797c2179c55072a Mon Sep 17 00:00:00 2001
From: simonfaltum <simon.faltum@databricks.com>
Date: Tue, 9 Jun 2026 23:05:52 +0200
Subject: [PATCH 1/2] 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 <simon.faltum@databricks.com>
Date: Wed, 10 Jun 2026 07:31:06 +0200
Subject: [PATCH 2/2] 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