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/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 874a66bba2b..def34ce363c 100644 --- a/experimental/ssh/internal/client/releases.go +++ b/experimental/ssh/internal/client/releases.go @@ -121,6 +121,11 @@ func uploadReleases(ctx context.Context, workspaceFiler filer.Filer, getRelease } } + err := workspaceFiler.Mkdir(ctx, remoteSubFolder) + if err != nil { + return fmt.Errorf("failed to create directory %s in workspace: %w", remoteSubFolder, err) + } + releaseReader, err := getRelease(ctx, arch, version, releasesDir) if err != nil { return fmt.Errorf("failed to get archive for architecture %s: %w", arch, err) diff --git a/experimental/ssh/internal/client/releases_test.go b/experimental/ssh/internal/client/releases_test.go index 81984741834..e0e2c384d9d 100644 --- a/experimental/ssh/internal/client/releases_test.go +++ b/experimental/ssh/internal/client/releases_test.go @@ -6,11 +6,15 @@ import ( "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 +145,69 @@ func TestUploadReleasesWithExistingBinary(t *testing.T) { }) } } + +func TestUploadReleasesSendsEachArchiveOnce(t *testing.T) { + 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", + }, + }, + } + + 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) + + 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) + }) + } +}