From 2c6d8d4af53a7efa1b2a5ac6d1b15bf247936dd8 Mon Sep 17 00:00:00 2001 From: Abdurrahman Abdulaziz Date: Wed, 21 Oct 2026 17:08:30 +0200 Subject: [PATCH 1/3] ssh: create the release folder before uploading the tunnel binaries import-file reports a missing parent folder with a 404 only after the whole archive has been sent. On the first upload of every new version, CreateParentDirectories then created the folder and sent the archive again, so each ~17 MB archive was uploaded twice and half of the import-file calls failed. Create the release folder with mkdirs, which succeeds when it already exists, and upload once. Co-authored-by: Isaac --- experimental/ssh/internal/client/releases.go | 9 +++- .../ssh/internal/client/releases_test.go | 42 +++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/experimental/ssh/internal/client/releases.go b/experimental/ssh/internal/client/releases.go index 874a66bba2b..dd7972372a4 100644 --- a/experimental/ssh/internal/client/releases.go +++ b/experimental/ssh/internal/client/releases.go @@ -127,10 +127,17 @@ func uploadReleases(ctx context.Context, workspaceFiler filer.Filer, getRelease } defer releaseReader.Close() + // import-file reports a missing parent folder only after receiving the whole archive, so + // create the folder up front instead of letting CreateParentDirectories upload it twice. + err = workspaceFiler.Mkdir(ctx, remoteSubFolder) + if err != nil { + return fmt.Errorf("failed to create directory %s in workspace: %w", remoteSubFolder, err) + } + log.Infof(ctx, "Uploading %s to the workspace", fileName) // workspace-files/import-file API will automatically unzip the payload, // producing the filerRoot/remoteSubFolder/*archive-contents* structure, with 'databricks' binary inside. - err = workspaceFiler.Write(ctx, remoteArchivePath, releaseReader, filer.OverwriteIfExists, filer.CreateParentDirectories) + err = workspaceFiler.Write(ctx, remoteArchivePath, releaseReader, filer.OverwriteIfExists) if err != nil { if isProxyUploadError(err) { return fmt.Errorf("failed to upload file %s to workspace: %w\n\n"+ diff --git a/experimental/ssh/internal/client/releases_test.go b/experimental/ssh/internal/client/releases_test.go index 81984741834..bb543cced72 100644 --- a/experimental/ssh/internal/client/releases_test.go +++ b/experimental/ssh/internal/client/releases_test.go @@ -1,16 +1,21 @@ package client import ( + "cmp" "context" "errors" "fmt" "io" "net/http" + "path" "strings" + "sync" "syscall" "testing" "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/config" "github.com/stretchr/testify/assert" @@ -141,3 +146,40 @@ func TestUploadReleasesWithExistingBinary(t *testing.T) { }) } } + +func TestUploadReleasesSendsEachArchiveOnce(t *testing.T) { + const versionedDir = "/Workspace/Users/tester@databricks.test/.databricks/ssh-tunnel/1.12.0" + + server := testserver.New(t) + var mu sync.Mutex + var imports []string + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any { + overwrite := req.URL.Query().Get("overwrite") == "true" + resp := req.Workspace.WorkspaceFilesImportFile(req.Vars["path"], req.Body, overwrite) + mu.Lock() + defer mu.Unlock() + imports = append(imports, fmt.Sprintf("%s %d", path.Base(path.Dir(req.Vars["path"])), cmp.Or(resp.StatusCode, http.StatusOK))) + return resp + }) + testserver.AddDefaultHandlers(server) + client, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: server.URL, + Token: "testtoken", + }) + require.NoError(t, err) + workspaceFiler, err := filer.NewWorkspaceFilesClient(client, versionedDir) + require.NoError(t, err) + + getRelease := func(ctx context.Context, architecture, version, releasesDir string) (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader("archive")), nil + } + err = uploadReleases(t.Context(), workspaceFiler, getRelease, "1.12.0", "") + require.NoError(t, err) + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, []string{ + "databricks_cli_1.12.0_linux_amd64 200", + "databricks_cli_1.12.0_linux_arm64 200", + }, imports) +} From b288cfbaa5529e3fba2c90f6ccfcf7ae0ec92d5b Mon Sep 17 00:00:00 2001 From: Abdurrahman Abdulaziz Date: Wed, 21 Oct 2026 17:59:21 +0200 Subject: [PATCH 2/3] ssh: record the release uploads in an acceptance test - connect-serverless-feature-disabled now prints the import-file requests, showing one upload per architecture. - Keep CreateParentDirectories as the fallback for a 404 after the folder was created, and drop the comment. - The unit test reads status codes through the test server's ResponseCallback instead of copying the import-file handler. Co-authored-by: Isaac --- .../output.txt | 18 +++++++++++++++++- .../connect-serverless-feature-disabled/script | 3 ++- experimental/ssh/internal/client/releases.go | 4 +--- .../ssh/internal/client/releases_test.go | 15 +++++++-------- 4 files changed, 27 insertions(+), 13 deletions(-) diff --git a/acceptance/ssh/connect-serverless-feature-disabled/output.txt b/acceptance/ssh/connect-serverless-feature-disabled/output.txt index 678cdab2437..9957a1c7e18 100644 --- a/acceptance/ssh/connect-serverless-feature-disabled/output.txt +++ b/acceptance/ssh/connect-serverless-feature-disabled/output.txt @@ -8,7 +8,23 @@ >>> jq -r select(.path == "/telemetry-ext") | .body.protoLogs[] | fromjson | .entry.databricks_cli_log.ssh_tunnel_event.error_category [OUTPUT_DIR]/out.requests.txt DRIVER_PROXY_FORBIDDEN ->>> print_requests.py //api/2.2/jobs/runs/cancel +=== Each release archive is uploaded once, and the submitted run is cancelled + +>>> print_requests.py //api/2.0/workspace-files/import-file //api/2.2/jobs/runs/cancel +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.databricks/ssh-tunnel/[CLI_VERSION]/databricks_cli_linux_amd64/databricks.zip", + "q": { + "overwrite": "true" + } +} +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.databricks/ssh-tunnel/[CLI_VERSION]/databricks_cli_linux_arm64/databricks.zip", + "q": { + "overwrite": "true" + } +} { "method": "POST", "path": "/api/2.2/jobs/runs/cancel", diff --git a/acceptance/ssh/connect-serverless-feature-disabled/script b/acceptance/ssh/connect-serverless-feature-disabled/script index 68bc8b23528..9f8fe5c1b92 100644 --- a/acceptance/ssh/connect-serverless-feature-disabled/script +++ b/acceptance/ssh/connect-serverless-feature-disabled/script @@ -18,4 +18,5 @@ add_repl "$(grep -o 'run ID: [0-9]*' out.connect.txt | grep -o '[0-9]*$')" RUN_I title "Reported error category\n" trace jq -r 'select(.path == "/telemetry-ext") | .body.protoLogs[] | fromjson | .entry.databricks_cli_log.ssh_tunnel_event.error_category' "$OUT_REQUESTS" -trace print_requests.py //api/2.2/jobs/runs/cancel +title "Each release archive is uploaded once, and the submitted run is cancelled\n" +trace print_requests.py //api/2.0/workspace-files/import-file //api/2.2/jobs/runs/cancel diff --git a/experimental/ssh/internal/client/releases.go b/experimental/ssh/internal/client/releases.go index dd7972372a4..9a28c408a9f 100644 --- a/experimental/ssh/internal/client/releases.go +++ b/experimental/ssh/internal/client/releases.go @@ -127,8 +127,6 @@ func uploadReleases(ctx context.Context, workspaceFiler filer.Filer, getRelease } defer releaseReader.Close() - // import-file reports a missing parent folder only after receiving the whole archive, so - // create the folder up front instead of letting CreateParentDirectories upload it twice. err = workspaceFiler.Mkdir(ctx, remoteSubFolder) if err != nil { return fmt.Errorf("failed to create directory %s in workspace: %w", remoteSubFolder, err) @@ -137,7 +135,7 @@ func uploadReleases(ctx context.Context, workspaceFiler filer.Filer, getRelease log.Infof(ctx, "Uploading %s to the workspace", fileName) // workspace-files/import-file API will automatically unzip the payload, // producing the filerRoot/remoteSubFolder/*archive-contents* structure, with 'databricks' binary inside. - err = workspaceFiler.Write(ctx, remoteArchivePath, releaseReader, filer.OverwriteIfExists) + err = workspaceFiler.Write(ctx, remoteArchivePath, releaseReader, filer.OverwriteIfExists, filer.CreateParentDirectories) if err != nil { if isProxyUploadError(err) { return fmt.Errorf("failed to upload file %s to workspace: %w\n\n"+ diff --git a/experimental/ssh/internal/client/releases_test.go b/experimental/ssh/internal/client/releases_test.go index bb543cced72..95229ed422a 100644 --- a/experimental/ssh/internal/client/releases_test.go +++ b/experimental/ssh/internal/client/releases_test.go @@ -1,7 +1,6 @@ package client import ( - "cmp" "context" "errors" "fmt" @@ -151,17 +150,17 @@ func TestUploadReleasesSendsEachArchiveOnce(t *testing.T) { const versionedDir = "/Workspace/Users/tester@databricks.test/.databricks/ssh-tunnel/1.12.0" server := testserver.New(t) + testserver.AddDefaultHandlers(server) var mu sync.Mutex var imports []string - server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any { - overwrite := req.URL.Query().Get("overwrite") == "true" - resp := req.Workspace.WorkspaceFilesImportFile(req.Vars["path"], req.Body, overwrite) + server.ResponseCallback = func(req *testserver.Request, resp *testserver.EncodedResponse) { + if !strings.HasPrefix(req.URL.Path, "/api/2.0/workspace-files/import-file/") { + return + } mu.Lock() defer mu.Unlock() - imports = append(imports, fmt.Sprintf("%s %d", path.Base(path.Dir(req.Vars["path"])), cmp.Or(resp.StatusCode, http.StatusOK))) - return resp - }) - testserver.AddDefaultHandlers(server) + imports = append(imports, fmt.Sprintf("%s %d", path.Base(path.Dir(req.URL.Path)), resp.StatusCode)) + } client, err := databricks.NewWorkspaceClient(&databricks.Config{ Host: server.URL, Token: "testtoken", From 689064c8998bf9fd2d79d7c89d9558b564b1c028 Mon Sep 17 00:00:00 2001 From: Abdurrahman Abdulaziz Date: Wed, 21 Oct 2026 18:48:21 +0200 Subject: [PATCH 3/3] ssh: create the release folder before fetching the archive Run mkdirs before getting the release, so a failed mkdirs doesn't leave a started GitHub download open. Cover a dev re-upload into existing folders in the unit test, keeping each case's version in one place, and mention the upload check in the RecordRequests comment of connect-serverless-feature-disabled. Co-authored-by: Isaac --- .../test.toml | 3 +- experimental/ssh/internal/client/releases.go | 10 +- .../ssh/internal/client/releases_test.go | 91 ++++++++++++------- 3 files changed, 67 insertions(+), 37 deletions(-) diff --git a/acceptance/ssh/connect-serverless-feature-disabled/test.toml b/acceptance/ssh/connect-serverless-feature-disabled/test.toml index 36257965b75..3ed947e3902 100644 --- a/acceptance/ssh/connect-serverless-feature-disabled/test.toml +++ b/acceptance/ssh/connect-serverless-feature-disabled/test.toml @@ -1,6 +1,7 @@ Cloud = false -# Record requests so we can assert that the CLI cancels the run it submitted. +# Record requests so we can assert that the CLI uploads each release archive once and cancels +# the run it submitted. RecordRequests = true # Stand-in release archives staged by the script; not golden output. diff --git a/experimental/ssh/internal/client/releases.go b/experimental/ssh/internal/client/releases.go index 9a28c408a9f..def34ce363c 100644 --- a/experimental/ssh/internal/client/releases.go +++ b/experimental/ssh/internal/client/releases.go @@ -121,16 +121,16 @@ func uploadReleases(ctx context.Context, workspaceFiler filer.Filer, getRelease } } - releaseReader, err := getRelease(ctx, arch, version, releasesDir) + err := workspaceFiler.Mkdir(ctx, remoteSubFolder) if err != nil { - return fmt.Errorf("failed to get archive for architecture %s: %w", arch, err) + return fmt.Errorf("failed to create directory %s in workspace: %w", remoteSubFolder, err) } - defer releaseReader.Close() - err = workspaceFiler.Mkdir(ctx, remoteSubFolder) + releaseReader, err := getRelease(ctx, arch, version, releasesDir) if err != nil { - return fmt.Errorf("failed to create directory %s in workspace: %w", remoteSubFolder, err) + return fmt.Errorf("failed to get archive for architecture %s: %w", arch, err) } + defer releaseReader.Close() log.Infof(ctx, "Uploading %s to the workspace", fileName) // workspace-files/import-file API will automatically unzip the payload, diff --git a/experimental/ssh/internal/client/releases_test.go b/experimental/ssh/internal/client/releases_test.go index 95229ed422a..e0e2c384d9d 100644 --- a/experimental/ssh/internal/client/releases_test.go +++ b/experimental/ssh/internal/client/releases_test.go @@ -147,38 +147,67 @@ func TestUploadReleasesWithExistingBinary(t *testing.T) { } func TestUploadReleasesSendsEachArchiveOnce(t *testing.T) { - const versionedDir = "/Workspace/Users/tester@databricks.test/.databricks/ssh-tunnel/1.12.0" - - server := testserver.New(t) - testserver.AddDefaultHandlers(server) - var mu sync.Mutex - var imports []string - server.ResponseCallback = func(req *testserver.Request, resp *testserver.EncodedResponse) { - if !strings.HasPrefix(req.URL.Path, "/api/2.0/workspace-files/import-file/") { - return - } - mu.Lock() - defer mu.Unlock() - imports = append(imports, fmt.Sprintf("%s %d", path.Base(path.Dir(req.URL.Path)), resp.StatusCode)) + tests := []struct { + name string + version string + uploads int + want []string + }{ + { + name: "release version into new folders", + version: "1.12.0", + uploads: 1, + want: []string{ + "databricks_cli_1.12.0_linux_amd64 200", + "databricks_cli_1.12.0_linux_arm64 200", + }, + }, + { + name: "dev version again into existing folders", + version: "1.12.1-dev+abcdef123456", + uploads: 2, + want: []string{ + "databricks_cli_linux_amd64 200", + "databricks_cli_linux_arm64 200", + "databricks_cli_linux_amd64 200", + "databricks_cli_linux_arm64 200", + }, + }, } - client, err := databricks.NewWorkspaceClient(&databricks.Config{ - Host: server.URL, - Token: "testtoken", - }) - require.NoError(t, err) - workspaceFiler, err := filer.NewWorkspaceFilesClient(client, versionedDir) - require.NoError(t, err) - getRelease := func(ctx context.Context, architecture, version, releasesDir string) (io.ReadCloser, error) { - return io.NopCloser(strings.NewReader("archive")), nil - } - err = uploadReleases(t.Context(), workspaceFiler, getRelease, "1.12.0", "") - require.NoError(t, err) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := testserver.New(t) + testserver.AddDefaultHandlers(server) + var mu sync.Mutex + var imports []string + server.ResponseCallback = func(req *testserver.Request, resp *testserver.EncodedResponse) { + if !strings.HasPrefix(req.URL.Path, "/api/2.0/workspace-files/import-file/") { + return + } + mu.Lock() + defer mu.Unlock() + imports = append(imports, fmt.Sprintf("%s %d", path.Base(path.Dir(req.URL.Path)), resp.StatusCode)) + } + client, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: server.URL, + Token: "testtoken", + }) + require.NoError(t, err) + workspaceFiler, err := filer.NewWorkspaceFilesClient(client, "/Workspace/Users/tester@databricks.test/.databricks/ssh-tunnel/"+tt.version) + require.NoError(t, err) - mu.Lock() - defer mu.Unlock() - assert.Equal(t, []string{ - "databricks_cli_1.12.0_linux_amd64 200", - "databricks_cli_1.12.0_linux_arm64 200", - }, imports) + getRelease := func(ctx context.Context, architecture, version, releasesDir string) (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader("archive")), nil + } + for range tt.uploads { + err = uploadReleases(t.Context(), workspaceFiler, getRelease, tt.version, "") + require.NoError(t, err) + } + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, tt.want, imports) + }) + } }