From bf2fc1fc7316bc842e5da756d33d0efc34cc93f3 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 14 Apr 2025 15:34:30 +0200 Subject: [PATCH 01/57] acc: Refactor starting test server in a separate file --- acceptance/acceptance_test.go | 157 ++----------- acceptance/internal/handlers.go | 261 ++++++++++++++++++++++ acceptance/internal/server.go | 377 +++++++++++++------------------- 3 files changed, 426 insertions(+), 369 deletions(-) create mode 100644 acceptance/internal/handlers.go diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 917561e16a..cfea953de7 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -26,12 +26,8 @@ import ( "github.com/databricks/cli/acceptance/internal" "github.com/databricks/cli/internal/testutil" - "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testdiff" - "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/iam" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -144,13 +140,7 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int { cloudEnv := os.Getenv("CLOUD_ENV") if cloudEnv == "" { - defaultServer := testserver.New(t) - internal.AddHandlers(defaultServer) - t.Setenv("DATABRICKS_DEFAULT_HOST", defaultServer.URL) - - homeDir := t.TempDir() - // Do not read user's ~/.databrickscfg - t.Setenv(env.HomeEnvVar(), homeDir) + internal.StartDefaultServer(t) } terraformrcPath := filepath.Join(buildDir, ".terraformrc") @@ -371,96 +361,26 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont cmd.Env = append(cmd.Env, "UNIQUE_NAME="+uniqueName) cmd.Env = append(cmd.Env, "TEST_TMP_DIR="+tmpDir) - var workspaceClient *databricks.WorkspaceClient - var user iam.User - - // Start a new server with a custom configuration if the acceptance test - // specifies a custom server stubs. - var server *testserver.Server - - if !isRunningOnCloud { - // Start a new server for this test if either: - // 1. A custom server spec is defined in the test configuration. - // 2. The test is configured to record requests and assert on them. We need - // a duplicate of the default server to record requests because the default - // server otherwise is a shared resource. - - databricksLocalHost := os.Getenv("DATABRICKS_DEFAULT_HOST") - - if len(config.Server) > 0 || isTruePtr(config.RecordRequests) { - server = testserver.New(t) - if isTruePtr(config.RecordRequests) { - requestsPath := filepath.Join(tmpDir, "out.requests.txt") - server.RequestCallback = func(request *testserver.Request) { - req := getLoggedRequest(request, config.IncludeRequestHeaders) - reqJson, err := json.MarshalIndent(req, "", " ") - assert.NoErrorf(t, err, "Failed to json-encode: %#v", req) - - f, err := os.OpenFile(requestsPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) - assert.NoError(t, err) - defer f.Close() - - _, err = f.WriteString(string(reqJson) + "\n") - assert.NoError(t, err) - } - } - - if LogRequests { - server.ResponseCallback = func(request *testserver.Request, response *testserver.EncodedResponse) { - t.Logf("%d %s %s\n%s\n%s", - response.StatusCode, request.Method, request.URL, - formatHeadersAndBody("> ", request.Headers, request.Body), - formatHeadersAndBody("# ", response.Headers, response.Body), - ) - } - } - - // We want later stubs takes precedence, because then leaf configs take precedence over parent directory configs - // In gorilla/mux earlier handlers take precedence, so we need to reverse the order - slices.Reverse(config.Server) - - for _, stub := range config.Server { - require.NotEmpty(t, stub.Pattern) - items := strings.Split(stub.Pattern, " ") - require.Len(t, items, 2) - server.Handle(items[0], items[1], func(req testserver.Request) any { - time.Sleep(stub.Delay) - return stub.Response - }) - } - - // The earliest handlers take precedence, add default handlers last - internal.AddHandlers(server) - databricksLocalHost = server.URL - } - - // Each local test should use a new token that will result into a new fake workspace, - // so that test don't interfere with each other. - tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") - config := databricks.Config{ - Host: databricksLocalHost, - Token: "dbapi" + tokenSuffix, - } - workspaceClient, err = databricks.NewWorkspaceClient(&config) - require.NoError(t, err) + workspaceClient := internal.ResolveServer(t, config, LogRequests, tmpDir) + testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient) - cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+config.Host) - cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+config.Token) + // Configure resolved credentials in the environment. + cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+workspaceClient.Config.Host) + if workspaceClient.Config.Token != "" { + cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+workspaceClient.Config.Token) + } - // For the purposes of replacements, use testUser. - // Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement: - user = internal.TestUser - } else { - // Use whatever authentication mechanism is configured by the test runner. - workspaceClient, err = databricks.NewWorkspaceClient(&databricks.Config{}) - require.NoError(t, err) + var user iam.User + if isRunningOnCloud { pUser, err := workspaceClient.CurrentUser.Me(context.Background()) require.NoError(t, err, "Failed to get current user") user = *pUser + } else { + // For the purposes of replacements, use testUser for local runs. + // Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement: + user = internal.TestUser } - testdiff.PrepareReplacementsUser(t, &repls, user) - testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient) // Must be added PrepareReplacementsUser, otherwise conflicts with [USERNAME] testdiff.PrepareReplacementsUUID(t, &repls) @@ -835,33 +755,6 @@ type LoggedRequest struct { RawBody string `json:"raw_body,omitempty"` } -func getLoggedRequest(req *testserver.Request, includedHeaders []string) LoggedRequest { - result := LoggedRequest{ - Method: req.Method, - Path: req.URL.Path, - Headers: filterHeaders(req.Headers, includedHeaders), - } - - if json.Valid(req.Body) { - result.Body = json.RawMessage(req.Body) - } else { - result.RawBody = string(req.Body) - } - - return result -} - -func filterHeaders(h http.Header, includedHeaders []string) http.Header { - headers := make(http.Header) - for k, v := range h { - if !slices.Contains(includedHeaders, k) { - continue - } - headers[k] = v - } - return headers -} - func isTruePtr(value *bool) bool { return value != nil && *value } @@ -968,25 +861,3 @@ func buildDatabricksBundlesWheel(t *testing.T, buildDir string) (string, error) return "", fmt.Errorf("databricks-bundles wheel not found in %s", buildDir) } } - -func formatHeadersAndBody(prefix string, headers http.Header, body []byte) string { - var result []string - for key, values := range headers { - if len(values) == 1 { - result = append(result, fmt.Sprintf("%s%s: %s", prefix, key, values[0])) - } else { - result = append(result, fmt.Sprintf("%s%s: %s", prefix, key, values)) - } - } - if len(body) > 0 { - var s string - if utf8.Valid(body) { - s = string(body) - } else { - s = fmt.Sprintf("[Binary %d bytes]", len(body)) - } - s = strings.ReplaceAll(s, "\n", "\n"+prefix) - result = append(result, prefix+s) - } - return strings.Join(result, "\n") -} diff --git a/acceptance/internal/handlers.go b/acceptance/internal/handlers.go new file mode 100644 index 0000000000..4f1665b027 --- /dev/null +++ b/acceptance/internal/handlers.go @@ -0,0 +1,261 @@ +package internal + +import ( + "encoding/json" + "fmt" + "net/http" + + "github.com/databricks/databricks-sdk-go/service/catalog" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/pipelines" + + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/jobs" + + "github.com/databricks/cli/libs/telemetry" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go/service/workspace" +) + +var TestUser = iam.User{ + Id: "1000012345", + UserName: "tester@databricks.com", +} + +var TestMetastore = catalog.MetastoreAssignment{ + DefaultCatalogName: "hive_metastore", + MetastoreId: "120efa64-9b68-46ba-be38-f319458430d2", + WorkspaceId: 470123456789500, +} + +func addDefaultHandlers(server *testserver.Server) { + server.Handle("GET", "/api/2.0/policies/clusters/list", func(req testserver.Request) any { + return compute.ListPoliciesResponse{ + Policies: []compute.Policy{ + { + PolicyId: "5678", + Name: "wrong-cluster-policy", + }, + { + PolicyId: "9876", + Name: "some-test-cluster-policy", + }, + }, + } + }) + + server.Handle("GET", "/api/2.0/instance-pools/list", func(req testserver.Request) any { + return compute.ListInstancePools{ + InstancePools: []compute.InstancePoolAndStats{ + { + InstancePoolName: "some-test-instance-pool", + InstancePoolId: "1234", + }, + }, + } + }) + + server.Handle("GET", "/api/2.1/clusters/list", func(req testserver.Request) any { + return compute.ListClustersResponse{ + Clusters: []compute.ClusterDetails{ + { + ClusterName: "some-test-cluster", + ClusterId: "4321", + }, + { + ClusterName: "some-other-cluster", + ClusterId: "9876", + }, + }, + } + }) + + server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(req testserver.Request) any { + return testserver.Response{ + Headers: map[string][]string{"X-Databricks-Org-Id": {"900800700600"}}, + Body: TestUser, + } + }) + + server.Handle("GET", "/api/2.0/workspace/get-status", func(req testserver.Request) any { + path := req.URL.Query().Get("path") + return req.Workspace.WorkspaceGetStatus(path) + }) + + server.Handle("POST", "/api/2.0/workspace/mkdirs", func(req testserver.Request) any { + var request workspace.Mkdirs + if err := json.Unmarshal(req.Body, &request); err != nil { + return testserver.Response{ + Body: fmt.Sprintf("internal error: %s", err), + StatusCode: http.StatusInternalServerError, + } + } + + req.Workspace.WorkspaceMkdirs(request) + return "" + }) + + server.Handle("GET", "/api/2.0/workspace/export", func(req testserver.Request) any { + path := req.URL.Query().Get("path") + return req.Workspace.WorkspaceExport(path) + }) + + server.Handle("POST", "/api/2.0/workspace/delete", func(req testserver.Request) any { + path := req.URL.Query().Get("path") + recursive := req.URL.Query().Get("recursive") == "true" + req.Workspace.WorkspaceDelete(path, recursive) + return "" + }) + + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path:.*}", func(req testserver.Request) any { + path := req.Vars["path"] + req.Workspace.WorkspaceFilesImportFile(path, req.Body) + return "" + }) + + server.Handle("GET", "/api/2.0/workspace-files/{path:.*}", func(req testserver.Request) any { + path := req.Vars["path"] + return req.Workspace.WorkspaceFilesExportFile(path) + }) + + server.Handle("GET", "/api/2.1/unity-catalog/current-metastore-assignment", func(req testserver.Request) any { + return TestMetastore + }) + + server.Handle("GET", "/api/2.0/permissions/directories/{objectId}", func(req testserver.Request) any { + objectId := req.Vars["objectId"] + return workspace.WorkspaceObjectPermissions{ + ObjectId: objectId, + ObjectType: "DIRECTORY", + AccessControlList: []workspace.WorkspaceObjectAccessControlResponse{ + { + UserName: "tester@databricks.com", + AllPermissions: []workspace.WorkspaceObjectPermission{ + { + PermissionLevel: "CAN_MANAGE", + }, + }, + }, + }, + } + }) + + server.Handle("POST", "/api/2.2/jobs/create", func(req testserver.Request) any { + var request jobs.CreateJob + if err := json.Unmarshal(req.Body, &request); err != nil { + return testserver.Response{ + Body: fmt.Sprintf("internal error: %s", err), + StatusCode: 500, + } + } + + return req.Workspace.JobsCreate(request) + }) + + server.Handle("POST", "/api/2.0/pipelines", func(req testserver.Request) any { + var request pipelines.PipelineSpec + if err := json.Unmarshal(req.Body, &request); err != nil { + return testserver.Response{ + Body: fmt.Sprintf("internal error: %s", err), + StatusCode: 400, + } + } + + return req.Workspace.PipelinesCreate(request) + }) + + server.Handle("GET", "/api/2.2/jobs/get", func(req testserver.Request) any { + jobId := req.URL.Query().Get("job_id") + return req.Workspace.JobsGet(jobId) + }) + + server.Handle("GET", "/api/2.0/pipelines/{pipeline_id}", func(req testserver.Request) any { + pipelineId := req.Vars["pipeline_id"] + return req.Workspace.PipelinesGet(pipelineId) + }) + + server.Handle("GET", "/api/2.2/jobs/list", func(req testserver.Request) any { + return req.Workspace.JobsList() + }) + + server.Handle("GET", "/api/2.2/jobs/list", func(req testserver.Request) any { + return req.Workspace.JobsList() + }) + + server.Handle("GET", "/oidc/.well-known/oauth-authorization-server", func(_ testserver.Request) any { + return map[string]string{ + "authorization_endpoint": server.URL + "oidc/v1/authorize", + "token_endpoint": server.URL + "/oidc/v1/token", + } + }) + + server.Handle("POST", "/oidc/v1/token", func(_ testserver.Request) any { + return map[string]string{ + "access_token": "oauth-token", + "expires_in": "3600", + "scope": "all-apis", + "token_type": "Bearer", + } + }) + + server.Handle("POST", "/telemetry-ext", func(_ testserver.Request) any { + return telemetry.ResponseBody{ + Errors: []telemetry.LogError{}, + NumProtoSuccess: 1, + } + }) + + // Quality monitors: + + server.Handle("GET", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { + return testserver.MapGet(req.Workspace.Monitors, req.Vars["table_name"]) + }) + + server.Handle("POST", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { + return req.Workspace.QualityMonitorUpsert(req, req.Vars["table_name"], false) + }) + + server.Handle("PUT", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { + return req.Workspace.QualityMonitorUpsert(req, req.Vars["table_name"], true) + }) + + server.Handle("DELETE", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { + return testserver.MapDelete(req.Workspace.Monitors, req.Vars["table_name"]) + }) + + // Apps: + + server.Handle("GET", "/api/2.0/apps/{name}", func(req testserver.Request) any { + return testserver.MapGet(req.Workspace.Apps, req.Vars["name"]) + }) + + server.Handle("POST", "/api/2.0/apps", func(req testserver.Request) any { + return req.Workspace.AppsUpsert(req, "") + }) + + server.Handle("PATCH", "/api/2.0/apps/{name}", func(req testserver.Request) any { + return req.Workspace.AppsUpsert(req, req.Vars["name"]) + }) + + server.Handle("DELETE", "/api/2.0/apps/{name}", func(req testserver.Request) any { + return testserver.MapDelete(req.Workspace.Apps, req.Vars["name"]) + }) + + // Schemas: + + server.Handle("GET", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { + return testserver.MapGet(req.Workspace.Schemas, req.Vars["full_name"]) + }) + + server.Handle("POST", "/api/2.1/unity-catalog/schemas", func(req testserver.Request) any { + return req.Workspace.SchemasCreate(req) + }) + + server.Handle("PATCH", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { + return req.Workspace.SchemasUpdate(req, req.Vars["full_name"]) + }) + + server.Handle("DELETE", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { + return testserver.MapDelete(req.Workspace.Schemas, req.Vars["full_name"]) + }) +} diff --git a/acceptance/internal/server.go b/acceptance/internal/server.go index 3156eeca74..5d7dfcc443 100644 --- a/acceptance/internal/server.go +++ b/acceptance/internal/server.go @@ -4,258 +4,183 @@ import ( "encoding/json" "fmt" "net/http" - - "github.com/databricks/databricks-sdk-go/service/catalog" - "github.com/databricks/databricks-sdk-go/service/iam" - "github.com/databricks/databricks-sdk-go/service/pipelines" - - "github.com/databricks/databricks-sdk-go/service/compute" - "github.com/databricks/databricks-sdk-go/service/jobs" - - "github.com/databricks/cli/libs/telemetry" + "os" + "path/filepath" + "slices" + "strings" + "testing" + "time" + "unicode/utf8" + + "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go/service/workspace" + "github.com/databricks/databricks-sdk-go" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) -var TestUser = iam.User{ - Id: "1000012345", - UserName: "tester@databricks.com", -} - -var TestMetastore = catalog.MetastoreAssignment{ - DefaultCatalogName: "hive_metastore", - MetastoreId: "120efa64-9b68-46ba-be38-f319458430d2", - WorkspaceId: 470123456789500, -} - -func AddHandlers(server *testserver.Server) { - server.Handle("GET", "/api/2.0/policies/clusters/list", func(req testserver.Request) any { - return compute.ListPoliciesResponse{ - Policies: []compute.Policy{ - { - PolicyId: "5678", - Name: "wrong-cluster-policy", - }, - { - PolicyId: "9876", - Name: "some-test-cluster-policy", - }, - }, - } - }) +func StartDefaultServer(t *testing.T) { + s := testserver.New(t) + addDefaultHandlers(s) - server.Handle("GET", "/api/2.0/instance-pools/list", func(req testserver.Request) any { - return compute.ListInstancePools{ - InstancePools: []compute.InstancePoolAndStats{ - { - InstancePoolName: "some-test-instance-pool", - InstancePoolId: "1234", - }, - }, - } - }) + t.Setenv("DATABRICKS_DEFAULT_HOST", s.URL) - server.Handle("GET", "/api/2.1/clusters/list", func(req testserver.Request) any { - return compute.ListClustersResponse{ - Clusters: []compute.ClusterDetails{ - { - ClusterName: "some-test-cluster", - ClusterId: "4321", - }, - { - ClusterName: "some-other-cluster", - ClusterId: "9876", - }, - }, - } - }) + // Do not read user's ~/.databrickscfg + homeDir := t.TempDir() + t.Setenv(env.HomeEnvVar(), homeDir) +} - server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(req testserver.Request) any { - return testserver.Response{ - Headers: map[string][]string{"X-Databricks-Org-Id": {"900800700600"}}, - Body: TestUser, - } - }) +func isTruePtr(value *bool) bool { + return value != nil && *value +} - server.Handle("GET", "/api/2.0/workspace/get-status", func(req testserver.Request) any { - path := req.URL.Query().Get("path") - return req.Workspace.WorkspaceGetStatus(path) - }) +func ResolveServer(t *testing.T, config TestConfig, logRequests bool, outputDir string) *databricks.WorkspaceClient { + cloudEnv := os.Getenv("CLOUD_ENV") - server.Handle("POST", "/api/2.0/workspace/mkdirs", func(req testserver.Request) any { - var request workspace.Mkdirs - if err := json.Unmarshal(req.Body, &request); err != nil { - return testserver.Response{ - Body: fmt.Sprintf("internal error: %s", err), - StatusCode: http.StatusInternalServerError, - } - } + // If we are running on a cloud environment, use the host configured in the + // environment. + if cloudEnv != "" { + w, err := databricks.NewWorkspaceClient(&databricks.Config{}) + require.NoError(t, err) - req.Workspace.WorkspaceMkdirs(request) - return "" - }) + return w + } - server.Handle("GET", "/api/2.0/workspace/export", func(req testserver.Request) any { - path := req.URL.Query().Get("path") - return req.Workspace.WorkspaceExport(path) - }) + recordRequests := isTruePtr(config.RecordRequests) - server.Handle("POST", "/api/2.0/workspace/delete", func(req testserver.Request) any { - path := req.URL.Query().Get("path") - recursive := req.URL.Query().Get("recursive") == "true" - req.Workspace.WorkspaceDelete(path, recursive) - return "" - }) + tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") + token := fmt.Sprintf("dbapi%s", tokenSuffix) - server.Handle("POST", "/api/2.0/workspace-files/import-file/{path:.*}", func(req testserver.Request) any { - path := req.Vars["path"] - req.Workspace.WorkspaceFilesImportFile(path, req.Body) - return "" - }) + // If we are not recording requests, and no custom server server stubs are configured, + // use the default shared server. + if len(config.Server) == 0 && !recordRequests { + w, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: os.Getenv("DATABRICKS_DEFAULT_HOST"), + Token: token, + }) + require.NoError(t, err) - server.Handle("GET", "/api/2.0/workspace-files/{path:.*}", func(req testserver.Request) any { - path := req.Vars["path"] - return req.Workspace.WorkspaceFilesExportFile(path) - }) + return w + } - server.Handle("GET", "/api/2.1/unity-catalog/current-metastore-assignment", func(req testserver.Request) any { - return TestMetastore - }) + // We want later stubs takes precedence, because then leaf configs take precedence over parent directory configs + // In gorilla/mux earlier handlers take precedence, so we need to reverse the order + slices.Reverse(config.Server) + host := startDedicatedServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) - server.Handle("GET", "/api/2.0/permissions/directories/{objectId}", func(req testserver.Request) any { - objectId := req.Vars["objectId"] - return workspace.WorkspaceObjectPermissions{ - ObjectId: objectId, - ObjectType: "DIRECTORY", - AccessControlList: []workspace.WorkspaceObjectAccessControlResponse{ - { - UserName: "tester@databricks.com", - AllPermissions: []workspace.WorkspaceObjectPermission{ - { - PermissionLevel: "CAN_MANAGE", - }, - }, - }, - }, - } + w, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: host, + Token: token, }) + require.NoError(t, err) + return w +} - server.Handle("POST", "/api/2.2/jobs/create", func(req testserver.Request) any { - var request jobs.CreateJob - if err := json.Unmarshal(req.Body, &request); err != nil { - return testserver.Response{ - Body: fmt.Sprintf("internal error: %s", err), - StatusCode: 500, - } +func startDedicatedServer(t *testing.T, + stubs []ServerStub, + recordRequests bool, + logRequests bool, + includedHeaders []string, + outputDir string, +) string { + s := testserver.New(t) + + if recordRequests { + requestsPath := filepath.Join(outputDir, "out.requests.txt") + s.RequestCallback = func(request *testserver.Request) { + req := getLoggedRequest(request, includedHeaders) + reqJson, err := json.MarshalIndent(req, "", " ") + assert.NoErrorf(t, err, "Failed to json-encode: %#v", req) + + f, err := os.OpenFile(requestsPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) + assert.NoError(t, err) + defer f.Close() + + _, err = f.WriteString(string(reqJson) + "\n") + assert.NoError(t, err) } - - return req.Workspace.JobsCreate(request) - }) - - server.Handle("POST", "/api/2.0/pipelines", func(req testserver.Request) any { - var request pipelines.PipelineSpec - if err := json.Unmarshal(req.Body, &request); err != nil { - return testserver.Response{ - Body: fmt.Sprintf("internal error: %s", err), - StatusCode: 400, - } + } + + if logRequests { + s.ResponseCallback = func(request *testserver.Request, response *testserver.EncodedResponse) { + t.Logf("%d %s %s\n%s\n%s", + response.StatusCode, request.Method, request.URL, + formatHeadersAndBody("> ", request.Headers, request.Body), + formatHeadersAndBody("# ", response.Headers, response.Body), + ) } + } + + for _, stub := range stubs { + require.NotEmpty(t, stub.Pattern) + items := strings.Split(stub.Pattern, " ") + require.Len(t, items, 2) + s.Handle(items[0], items[1], func(req testserver.Request) any { + time.Sleep(stub.Delay) + return stub.Response + }) + } + + // The earliest handlers take precedence, add default handlers last + addDefaultHandlers(s) + + return s.URL +} - return req.Workspace.PipelinesCreate(request) - }) +type LoggedRequest struct { + Headers http.Header `json:"headers,omitempty"` + Method string `json:"method"` + Path string `json:"path"` + Body any `json:"body,omitempty"` + RawBody string `json:"raw_body,omitempty"` +} - server.Handle("GET", "/api/2.2/jobs/get", func(req testserver.Request) any { - jobId := req.URL.Query().Get("job_id") - return req.Workspace.JobsGet(jobId) - }) +func getLoggedRequest(req *testserver.Request, includedHeaders []string) LoggedRequest { + result := LoggedRequest{ + Method: req.Method, + Path: req.URL.Path, + Headers: filterHeaders(req.Headers, includedHeaders), + } - server.Handle("GET", "/api/2.0/pipelines/{pipeline_id}", func(req testserver.Request) any { - pipelineId := req.Vars["pipeline_id"] - return req.Workspace.PipelinesGet(pipelineId) - }) + if json.Valid(req.Body) { + result.Body = json.RawMessage(req.Body) + } else { + result.RawBody = string(req.Body) + } - server.Handle("GET", "/api/2.2/jobs/list", func(req testserver.Request) any { - return req.Workspace.JobsList() - }) - - server.Handle("GET", "/api/2.2/jobs/list", func(req testserver.Request) any { - return req.Workspace.JobsList() - }) + return result +} - server.Handle("GET", "/oidc/.well-known/oauth-authorization-server", func(_ testserver.Request) any { - return map[string]string{ - "authorization_endpoint": server.URL + "oidc/v1/authorize", - "token_endpoint": server.URL + "/oidc/v1/token", +func filterHeaders(h http.Header, includedHeaders []string) http.Header { + headers := make(http.Header) + for k, v := range h { + if !slices.Contains(includedHeaders, k) { + continue } - }) + headers[k] = v + } + return headers +} - server.Handle("POST", "/oidc/v1/token", func(_ testserver.Request) any { - return map[string]string{ - "access_token": "oauth-token", - "expires_in": "3600", - "scope": "all-apis", - "token_type": "Bearer", +func formatHeadersAndBody(prefix string, headers http.Header, body []byte) string { + var result []string + for key, values := range headers { + if len(values) == 1 { + result = append(result, fmt.Sprintf("%s%s: %s", prefix, key, values[0])) + } else { + result = append(result, fmt.Sprintf("%s%s: %s", prefix, key, values)) } - }) - - server.Handle("POST", "/telemetry-ext", func(_ testserver.Request) any { - return telemetry.ResponseBody{ - Errors: []telemetry.LogError{}, - NumProtoSuccess: 1, + } + if len(body) > 0 { + var s string + if utf8.Valid(body) { + s = string(body) + } else { + s = fmt.Sprintf("[Binary %d bytes]", len(body)) } - }) - - // Quality monitors: - - server.Handle("GET", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { - return testserver.MapGet(req.Workspace.Monitors, req.Vars["table_name"]) - }) - - server.Handle("POST", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { - return req.Workspace.QualityMonitorUpsert(req, req.Vars["table_name"], false) - }) - - server.Handle("PUT", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { - return req.Workspace.QualityMonitorUpsert(req, req.Vars["table_name"], true) - }) - - server.Handle("DELETE", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { - return testserver.MapDelete(req.Workspace.Monitors, req.Vars["table_name"]) - }) - - // Apps: - - server.Handle("GET", "/api/2.0/apps/{name}", func(req testserver.Request) any { - return testserver.MapGet(req.Workspace.Apps, req.Vars["name"]) - }) - - server.Handle("POST", "/api/2.0/apps", func(req testserver.Request) any { - return req.Workspace.AppsUpsert(req, "") - }) - - server.Handle("PATCH", "/api/2.0/apps/{name}", func(req testserver.Request) any { - return req.Workspace.AppsUpsert(req, req.Vars["name"]) - }) - - server.Handle("DELETE", "/api/2.0/apps/{name}", func(req testserver.Request) any { - return testserver.MapDelete(req.Workspace.Apps, req.Vars["name"]) - }) - - // Schemas: - - server.Handle("GET", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { - return testserver.MapGet(req.Workspace.Schemas, req.Vars["full_name"]) - }) - - server.Handle("POST", "/api/2.1/unity-catalog/schemas", func(req testserver.Request) any { - return req.Workspace.SchemasCreate(req) - }) - - server.Handle("PATCH", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { - return req.Workspace.SchemasUpdate(req, req.Vars["full_name"]) - }) - - server.Handle("DELETE", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { - return testserver.MapDelete(req.Workspace.Schemas, req.Vars["full_name"]) - }) + s = strings.ReplaceAll(s, "\n", "\n"+prefix) + result = append(result, prefix+s) + } + return strings.Join(result, "\n") } From 4f781aba843047f0ee1b1caea68ac3ad775671dd Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 14 Apr 2025 15:38:09 +0200 Subject: [PATCH 02/57] - --- acceptance/internal/handlers.go | 261 ------------------ acceptance/internal/resolve_server.go | 186 +++++++++++++ acceptance/internal/server.go | 377 +++++++++++++++----------- 3 files changed, 412 insertions(+), 412 deletions(-) delete mode 100644 acceptance/internal/handlers.go create mode 100644 acceptance/internal/resolve_server.go diff --git a/acceptance/internal/handlers.go b/acceptance/internal/handlers.go deleted file mode 100644 index 4f1665b027..0000000000 --- a/acceptance/internal/handlers.go +++ /dev/null @@ -1,261 +0,0 @@ -package internal - -import ( - "encoding/json" - "fmt" - "net/http" - - "github.com/databricks/databricks-sdk-go/service/catalog" - "github.com/databricks/databricks-sdk-go/service/iam" - "github.com/databricks/databricks-sdk-go/service/pipelines" - - "github.com/databricks/databricks-sdk-go/service/compute" - "github.com/databricks/databricks-sdk-go/service/jobs" - - "github.com/databricks/cli/libs/telemetry" - "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go/service/workspace" -) - -var TestUser = iam.User{ - Id: "1000012345", - UserName: "tester@databricks.com", -} - -var TestMetastore = catalog.MetastoreAssignment{ - DefaultCatalogName: "hive_metastore", - MetastoreId: "120efa64-9b68-46ba-be38-f319458430d2", - WorkspaceId: 470123456789500, -} - -func addDefaultHandlers(server *testserver.Server) { - server.Handle("GET", "/api/2.0/policies/clusters/list", func(req testserver.Request) any { - return compute.ListPoliciesResponse{ - Policies: []compute.Policy{ - { - PolicyId: "5678", - Name: "wrong-cluster-policy", - }, - { - PolicyId: "9876", - Name: "some-test-cluster-policy", - }, - }, - } - }) - - server.Handle("GET", "/api/2.0/instance-pools/list", func(req testserver.Request) any { - return compute.ListInstancePools{ - InstancePools: []compute.InstancePoolAndStats{ - { - InstancePoolName: "some-test-instance-pool", - InstancePoolId: "1234", - }, - }, - } - }) - - server.Handle("GET", "/api/2.1/clusters/list", func(req testserver.Request) any { - return compute.ListClustersResponse{ - Clusters: []compute.ClusterDetails{ - { - ClusterName: "some-test-cluster", - ClusterId: "4321", - }, - { - ClusterName: "some-other-cluster", - ClusterId: "9876", - }, - }, - } - }) - - server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(req testserver.Request) any { - return testserver.Response{ - Headers: map[string][]string{"X-Databricks-Org-Id": {"900800700600"}}, - Body: TestUser, - } - }) - - server.Handle("GET", "/api/2.0/workspace/get-status", func(req testserver.Request) any { - path := req.URL.Query().Get("path") - return req.Workspace.WorkspaceGetStatus(path) - }) - - server.Handle("POST", "/api/2.0/workspace/mkdirs", func(req testserver.Request) any { - var request workspace.Mkdirs - if err := json.Unmarshal(req.Body, &request); err != nil { - return testserver.Response{ - Body: fmt.Sprintf("internal error: %s", err), - StatusCode: http.StatusInternalServerError, - } - } - - req.Workspace.WorkspaceMkdirs(request) - return "" - }) - - server.Handle("GET", "/api/2.0/workspace/export", func(req testserver.Request) any { - path := req.URL.Query().Get("path") - return req.Workspace.WorkspaceExport(path) - }) - - server.Handle("POST", "/api/2.0/workspace/delete", func(req testserver.Request) any { - path := req.URL.Query().Get("path") - recursive := req.URL.Query().Get("recursive") == "true" - req.Workspace.WorkspaceDelete(path, recursive) - return "" - }) - - server.Handle("POST", "/api/2.0/workspace-files/import-file/{path:.*}", func(req testserver.Request) any { - path := req.Vars["path"] - req.Workspace.WorkspaceFilesImportFile(path, req.Body) - return "" - }) - - server.Handle("GET", "/api/2.0/workspace-files/{path:.*}", func(req testserver.Request) any { - path := req.Vars["path"] - return req.Workspace.WorkspaceFilesExportFile(path) - }) - - server.Handle("GET", "/api/2.1/unity-catalog/current-metastore-assignment", func(req testserver.Request) any { - return TestMetastore - }) - - server.Handle("GET", "/api/2.0/permissions/directories/{objectId}", func(req testserver.Request) any { - objectId := req.Vars["objectId"] - return workspace.WorkspaceObjectPermissions{ - ObjectId: objectId, - ObjectType: "DIRECTORY", - AccessControlList: []workspace.WorkspaceObjectAccessControlResponse{ - { - UserName: "tester@databricks.com", - AllPermissions: []workspace.WorkspaceObjectPermission{ - { - PermissionLevel: "CAN_MANAGE", - }, - }, - }, - }, - } - }) - - server.Handle("POST", "/api/2.2/jobs/create", func(req testserver.Request) any { - var request jobs.CreateJob - if err := json.Unmarshal(req.Body, &request); err != nil { - return testserver.Response{ - Body: fmt.Sprintf("internal error: %s", err), - StatusCode: 500, - } - } - - return req.Workspace.JobsCreate(request) - }) - - server.Handle("POST", "/api/2.0/pipelines", func(req testserver.Request) any { - var request pipelines.PipelineSpec - if err := json.Unmarshal(req.Body, &request); err != nil { - return testserver.Response{ - Body: fmt.Sprintf("internal error: %s", err), - StatusCode: 400, - } - } - - return req.Workspace.PipelinesCreate(request) - }) - - server.Handle("GET", "/api/2.2/jobs/get", func(req testserver.Request) any { - jobId := req.URL.Query().Get("job_id") - return req.Workspace.JobsGet(jobId) - }) - - server.Handle("GET", "/api/2.0/pipelines/{pipeline_id}", func(req testserver.Request) any { - pipelineId := req.Vars["pipeline_id"] - return req.Workspace.PipelinesGet(pipelineId) - }) - - server.Handle("GET", "/api/2.2/jobs/list", func(req testserver.Request) any { - return req.Workspace.JobsList() - }) - - server.Handle("GET", "/api/2.2/jobs/list", func(req testserver.Request) any { - return req.Workspace.JobsList() - }) - - server.Handle("GET", "/oidc/.well-known/oauth-authorization-server", func(_ testserver.Request) any { - return map[string]string{ - "authorization_endpoint": server.URL + "oidc/v1/authorize", - "token_endpoint": server.URL + "/oidc/v1/token", - } - }) - - server.Handle("POST", "/oidc/v1/token", func(_ testserver.Request) any { - return map[string]string{ - "access_token": "oauth-token", - "expires_in": "3600", - "scope": "all-apis", - "token_type": "Bearer", - } - }) - - server.Handle("POST", "/telemetry-ext", func(_ testserver.Request) any { - return telemetry.ResponseBody{ - Errors: []telemetry.LogError{}, - NumProtoSuccess: 1, - } - }) - - // Quality monitors: - - server.Handle("GET", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { - return testserver.MapGet(req.Workspace.Monitors, req.Vars["table_name"]) - }) - - server.Handle("POST", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { - return req.Workspace.QualityMonitorUpsert(req, req.Vars["table_name"], false) - }) - - server.Handle("PUT", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { - return req.Workspace.QualityMonitorUpsert(req, req.Vars["table_name"], true) - }) - - server.Handle("DELETE", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { - return testserver.MapDelete(req.Workspace.Monitors, req.Vars["table_name"]) - }) - - // Apps: - - server.Handle("GET", "/api/2.0/apps/{name}", func(req testserver.Request) any { - return testserver.MapGet(req.Workspace.Apps, req.Vars["name"]) - }) - - server.Handle("POST", "/api/2.0/apps", func(req testserver.Request) any { - return req.Workspace.AppsUpsert(req, "") - }) - - server.Handle("PATCH", "/api/2.0/apps/{name}", func(req testserver.Request) any { - return req.Workspace.AppsUpsert(req, req.Vars["name"]) - }) - - server.Handle("DELETE", "/api/2.0/apps/{name}", func(req testserver.Request) any { - return testserver.MapDelete(req.Workspace.Apps, req.Vars["name"]) - }) - - // Schemas: - - server.Handle("GET", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { - return testserver.MapGet(req.Workspace.Schemas, req.Vars["full_name"]) - }) - - server.Handle("POST", "/api/2.1/unity-catalog/schemas", func(req testserver.Request) any { - return req.Workspace.SchemasCreate(req) - }) - - server.Handle("PATCH", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { - return req.Workspace.SchemasUpdate(req, req.Vars["full_name"]) - }) - - server.Handle("DELETE", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { - return testserver.MapDelete(req.Workspace.Schemas, req.Vars["full_name"]) - }) -} diff --git a/acceptance/internal/resolve_server.go b/acceptance/internal/resolve_server.go new file mode 100644 index 0000000000..5d7dfcc443 --- /dev/null +++ b/acceptance/internal/resolve_server.go @@ -0,0 +1,186 @@ +package internal + +import ( + "encoding/json" + "fmt" + "net/http" + "os" + "path/filepath" + "slices" + "strings" + "testing" + "time" + "unicode/utf8" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func StartDefaultServer(t *testing.T) { + s := testserver.New(t) + addDefaultHandlers(s) + + t.Setenv("DATABRICKS_DEFAULT_HOST", s.URL) + + // Do not read user's ~/.databrickscfg + homeDir := t.TempDir() + t.Setenv(env.HomeEnvVar(), homeDir) +} + +func isTruePtr(value *bool) bool { + return value != nil && *value +} + +func ResolveServer(t *testing.T, config TestConfig, logRequests bool, outputDir string) *databricks.WorkspaceClient { + cloudEnv := os.Getenv("CLOUD_ENV") + + // If we are running on a cloud environment, use the host configured in the + // environment. + if cloudEnv != "" { + w, err := databricks.NewWorkspaceClient(&databricks.Config{}) + require.NoError(t, err) + + return w + } + + recordRequests := isTruePtr(config.RecordRequests) + + tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") + token := fmt.Sprintf("dbapi%s", tokenSuffix) + + // If we are not recording requests, and no custom server server stubs are configured, + // use the default shared server. + if len(config.Server) == 0 && !recordRequests { + w, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: os.Getenv("DATABRICKS_DEFAULT_HOST"), + Token: token, + }) + require.NoError(t, err) + + return w + } + + // We want later stubs takes precedence, because then leaf configs take precedence over parent directory configs + // In gorilla/mux earlier handlers take precedence, so we need to reverse the order + slices.Reverse(config.Server) + host := startDedicatedServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + + w, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: host, + Token: token, + }) + require.NoError(t, err) + return w +} + +func startDedicatedServer(t *testing.T, + stubs []ServerStub, + recordRequests bool, + logRequests bool, + includedHeaders []string, + outputDir string, +) string { + s := testserver.New(t) + + if recordRequests { + requestsPath := filepath.Join(outputDir, "out.requests.txt") + s.RequestCallback = func(request *testserver.Request) { + req := getLoggedRequest(request, includedHeaders) + reqJson, err := json.MarshalIndent(req, "", " ") + assert.NoErrorf(t, err, "Failed to json-encode: %#v", req) + + f, err := os.OpenFile(requestsPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) + assert.NoError(t, err) + defer f.Close() + + _, err = f.WriteString(string(reqJson) + "\n") + assert.NoError(t, err) + } + } + + if logRequests { + s.ResponseCallback = func(request *testserver.Request, response *testserver.EncodedResponse) { + t.Logf("%d %s %s\n%s\n%s", + response.StatusCode, request.Method, request.URL, + formatHeadersAndBody("> ", request.Headers, request.Body), + formatHeadersAndBody("# ", response.Headers, response.Body), + ) + } + } + + for _, stub := range stubs { + require.NotEmpty(t, stub.Pattern) + items := strings.Split(stub.Pattern, " ") + require.Len(t, items, 2) + s.Handle(items[0], items[1], func(req testserver.Request) any { + time.Sleep(stub.Delay) + return stub.Response + }) + } + + // The earliest handlers take precedence, add default handlers last + addDefaultHandlers(s) + + return s.URL +} + +type LoggedRequest struct { + Headers http.Header `json:"headers,omitempty"` + Method string `json:"method"` + Path string `json:"path"` + Body any `json:"body,omitempty"` + RawBody string `json:"raw_body,omitempty"` +} + +func getLoggedRequest(req *testserver.Request, includedHeaders []string) LoggedRequest { + result := LoggedRequest{ + Method: req.Method, + Path: req.URL.Path, + Headers: filterHeaders(req.Headers, includedHeaders), + } + + if json.Valid(req.Body) { + result.Body = json.RawMessage(req.Body) + } else { + result.RawBody = string(req.Body) + } + + return result +} + +func filterHeaders(h http.Header, includedHeaders []string) http.Header { + headers := make(http.Header) + for k, v := range h { + if !slices.Contains(includedHeaders, k) { + continue + } + headers[k] = v + } + return headers +} + +func formatHeadersAndBody(prefix string, headers http.Header, body []byte) string { + var result []string + for key, values := range headers { + if len(values) == 1 { + result = append(result, fmt.Sprintf("%s%s: %s", prefix, key, values[0])) + } else { + result = append(result, fmt.Sprintf("%s%s: %s", prefix, key, values)) + } + } + if len(body) > 0 { + var s string + if utf8.Valid(body) { + s = string(body) + } else { + s = fmt.Sprintf("[Binary %d bytes]", len(body)) + } + s = strings.ReplaceAll(s, "\n", "\n"+prefix) + result = append(result, prefix+s) + } + return strings.Join(result, "\n") +} diff --git a/acceptance/internal/server.go b/acceptance/internal/server.go index 5d7dfcc443..4f1665b027 100644 --- a/acceptance/internal/server.go +++ b/acceptance/internal/server.go @@ -4,183 +4,258 @@ import ( "encoding/json" "fmt" "net/http" - "os" - "path/filepath" - "slices" - "strings" - "testing" - "time" - "unicode/utf8" - - "github.com/databricks/cli/libs/env" - "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go" - "github.com/google/uuid" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) -func StartDefaultServer(t *testing.T) { - s := testserver.New(t) - addDefaultHandlers(s) + "github.com/databricks/databricks-sdk-go/service/catalog" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/pipelines" + + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/jobs" - t.Setenv("DATABRICKS_DEFAULT_HOST", s.URL) + "github.com/databricks/cli/libs/telemetry" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go/service/workspace" +) - // Do not read user's ~/.databrickscfg - homeDir := t.TempDir() - t.Setenv(env.HomeEnvVar(), homeDir) +var TestUser = iam.User{ + Id: "1000012345", + UserName: "tester@databricks.com", } -func isTruePtr(value *bool) bool { - return value != nil && *value +var TestMetastore = catalog.MetastoreAssignment{ + DefaultCatalogName: "hive_metastore", + MetastoreId: "120efa64-9b68-46ba-be38-f319458430d2", + WorkspaceId: 470123456789500, } -func ResolveServer(t *testing.T, config TestConfig, logRequests bool, outputDir string) *databricks.WorkspaceClient { - cloudEnv := os.Getenv("CLOUD_ENV") +func addDefaultHandlers(server *testserver.Server) { + server.Handle("GET", "/api/2.0/policies/clusters/list", func(req testserver.Request) any { + return compute.ListPoliciesResponse{ + Policies: []compute.Policy{ + { + PolicyId: "5678", + Name: "wrong-cluster-policy", + }, + { + PolicyId: "9876", + Name: "some-test-cluster-policy", + }, + }, + } + }) - // If we are running on a cloud environment, use the host configured in the - // environment. - if cloudEnv != "" { - w, err := databricks.NewWorkspaceClient(&databricks.Config{}) - require.NoError(t, err) + server.Handle("GET", "/api/2.0/instance-pools/list", func(req testserver.Request) any { + return compute.ListInstancePools{ + InstancePools: []compute.InstancePoolAndStats{ + { + InstancePoolName: "some-test-instance-pool", + InstancePoolId: "1234", + }, + }, + } + }) - return w - } + server.Handle("GET", "/api/2.1/clusters/list", func(req testserver.Request) any { + return compute.ListClustersResponse{ + Clusters: []compute.ClusterDetails{ + { + ClusterName: "some-test-cluster", + ClusterId: "4321", + }, + { + ClusterName: "some-other-cluster", + ClusterId: "9876", + }, + }, + } + }) - recordRequests := isTruePtr(config.RecordRequests) + server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(req testserver.Request) any { + return testserver.Response{ + Headers: map[string][]string{"X-Databricks-Org-Id": {"900800700600"}}, + Body: TestUser, + } + }) - tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") - token := fmt.Sprintf("dbapi%s", tokenSuffix) + server.Handle("GET", "/api/2.0/workspace/get-status", func(req testserver.Request) any { + path := req.URL.Query().Get("path") + return req.Workspace.WorkspaceGetStatus(path) + }) - // If we are not recording requests, and no custom server server stubs are configured, - // use the default shared server. - if len(config.Server) == 0 && !recordRequests { - w, err := databricks.NewWorkspaceClient(&databricks.Config{ - Host: os.Getenv("DATABRICKS_DEFAULT_HOST"), - Token: token, - }) - require.NoError(t, err) + server.Handle("POST", "/api/2.0/workspace/mkdirs", func(req testserver.Request) any { + var request workspace.Mkdirs + if err := json.Unmarshal(req.Body, &request); err != nil { + return testserver.Response{ + Body: fmt.Sprintf("internal error: %s", err), + StatusCode: http.StatusInternalServerError, + } + } - return w - } + req.Workspace.WorkspaceMkdirs(request) + return "" + }) - // We want later stubs takes precedence, because then leaf configs take precedence over parent directory configs - // In gorilla/mux earlier handlers take precedence, so we need to reverse the order - slices.Reverse(config.Server) - host := startDedicatedServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + server.Handle("GET", "/api/2.0/workspace/export", func(req testserver.Request) any { + path := req.URL.Query().Get("path") + return req.Workspace.WorkspaceExport(path) + }) - w, err := databricks.NewWorkspaceClient(&databricks.Config{ - Host: host, - Token: token, + server.Handle("POST", "/api/2.0/workspace/delete", func(req testserver.Request) any { + path := req.URL.Query().Get("path") + recursive := req.URL.Query().Get("recursive") == "true" + req.Workspace.WorkspaceDelete(path, recursive) + return "" }) - require.NoError(t, err) - return w -} -func startDedicatedServer(t *testing.T, - stubs []ServerStub, - recordRequests bool, - logRequests bool, - includedHeaders []string, - outputDir string, -) string { - s := testserver.New(t) - - if recordRequests { - requestsPath := filepath.Join(outputDir, "out.requests.txt") - s.RequestCallback = func(request *testserver.Request) { - req := getLoggedRequest(request, includedHeaders) - reqJson, err := json.MarshalIndent(req, "", " ") - assert.NoErrorf(t, err, "Failed to json-encode: %#v", req) - - f, err := os.OpenFile(requestsPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) - assert.NoError(t, err) - defer f.Close() - - _, err = f.WriteString(string(reqJson) + "\n") - assert.NoError(t, err) + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path:.*}", func(req testserver.Request) any { + path := req.Vars["path"] + req.Workspace.WorkspaceFilesImportFile(path, req.Body) + return "" + }) + + server.Handle("GET", "/api/2.0/workspace-files/{path:.*}", func(req testserver.Request) any { + path := req.Vars["path"] + return req.Workspace.WorkspaceFilesExportFile(path) + }) + + server.Handle("GET", "/api/2.1/unity-catalog/current-metastore-assignment", func(req testserver.Request) any { + return TestMetastore + }) + + server.Handle("GET", "/api/2.0/permissions/directories/{objectId}", func(req testserver.Request) any { + objectId := req.Vars["objectId"] + return workspace.WorkspaceObjectPermissions{ + ObjectId: objectId, + ObjectType: "DIRECTORY", + AccessControlList: []workspace.WorkspaceObjectAccessControlResponse{ + { + UserName: "tester@databricks.com", + AllPermissions: []workspace.WorkspaceObjectPermission{ + { + PermissionLevel: "CAN_MANAGE", + }, + }, + }, + }, } - } - - if logRequests { - s.ResponseCallback = func(request *testserver.Request, response *testserver.EncodedResponse) { - t.Logf("%d %s %s\n%s\n%s", - response.StatusCode, request.Method, request.URL, - formatHeadersAndBody("> ", request.Headers, request.Body), - formatHeadersAndBody("# ", response.Headers, response.Body), - ) + }) + + server.Handle("POST", "/api/2.2/jobs/create", func(req testserver.Request) any { + var request jobs.CreateJob + if err := json.Unmarshal(req.Body, &request); err != nil { + return testserver.Response{ + Body: fmt.Sprintf("internal error: %s", err), + StatusCode: 500, + } } - } - - for _, stub := range stubs { - require.NotEmpty(t, stub.Pattern) - items := strings.Split(stub.Pattern, " ") - require.Len(t, items, 2) - s.Handle(items[0], items[1], func(req testserver.Request) any { - time.Sleep(stub.Delay) - return stub.Response - }) - } - - // The earliest handlers take precedence, add default handlers last - addDefaultHandlers(s) - - return s.URL -} -type LoggedRequest struct { - Headers http.Header `json:"headers,omitempty"` - Method string `json:"method"` - Path string `json:"path"` - Body any `json:"body,omitempty"` - RawBody string `json:"raw_body,omitempty"` -} + return req.Workspace.JobsCreate(request) + }) -func getLoggedRequest(req *testserver.Request, includedHeaders []string) LoggedRequest { - result := LoggedRequest{ - Method: req.Method, - Path: req.URL.Path, - Headers: filterHeaders(req.Headers, includedHeaders), - } + server.Handle("POST", "/api/2.0/pipelines", func(req testserver.Request) any { + var request pipelines.PipelineSpec + if err := json.Unmarshal(req.Body, &request); err != nil { + return testserver.Response{ + Body: fmt.Sprintf("internal error: %s", err), + StatusCode: 400, + } + } - if json.Valid(req.Body) { - result.Body = json.RawMessage(req.Body) - } else { - result.RawBody = string(req.Body) - } + return req.Workspace.PipelinesCreate(request) + }) - return result -} + server.Handle("GET", "/api/2.2/jobs/get", func(req testserver.Request) any { + jobId := req.URL.Query().Get("job_id") + return req.Workspace.JobsGet(jobId) + }) + + server.Handle("GET", "/api/2.0/pipelines/{pipeline_id}", func(req testserver.Request) any { + pipelineId := req.Vars["pipeline_id"] + return req.Workspace.PipelinesGet(pipelineId) + }) -func filterHeaders(h http.Header, includedHeaders []string) http.Header { - headers := make(http.Header) - for k, v := range h { - if !slices.Contains(includedHeaders, k) { - continue + server.Handle("GET", "/api/2.2/jobs/list", func(req testserver.Request) any { + return req.Workspace.JobsList() + }) + + server.Handle("GET", "/api/2.2/jobs/list", func(req testserver.Request) any { + return req.Workspace.JobsList() + }) + + server.Handle("GET", "/oidc/.well-known/oauth-authorization-server", func(_ testserver.Request) any { + return map[string]string{ + "authorization_endpoint": server.URL + "oidc/v1/authorize", + "token_endpoint": server.URL + "/oidc/v1/token", } - headers[k] = v - } - return headers -} + }) -func formatHeadersAndBody(prefix string, headers http.Header, body []byte) string { - var result []string - for key, values := range headers { - if len(values) == 1 { - result = append(result, fmt.Sprintf("%s%s: %s", prefix, key, values[0])) - } else { - result = append(result, fmt.Sprintf("%s%s: %s", prefix, key, values)) + server.Handle("POST", "/oidc/v1/token", func(_ testserver.Request) any { + return map[string]string{ + "access_token": "oauth-token", + "expires_in": "3600", + "scope": "all-apis", + "token_type": "Bearer", } - } - if len(body) > 0 { - var s string - if utf8.Valid(body) { - s = string(body) - } else { - s = fmt.Sprintf("[Binary %d bytes]", len(body)) + }) + + server.Handle("POST", "/telemetry-ext", func(_ testserver.Request) any { + return telemetry.ResponseBody{ + Errors: []telemetry.LogError{}, + NumProtoSuccess: 1, } - s = strings.ReplaceAll(s, "\n", "\n"+prefix) - result = append(result, prefix+s) - } - return strings.Join(result, "\n") + }) + + // Quality monitors: + + server.Handle("GET", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { + return testserver.MapGet(req.Workspace.Monitors, req.Vars["table_name"]) + }) + + server.Handle("POST", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { + return req.Workspace.QualityMonitorUpsert(req, req.Vars["table_name"], false) + }) + + server.Handle("PUT", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { + return req.Workspace.QualityMonitorUpsert(req, req.Vars["table_name"], true) + }) + + server.Handle("DELETE", "/api/2.1/unity-catalog/tables/{table_name}/monitor", func(req testserver.Request) any { + return testserver.MapDelete(req.Workspace.Monitors, req.Vars["table_name"]) + }) + + // Apps: + + server.Handle("GET", "/api/2.0/apps/{name}", func(req testserver.Request) any { + return testserver.MapGet(req.Workspace.Apps, req.Vars["name"]) + }) + + server.Handle("POST", "/api/2.0/apps", func(req testserver.Request) any { + return req.Workspace.AppsUpsert(req, "") + }) + + server.Handle("PATCH", "/api/2.0/apps/{name}", func(req testserver.Request) any { + return req.Workspace.AppsUpsert(req, req.Vars["name"]) + }) + + server.Handle("DELETE", "/api/2.0/apps/{name}", func(req testserver.Request) any { + return testserver.MapDelete(req.Workspace.Apps, req.Vars["name"]) + }) + + // Schemas: + + server.Handle("GET", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { + return testserver.MapGet(req.Workspace.Schemas, req.Vars["full_name"]) + }) + + server.Handle("POST", "/api/2.1/unity-catalog/schemas", func(req testserver.Request) any { + return req.Workspace.SchemasCreate(req) + }) + + server.Handle("PATCH", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { + return req.Workspace.SchemasUpdate(req, req.Vars["full_name"]) + }) + + server.Handle("DELETE", "/api/2.1/unity-catalog/schemas/{full_name}", func(req testserver.Request) any { + return testserver.MapDelete(req.Workspace.Schemas, req.Vars["full_name"]) + }) } From 2f5aca33e4e6f17df510d8edffcd1edac86e883f Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 14 Apr 2025 15:38:36 +0200 Subject: [PATCH 03/57] - --- acceptance/internal/{server.go => handlers.go} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename acceptance/internal/{server.go => handlers.go} (100%) diff --git a/acceptance/internal/server.go b/acceptance/internal/handlers.go similarity index 100% rename from acceptance/internal/server.go rename to acceptance/internal/handlers.go From 0892552f4c1af1faaf254098840e26fd6bcf917f Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 14 Apr 2025 15:43:09 +0200 Subject: [PATCH 04/57] - --- acceptance/internal/resolve_server.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/acceptance/internal/resolve_server.go b/acceptance/internal/resolve_server.go index 5d7dfcc443..800f862e0c 100644 --- a/acceptance/internal/resolve_server.go +++ b/acceptance/internal/resolve_server.go @@ -81,7 +81,7 @@ func startDedicatedServer(t *testing.T, stubs []ServerStub, recordRequests bool, logRequests bool, - includedHeaders []string, + includeHeaders []string, outputDir string, ) string { s := testserver.New(t) @@ -89,7 +89,7 @@ func startDedicatedServer(t *testing.T, if recordRequests { requestsPath := filepath.Join(outputDir, "out.requests.txt") s.RequestCallback = func(request *testserver.Request) { - req := getLoggedRequest(request, includedHeaders) + req := getLoggedRequest(request, includeHeaders) reqJson, err := json.MarshalIndent(req, "", " ") assert.NoErrorf(t, err, "Failed to json-encode: %#v", req) From 0e7806b8537f58ba86e0f90d177e7f3be2e26483 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 14 Apr 2025 15:43:55 +0200 Subject: [PATCH 05/57] - --- acceptance/internal/resolve_server.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/acceptance/internal/resolve_server.go b/acceptance/internal/resolve_server.go index 800f862e0c..1dab7ae170 100644 --- a/acceptance/internal/resolve_server.go +++ b/acceptance/internal/resolve_server.go @@ -50,7 +50,7 @@ func ResolveServer(t *testing.T, config TestConfig, logRequests bool, outputDir recordRequests := isTruePtr(config.RecordRequests) tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") - token := fmt.Sprintf("dbapi%s", tokenSuffix) + token := "dbapi" + tokenSuffix // If we are not recording requests, and no custom server server stubs are configured, // use the default shared server. From 465d74d07f9418d203e774975aa8e5c4015c4f51 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Tue, 15 Apr 2025 00:30:39 +0200 Subject: [PATCH 06/57] - --- acceptance/acceptance_test.go | 36 ++++----- acceptance/internal/cmd_server.go | 2 +- acceptance/internal/resolve_server.go | 73 +++++++++++++----- acceptance/selftest/record_cloud/output.txt | 5 ++ acceptance/selftest/record_cloud/script | 7 ++ acceptance/selftest/record_cloud/test.toml | 3 + libs/testserver/server.go | 84 ++++++++++++++++++++- 7 files changed, 166 insertions(+), 44 deletions(-) create mode 100644 acceptance/selftest/record_cloud/output.txt create mode 100644 acceptance/selftest/record_cloud/script create mode 100644 acceptance/selftest/record_cloud/test.toml diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index cfea953de7..be891deda1 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -2,7 +2,6 @@ package acceptance_test import ( "bufio" - "context" "encoding/base32" "encoding/json" "errors" @@ -26,8 +25,9 @@ import ( "github.com/databricks/cli/acceptance/internal" "github.com/databricks/cli/internal/testutil" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/testdiff" - "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go" "github.com/stretchr/testify/require" ) @@ -78,7 +78,7 @@ var Ignored = map[string]bool{ } func TestAccept(t *testing.T) { - testAccept(t, InprocessMode, "") + testAccept(t, InprocessMode, "selftest/record_cloud") } func TestInprocessMode(t *testing.T) { @@ -361,27 +361,21 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont cmd.Env = append(cmd.Env, "UNIQUE_NAME="+uniqueName) cmd.Env = append(cmd.Env, "TEST_TMP_DIR="+tmpDir) - workspaceClient := internal.ResolveServer(t, config, LogRequests, tmpDir) - testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient) + cfg, user := internal.ResolveServer(t, config, LogRequests, tmpDir) + testdiff.PrepareReplacementsUser(t, &repls, user) - // Configure resolved credentials in the environment. - cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+workspaceClient.Config.Host) - if workspaceClient.Config.Token != "" { - cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+workspaceClient.Config.Token) - } + w, err := databricks.NewWorkspaceClient((*databricks.Config)(cfg)) + require.NoError(t, err) + testdiff.PrepareReplacementsWorkspaceClient(t, &repls, w) - var user iam.User - if isRunningOnCloud { - pUser, err := workspaceClient.CurrentUser.Me(context.Background()) - require.NoError(t, err, "Failed to get current user") - user = *pUser - } else { - // For the purposes of replacements, use testUser for local runs. - // Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement: - user = internal.TestUser + // Configure resolved credentials in the environment. + for _, v := range auth.ProcessEnv(cfg) { + cmd.Env = append(cmd.Env, v) } - testdiff.PrepareReplacementsUser(t, &repls, user) - + // cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+workspaceClient.Config.Host) + // if workspaceClient.Config.Token != "" { + // cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+workspaceClient.Config.Token) + // } // Must be added PrepareReplacementsUser, otherwise conflicts with [USERNAME] testdiff.PrepareReplacementsUUID(t, &repls) diff --git a/acceptance/internal/cmd_server.go b/acceptance/internal/cmd_server.go index 69d4f02ce3..ba0083fd0e 100644 --- a/acceptance/internal/cmd_server.go +++ b/acceptance/internal/cmd_server.go @@ -13,7 +13,7 @@ import ( ) func StartCmdServer(t *testing.T) *testserver.Server { - server := testserver.New(t) + server := testserver.New(t, false) server.Handle("GET", "/", func(r testserver.Request) any { q := r.URL.Query() args := strings.Split(q.Get("args"), " ") diff --git a/acceptance/internal/resolve_server.go b/acceptance/internal/resolve_server.go index 1dab7ae170..ad49da60ae 100644 --- a/acceptance/internal/resolve_server.go +++ b/acceptance/internal/resolve_server.go @@ -1,6 +1,7 @@ package internal import ( + "context" "encoding/json" "fmt" "net/http" @@ -15,13 +16,16 @@ import ( "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" + sdkconfig "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +// TODO: Make CLI auth for with this. func StartDefaultServer(t *testing.T) { - s := testserver.New(t) + s := testserver.New(t, false) addDefaultHandlers(s) t.Setenv("DATABRICKS_DEFAULT_HOST", s.URL) @@ -35,46 +39,72 @@ func isTruePtr(value *bool) bool { return value != nil && *value } -func ResolveServer(t *testing.T, config TestConfig, logRequests bool, outputDir string) *databricks.WorkspaceClient { +func ResolveServer(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User) { cloudEnv := os.Getenv("CLOUD_ENV") + recordRequests := isTruePtr(config.RecordRequests) + + tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") + token := "dbapi" + tokenSuffix + + // If we are running in a cloud environment AND we are recording requests, + // start a dedicated server to act as a reverse proxy and record requests. + if cloudEnv != "" && recordRequests { + host := startDedicatedServer(t, config.Server, recordRequests, + logRequests, config.IncludeRequestHeaders, outputDir, + true) + + cfg := &sdkconfig.Config{ + Host: host, + Token: token, + } + + // Use a non-proxy client to fetch user info so that this API call is not recorded + // in out.requests.txt + w, err := databricks.NewWorkspaceClient() + require.NoError(t, err) + + user, err := w.CurrentUser.Me(context.Background()) + require.NoError(t, err) + + return cfg, *user + } // If we are running on a cloud environment, use the host configured in the // environment. if cloudEnv != "" { - w, err := databricks.NewWorkspaceClient(&databricks.Config{}) + cfg := &sdkconfig.Config{} + err := cfg.EnsureResolved() require.NoError(t, err) - return w + return cfg, TestUser } - recordRequests := isTruePtr(config.RecordRequests) - - tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") - token := "dbapi" + tokenSuffix - // If we are not recording requests, and no custom server server stubs are configured, // use the default shared server. if len(config.Server) == 0 && !recordRequests { - w, err := databricks.NewWorkspaceClient(&databricks.Config{ + cfg := &sdkconfig.Config{ Host: os.Getenv("DATABRICKS_DEFAULT_HOST"), Token: token, - }) - require.NoError(t, err) + } - return w + return cfg, TestUser } // We want later stubs takes precedence, because then leaf configs take precedence over parent directory configs // In gorilla/mux earlier handlers take precedence, so we need to reverse the order slices.Reverse(config.Server) - host := startDedicatedServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + host := startDedicatedServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir, false) - w, err := databricks.NewWorkspaceClient(&databricks.Config{ + cfg := &sdkconfig.Config{ Host: host, Token: token, - }) + } + err := cfg.EnsureResolved() require.NoError(t, err) - return w + + // For the purposes of replacements, use testUser for local runs. + // Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement: + return cfg, TestUser } func startDedicatedServer(t *testing.T, @@ -83,8 +113,9 @@ func startDedicatedServer(t *testing.T, logRequests bool, includeHeaders []string, outputDir string, + cloud bool, ) string { - s := testserver.New(t) + s := testserver.New(t, cloud) if recordRequests { requestsPath := filepath.Join(outputDir, "out.requests.txt") @@ -123,7 +154,9 @@ func startDedicatedServer(t *testing.T, } // The earliest handlers take precedence, add default handlers last - addDefaultHandlers(s) + if !cloud { + addDefaultHandlers(s) + } return s.URL } @@ -152,6 +185,8 @@ func getLoggedRequest(req *testserver.Request, includedHeaders []string) LoggedR return result } +// TODO CONTINUE: The proxy is working on aws-prod-ucws. Unblock other environments +// and especially azure-prod-ucws which seems to be setting func filterHeaders(h http.Header, includedHeaders []string) http.Header { headers := make(http.Header) for k, v := range h { diff --git a/acceptance/selftest/record_cloud/output.txt b/acceptance/selftest/record_cloud/output.txt new file mode 100644 index 0000000000..b5c12b09ce --- /dev/null +++ b/acceptance/selftest/record_cloud/output.txt @@ -0,0 +1,5 @@ + +>>> [CLI] current-user me +Error: inner token: Post "[ARM_TENANT_ID]/oauth2/token": Post "[ARM_TENANT_ID]/oauth2/token": unsupported protocol scheme "" + +Exit code: 1 diff --git a/acceptance/selftest/record_cloud/script b/acceptance/selftest/record_cloud/script new file mode 100644 index 0000000000..789a66fe3c --- /dev/null +++ b/acceptance/selftest/record_cloud/script @@ -0,0 +1,7 @@ +trace $CLI current-user me | jq .name + +# TODO CONTINUE(v3): I got it working for a short while but something went wrong again with the proxy. +# Figure out out and make this test pass after the new approach where workspaceClient -> config.Config. +trace cat out.requests.txt | jq 'select(has("path") and .path == "/api/2.0/preview/scim/v2/Me")' + +rm out.requests.txt diff --git a/acceptance/selftest/record_cloud/test.toml b/acceptance/selftest/record_cloud/test.toml new file mode 100644 index 0000000000..20d0343612 --- /dev/null +++ b/acceptance/selftest/record_cloud/test.toml @@ -0,0 +1,3 @@ +Cloud = true +Local = false +RecordRequests = true diff --git a/libs/testserver/server.go b/libs/testserver/server.go index feb907968f..b4b6f4da9a 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -1,12 +1,15 @@ package testserver import ( + "context" "encoding/json" + "errors" "fmt" "io" "net/http" "net/http/httptest" "net/url" + "os" "reflect" "strings" "sync" @@ -15,6 +18,9 @@ import ( "github.com/databricks/cli/internal/testutil" "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/require" ) type Server struct { @@ -23,6 +29,9 @@ type Server struct { t testutil.TestingT + workspaceUrl string + proxyClient *client.DatabricksClient + fakeWorkspaces map[string]*FakeWorkspace mu *sync.Mutex @@ -177,7 +186,7 @@ func getHeaders(value []byte) http.Header { } } -func New(t testutil.TestingT) *Server { +func New(t testutil.TestingT, cloud bool) *Server { router := mux.NewRouter() server := httptest.NewServer(router) t.Cleanup(server.Close) @@ -190,8 +199,25 @@ func New(t testutil.TestingT) *Server { fakeWorkspaces: map[string]*FakeWorkspace{}, } - // Set up the not found handler as fallback - router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // If cloud is true, we proxy all requests to the configured Databricks workspace + // instead of using local server stubs. + var err error + if cloud { + s.workspaceUrl = os.Getenv("DATABRICKS_HOST") + s.proxyClient, err = client.New(&config.Config{}) + require.NoError(t, err) + } + + // Set up the not found handler as fallback. If a proxy URL is configured then this handler + // will proxy the request to the proxy URL. + s.Router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // For integration tests forwards all requests to the workspace instead of returning a 404. + if s.proxyClient != nil { + s.ProxyToCloud(w, r) + return + } + + // Original Not Found Logic pattern := r.Method + " " + r.URL.Path bodyBytes, err := io.ReadAll(r.Body) var body string @@ -234,8 +260,60 @@ Response.Body = '' type HandlerFunc func(req Request) any +func (s *Server) ProxyToCloud(w http.ResponseWriter, r *http.Request) { + request := NewRequest(s.t, r, nil) + s.RequestCallback(&request) + + headers := make(map[string]string) + for k, v := range r.Header { + // Authorization headers will be set by the SDK. No need to pass them along here. + if k == "Authorization" { + continue + } + if k == "Accept-Encoding" { + continue + } + headers[k] = v[0] + } + + queryParams := make(map[string]any) + for k, v := range r.URL.Query() { + queryParams[k] = v[0] + } + + // TODO: Since the response is always JSON, this should be specified in the header. + respB := map[string]any{} + err := s.proxyClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, r.Body, &respB) + require.NoError(s.t, err) // todo remove + if err != nil { + // API errors from the SDK are expected to be of the type apierr.APIError. + apiErr := apierr.APIError{} + if errors.As(err, &apiErr) { + w.WriteHeader(apiErr.StatusCode) + w.Write(respB["message"].([]byte)) + } else { + // Something else went wrong. + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(err.Error())) + } + } + + // Successful response + w.WriteHeader(200) + b, err := json.Marshal(respB) + require.NoError(s.t, err) + w.Write(b) +} + func (s *Server) Handle(method, path string, handler HandlerFunc) { s.Router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + // Overriding the proxy URL is not supported. We can add a configuration option + // to disable the proxy and use the test.toml stub if and when such a use case arises. + if s.proxyClient != nil { + s.ProxyToCloud(w, r) + return + } + // For simplicity we process requests sequentially. It's fast enough because // we don't do any IO except reading and writing request/response bodies. s.mu.Lock() From 875ac171b380b5658afc1f49a63e4c7cced7ae72 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 23 Apr 2025 12:39:38 +0200 Subject: [PATCH 07/57] - --- acceptance/acceptance_test.go | 2 +- acceptance/internal/prepare_server.go | 73 ++++++--- acceptance/internal/resolve_server.go | 221 -------------------------- 3 files changed, 55 insertions(+), 241 deletions(-) delete mode 100644 acceptance/internal/resolve_server.go diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index be891deda1..9eb4c90140 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -361,7 +361,7 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont cmd.Env = append(cmd.Env, "UNIQUE_NAME="+uniqueName) cmd.Env = append(cmd.Env, "TEST_TMP_DIR="+tmpDir) - cfg, user := internal.ResolveServer(t, config, LogRequests, tmpDir) + cfg, user := internal.PrepareServerAndClient(t, config, LogRequests, tmpDir) testdiff.PrepareReplacementsUser(t, &repls, user) w, err := databricks.NewWorkspaceClient((*databricks.Config)(cfg)) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index e30b78cc8f..7c9cdda74e 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -1,6 +1,7 @@ package internal import ( + "context" "encoding/json" "fmt" "net/http" @@ -15,13 +16,16 @@ import ( "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" + sdkconfig "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +// TODO: Make CLI auth for with this. func StartDefaultServer(t *testing.T) { - s := testserver.New(t) + s := testserver.New(t, false) addDefaultHandlers(s) t.Setenv("DATABRICKS_DEFAULT_HOST", s.URL) @@ -35,46 +39,72 @@ func isTruePtr(value *bool) bool { return value != nil && *value } -func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string) *databricks.WorkspaceClient { +func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User) { cloudEnv := os.Getenv("CLOUD_ENV") + recordRequests := isTruePtr(config.RecordRequests) + + tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") + token := "dbapi" + tokenSuffix + + // If we are running in a cloud environment AND we are recording requests, + // start a dedicated server to act as a reverse proxy and record requests. + if cloudEnv != "" && recordRequests { + host := startDedicatedServer(t, config.Server, recordRequests, + logRequests, config.IncludeRequestHeaders, outputDir, + true) + + cfg := &sdkconfig.Config{ + Host: host, + Token: token, + } + + // Use a non-proxy client to fetch user info so that this API call is not recorded + // in out.requests.txt + w, err := databricks.NewWorkspaceClient() + require.NoError(t, err) + + user, err := w.CurrentUser.Me(context.Background()) + require.NoError(t, err) + + return cfg, *user + } // If we are running on a cloud environment, use the host configured in the // environment. if cloudEnv != "" { - w, err := databricks.NewWorkspaceClient(&databricks.Config{}) + cfg := &sdkconfig.Config{} + err := cfg.EnsureResolved() require.NoError(t, err) - return w + return cfg, TestUser } - recordRequests := isTruePtr(config.RecordRequests) - - tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") - token := "dbapi" + tokenSuffix - // If we are not recording requests, and no custom server server stubs are configured, // use the default shared server. if len(config.Server) == 0 && !recordRequests { - w, err := databricks.NewWorkspaceClient(&databricks.Config{ + cfg := &sdkconfig.Config{ Host: os.Getenv("DATABRICKS_DEFAULT_HOST"), Token: token, - }) - require.NoError(t, err) + } - return w + return cfg, TestUser } // We want later stubs takes precedence, because then leaf configs take precedence over parent directory configs // In gorilla/mux earlier handlers take precedence, so we need to reverse the order slices.Reverse(config.Server) - host := startDedicatedServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + host := startDedicatedServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir, false) - w, err := databricks.NewWorkspaceClient(&databricks.Config{ + cfg := &sdkconfig.Config{ Host: host, Token: token, - }) + } + err := cfg.EnsureResolved() require.NoError(t, err) - return w + + // For the purposes of replacements, use testUser for local runs. + // Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement: + return cfg, TestUser } func startDedicatedServer(t *testing.T, @@ -83,8 +113,9 @@ func startDedicatedServer(t *testing.T, logRequests bool, includeHeaders []string, outputDir string, + cloud bool, ) string { - s := testserver.New(t) + s := testserver.New(t, cloud) if recordRequests { requestsPath := filepath.Join(outputDir, "out.requests.txt") @@ -123,7 +154,9 @@ func startDedicatedServer(t *testing.T, } // The earliest handlers take precedence, add default handlers last - addDefaultHandlers(s) + if !cloud { + addDefaultHandlers(s) + } return s.URL } @@ -152,6 +185,8 @@ func getLoggedRequest(req *testserver.Request, includedHeaders []string) LoggedR return result } +// TODO CONTINUE: The proxy is working on aws-prod-ucws. Unblock other environments +// and especially azure-prod-ucws which seems to be setting func filterHeaders(h http.Header, includedHeaders []string) http.Header { headers := make(http.Header) for k, v := range h { diff --git a/acceptance/internal/resolve_server.go b/acceptance/internal/resolve_server.go deleted file mode 100644 index ad49da60ae..0000000000 --- a/acceptance/internal/resolve_server.go +++ /dev/null @@ -1,221 +0,0 @@ -package internal - -import ( - "context" - "encoding/json" - "fmt" - "net/http" - "os" - "path/filepath" - "slices" - "strings" - "testing" - "time" - "unicode/utf8" - - "github.com/databricks/cli/libs/env" - "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go" - sdkconfig "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/service/iam" - "github.com/google/uuid" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// TODO: Make CLI auth for with this. -func StartDefaultServer(t *testing.T) { - s := testserver.New(t, false) - addDefaultHandlers(s) - - t.Setenv("DATABRICKS_DEFAULT_HOST", s.URL) - - // Do not read user's ~/.databrickscfg - homeDir := t.TempDir() - t.Setenv(env.HomeEnvVar(), homeDir) -} - -func isTruePtr(value *bool) bool { - return value != nil && *value -} - -func ResolveServer(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User) { - cloudEnv := os.Getenv("CLOUD_ENV") - recordRequests := isTruePtr(config.RecordRequests) - - tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") - token := "dbapi" + tokenSuffix - - // If we are running in a cloud environment AND we are recording requests, - // start a dedicated server to act as a reverse proxy and record requests. - if cloudEnv != "" && recordRequests { - host := startDedicatedServer(t, config.Server, recordRequests, - logRequests, config.IncludeRequestHeaders, outputDir, - true) - - cfg := &sdkconfig.Config{ - Host: host, - Token: token, - } - - // Use a non-proxy client to fetch user info so that this API call is not recorded - // in out.requests.txt - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) - - user, err := w.CurrentUser.Me(context.Background()) - require.NoError(t, err) - - return cfg, *user - } - - // If we are running on a cloud environment, use the host configured in the - // environment. - if cloudEnv != "" { - cfg := &sdkconfig.Config{} - err := cfg.EnsureResolved() - require.NoError(t, err) - - return cfg, TestUser - } - - // If we are not recording requests, and no custom server server stubs are configured, - // use the default shared server. - if len(config.Server) == 0 && !recordRequests { - cfg := &sdkconfig.Config{ - Host: os.Getenv("DATABRICKS_DEFAULT_HOST"), - Token: token, - } - - return cfg, TestUser - } - - // We want later stubs takes precedence, because then leaf configs take precedence over parent directory configs - // In gorilla/mux earlier handlers take precedence, so we need to reverse the order - slices.Reverse(config.Server) - host := startDedicatedServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir, false) - - cfg := &sdkconfig.Config{ - Host: host, - Token: token, - } - err := cfg.EnsureResolved() - require.NoError(t, err) - - // For the purposes of replacements, use testUser for local runs. - // Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement: - return cfg, TestUser -} - -func startDedicatedServer(t *testing.T, - stubs []ServerStub, - recordRequests bool, - logRequests bool, - includeHeaders []string, - outputDir string, - cloud bool, -) string { - s := testserver.New(t, cloud) - - if recordRequests { - requestsPath := filepath.Join(outputDir, "out.requests.txt") - s.RequestCallback = func(request *testserver.Request) { - req := getLoggedRequest(request, includeHeaders) - reqJson, err := json.MarshalIndent(req, "", " ") - assert.NoErrorf(t, err, "Failed to json-encode: %#v", req) - - f, err := os.OpenFile(requestsPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) - assert.NoError(t, err) - defer f.Close() - - _, err = f.WriteString(string(reqJson) + "\n") - assert.NoError(t, err) - } - } - - if logRequests { - s.ResponseCallback = func(request *testserver.Request, response *testserver.EncodedResponse) { - t.Logf("%d %s %s\n%s\n%s", - response.StatusCode, request.Method, request.URL, - formatHeadersAndBody("> ", request.Headers, request.Body), - formatHeadersAndBody("# ", response.Headers, response.Body), - ) - } - } - - for _, stub := range stubs { - require.NotEmpty(t, stub.Pattern) - items := strings.Split(stub.Pattern, " ") - require.Len(t, items, 2) - s.Handle(items[0], items[1], func(req testserver.Request) any { - time.Sleep(stub.Delay) - return stub.Response - }) - } - - // The earliest handlers take precedence, add default handlers last - if !cloud { - addDefaultHandlers(s) - } - - return s.URL -} - -type LoggedRequest struct { - Headers http.Header `json:"headers,omitempty"` - Method string `json:"method"` - Path string `json:"path"` - Body any `json:"body,omitempty"` - RawBody string `json:"raw_body,omitempty"` -} - -func getLoggedRequest(req *testserver.Request, includedHeaders []string) LoggedRequest { - result := LoggedRequest{ - Method: req.Method, - Path: req.URL.Path, - Headers: filterHeaders(req.Headers, includedHeaders), - } - - if json.Valid(req.Body) { - result.Body = json.RawMessage(req.Body) - } else { - result.RawBody = string(req.Body) - } - - return result -} - -// TODO CONTINUE: The proxy is working on aws-prod-ucws. Unblock other environments -// and especially azure-prod-ucws which seems to be setting -func filterHeaders(h http.Header, includedHeaders []string) http.Header { - headers := make(http.Header) - for k, v := range h { - if !slices.Contains(includedHeaders, k) { - continue - } - headers[k] = v - } - return headers -} - -func formatHeadersAndBody(prefix string, headers http.Header, body []byte) string { - var result []string - for key, values := range headers { - if len(values) == 1 { - result = append(result, fmt.Sprintf("%s%s: %s", prefix, key, values[0])) - } else { - result = append(result, fmt.Sprintf("%s%s: %s", prefix, key, values)) - } - } - if len(body) > 0 { - var s string - if utf8.Valid(body) { - s = string(body) - } else { - s = fmt.Sprintf("[Binary %d bytes]", len(body)) - } - s = strings.ReplaceAll(s, "\n", "\n"+prefix) - result = append(result, prefix+s) - } - return strings.Join(result, "\n") -} From b8b22d4a10c63c05a3de70f32cbf2115bf82f901 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 23 Apr 2025 13:32:31 +0200 Subject: [PATCH 08/57] fix lint --- libs/telemetry/logger_test.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/telemetry/logger_test.go b/libs/telemetry/logger_test.go index f0f2a10020..d71350e3be 100644 --- a/libs/telemetry/logger_test.go +++ b/libs/telemetry/logger_test.go @@ -14,7 +14,7 @@ import ( ) func TestTelemetryUploadRetriesOnPartialSuccess(t *testing.T) { - server := testserver.New(t) + server := testserver.New(t, false) t.Cleanup(server.Close) count := 0 @@ -57,7 +57,7 @@ func TestTelemetryUploadRetriesOnPartialSuccess(t *testing.T) { } func uploadRetriesFor(t *testing.T, statusCode int) { - server := testserver.New(t) + server := testserver.New(t, false) t.Cleanup(server.Close) count := 0 @@ -118,7 +118,7 @@ func TestTelemetryUploadRetriesForStatusCodes(t *testing.T) { } func TestTelemetryUploadMaxRetries(t *testing.T) { - server := testserver.New(t) + server := testserver.New(t, false) t.Cleanup(server.Close) count := 0 From 9dec7393f7427a659487c7173e28be499ce6f051 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 23 Apr 2025 23:22:04 +0200 Subject: [PATCH 09/57] mvp --- acceptance/acceptance_test.go | 73 +++++++++++++++---- acceptance/internal/prepare_server.go | 14 ++-- .../selftest/record_cloud/out.requests.txt | 4 + acceptance/selftest/record_cloud/output.txt | 10 ++- acceptance/selftest/record_cloud/script | 8 +- libs/testdiff/replacement.go | 39 +++++----- 6 files changed, 105 insertions(+), 43 deletions(-) create mode 100644 acceptance/selftest/record_cloud/out.requests.txt diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 9eb4c90140..e3d3b0f334 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -11,6 +11,7 @@ import ( "net/http" "os" "os/exec" + "path" "path/filepath" "regexp" "runtime" @@ -25,9 +26,7 @@ import ( "github.com/databricks/cli/acceptance/internal" "github.com/databricks/cli/internal/testutil" - "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/testdiff" - "github.com/databricks/databricks-sdk-go" "github.com/stretchr/testify/require" ) @@ -77,7 +76,41 @@ var Ignored = map[string]bool{ ReplsFile: true, } +// Detects if test is run from "debug test" feature in VS Code. +func IsInDebug() bool { + ex, _ := os.Executable() + return strings.HasPrefix(path.Base(ex), "__debug_bin") +} + +// Loads debug environment from ~/.databricks/debug-env.json. +func loadDebugEnvIfRunFromIDE(t testutil.TestingT, key string) { + if !IsInDebug() { + return + } + home, err := os.UserHomeDir() + if err != nil { + t.Fatalf("cannot find user home: %s", err) + } + raw, err := os.ReadFile(filepath.Join(home, ".databricks/debug-env.json")) + if err != nil { + t.Fatalf("cannot load ~/.databricks/debug-env.json: %s", err) + } + var conf map[string]map[string]string + err = json.Unmarshal(raw, &conf) + if err != nil { + t.Fatalf("cannot parse ~/.databricks/debug-env.json: %s", err) + } + vars, ok := conf[key] + if !ok { + t.Fatalf("~/.databricks/debug-env.json#%s not configured", key) + } + for k, v := range vars { + os.Setenv(k, v) + } +} + func TestAccept(t *testing.T) { + // loadDebugEnvIfRunFromIDE(t, "workspace") testAccept(t, InprocessMode, "selftest/record_cloud") } @@ -85,7 +118,8 @@ func TestInprocessMode(t *testing.T) { if InprocessMode && !Forcerun { t.Skip("Already tested by TestAccept") } - require.Equal(t, 1, testAccept(t, true, "selftest/basic")) + loadDebugEnvIfRunFromIDE(t, "CLOUD_ENV") + require.Equal(t, 1, testAccept(t, true, "selftest/record_cloud")) require.Equal(t, 1, testAccept(t, true, "selftest/server")) } @@ -355,27 +389,34 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont err = CopyDir(dir, tmpDir, inputs, outputs) require.NoError(t, err) + cfg, user, env := internal.PrepareServerAndClient(t, config, LogRequests, tmpDir) + testdiff.PrepareReplacementsUser(t, &repls, user) + testdiff.PrepareReplacementsWorkspaceConfig(t, &repls, cfg) + + if env != nil && InprocessMode { + testutil.NullEnvironment(t) + for _, kv := range env { + parts := strings.SplitN(kv, "=", 2) + require.Len(t, parts, 2) + t.Setenv(parts[0], parts[1]) + } + } + args := []string{"bash", "-euo", "pipefail", EntryPointScript} cmd := exec.Command(args[0], args[1:]...) cmd.Env = os.Environ() + if env != nil { + cmd.Env = env + } cmd.Env = append(cmd.Env, "UNIQUE_NAME="+uniqueName) cmd.Env = append(cmd.Env, "TEST_TMP_DIR="+tmpDir) - cfg, user := internal.PrepareServerAndClient(t, config, LogRequests, tmpDir) - testdiff.PrepareReplacementsUser(t, &repls, user) - - w, err := databricks.NewWorkspaceClient((*databricks.Config)(cfg)) - require.NoError(t, err) - testdiff.PrepareReplacementsWorkspaceClient(t, &repls, w) - // Configure resolved credentials in the environment. - for _, v := range auth.ProcessEnv(cfg) { - cmd.Env = append(cmd.Env, v) + cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+cfg.Host) + if cfg.Token != "" { + cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+cfg.Token) } - // cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+workspaceClient.Config.Host) - // if workspaceClient.Config.Token != "" { - // cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+workspaceClient.Config.Token) - // } + // Must be added PrepareReplacementsUser, otherwise conflicts with [USERNAME] testdiff.PrepareReplacementsUUID(t, &repls) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 7c9cdda74e..3e18f62f48 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -13,6 +13,7 @@ import ( "time" "unicode/utf8" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" @@ -39,7 +40,7 @@ func isTruePtr(value *bool) bool { return value != nil && *value } -func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User) { +func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User, []string) { cloudEnv := os.Getenv("CLOUD_ENV") recordRequests := isTruePtr(config.RecordRequests) @@ -66,7 +67,10 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o user, err := w.CurrentUser.Me(context.Background()) require.NoError(t, err) - return cfg, *user + // TODO: I only need to set these environment variables in hte inprocess mode. + // Support both child process and inprocess mode. + env := auth.ProcessEnv(cfg) + return cfg, *user, env } // If we are running on a cloud environment, use the host configured in the @@ -76,7 +80,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o err := cfg.EnsureResolved() require.NoError(t, err) - return cfg, TestUser + return cfg, TestUser, nil } // If we are not recording requests, and no custom server server stubs are configured, @@ -87,7 +91,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o Token: token, } - return cfg, TestUser + return cfg, TestUser, nil } // We want later stubs takes precedence, because then leaf configs take precedence over parent directory configs @@ -104,7 +108,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o // For the purposes of replacements, use testUser for local runs. // Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement: - return cfg, TestUser + return cfg, TestUser, nil } func startDedicatedServer(t *testing.T, diff --git a/acceptance/selftest/record_cloud/out.requests.txt b/acceptance/selftest/record_cloud/out.requests.txt new file mode 100644 index 0000000000..ec41b6c5f6 --- /dev/null +++ b/acceptance/selftest/record_cloud/out.requests.txt @@ -0,0 +1,4 @@ +{ + "method": "GET", + "path": "/api/2.0/preview/scim/v2/Me" +} diff --git a/acceptance/selftest/record_cloud/output.txt b/acceptance/selftest/record_cloud/output.txt index b5c12b09ce..2e339b53a5 100644 --- a/acceptance/selftest/record_cloud/output.txt +++ b/acceptance/selftest/record_cloud/output.txt @@ -1,5 +1,11 @@ >>> [CLI] current-user me -Error: inner token: Post "[ARM_TENANT_ID]/oauth2/token": Post "[ARM_TENANT_ID]/oauth2/token": unsupported protocol scheme "" +{ + "givenName": "[USERNAME]" +} -Exit code: 1 +>>> cat out.requests.txt +{ + "method": "GET", + "path": "/api/2.0/preview/scim/v2/Me" +} diff --git a/acceptance/selftest/record_cloud/script b/acceptance/selftest/record_cloud/script index 789a66fe3c..7e1c296a0d 100644 --- a/acceptance/selftest/record_cloud/script +++ b/acceptance/selftest/record_cloud/script @@ -1,7 +1,9 @@ trace $CLI current-user me | jq .name -# TODO CONTINUE(v3): I got it working for a short while but something went wrong again with the proxy. -# Figure out out and make this test pass after the new approach where workspaceClient -> config.Config. +# TODO CONTINUE(v4): I think it works fine for debugging (with the right env vars set in the currnet process). +# does not seem to work for the child process. trace cat out.requests.txt | jq 'select(has("path") and .path == "/api/2.0/preview/scim/v2/Me")' -rm out.requests.txt + +# TODO: Cleanup this test. +# rm out.requests.txt diff --git a/libs/testdiff/replacement.go b/libs/testdiff/replacement.go index d4d5eb27bb..b8add9b9a8 100644 --- a/libs/testdiff/replacement.go +++ b/libs/testdiff/replacement.go @@ -11,6 +11,7 @@ import ( "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/iamutil" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/service/iam" "golang.org/x/mod/semver" ) @@ -136,29 +137,33 @@ func (r *ReplacementsContext) SetPathWithParents(old, new string) { r.SetPath(filepath.Dir(filepath.Dir(old)), new+"_GPARENT") } +// TODO (followup): remove this function. func PrepareReplacementsWorkspaceClient(t testutil.TestingT, r *ReplacementsContext, w *databricks.WorkspaceClient) { + t.Helper() + PrepareReplacementsWorkspaceConfig(t, r, w.Config) +} + +func PrepareReplacementsWorkspaceConfig(t testutil.TestingT, r *ReplacementsContext, cfg *config.Config) { t.Helper() // in some clouds (gcp) w.Config.Host includes "https://" prefix in others it's really just a host (azure) - host := strings.TrimPrefix(strings.TrimPrefix(w.Config.Host, "http://"), "https://") + host := strings.TrimPrefix(strings.TrimPrefix(cfg.Host, "http://"), "https://") r.Set("https://"+host, "[DATABRICKS_URL]") r.Set("http://"+host, "[DATABRICKS_URL]") r.Set(host, "[DATABRICKS_HOST]") - r.Set(w.Config.ClusterID, "[DATABRICKS_CLUSTER_ID]") - r.Set(w.Config.WarehouseID, "[DATABRICKS_WAREHOUSE_ID]") - r.Set(w.Config.ServerlessComputeID, "[DATABRICKS_SERVERLESS_COMPUTE_ID]") - r.Set(w.Config.AccountID, "[DATABRICKS_ACCOUNT_ID]") - r.Set(w.Config.Username, "[DATABRICKS_USERNAME]") - r.SetPath(w.Config.Profile, "[DATABRICKS_CONFIG_PROFILE]") - r.Set(w.Config.ConfigFile, "[DATABRICKS_CONFIG_FILE]") - r.Set(w.Config.GoogleServiceAccount, "[DATABRICKS_GOOGLE_SERVICE_ACCOUNT]") - r.Set(w.Config.AzureResourceID, "[DATABRICKS_AZURE_RESOURCE_ID]") - r.Set(w.Config.AzureClientID, testerName) - r.Set(w.Config.AzureTenantID, "[ARM_TENANT_ID]") - r.Set(w.Config.AzureEnvironment, "[ARM_ENVIRONMENT]") - r.Set(w.Config.ClientID, "[DATABRICKS_CLIENT_ID]") - r.SetPath(w.Config.DatabricksCliPath, "[DATABRICKS_CLI_PATH]") - // This is set to words like "path" that happen too frequently - // r.Set(w.Config.AuthType, "[DATABRICKS_AUTH_TYPE]") + r.Set(cfg.ClusterID, "[DATABRICKS_CLUSTER_ID]") + r.Set(cfg.WarehouseID, "[DATABRICKS_WAREHOUSE_ID]") + r.Set(cfg.ServerlessComputeID, "[DATABRICKS_SERVERLESS_COMPUTE_ID]") + r.Set(cfg.AccountID, "[DATABRICKS_ACCOUNT_ID]") + r.Set(cfg.Username, "[DATABRICKS_USERNAME]") + r.SetPath(cfg.Profile, "[DATABRICKS_CONFIG_PROFILE]") + r.Set(cfg.ConfigFile, "[DATABRICKS_CONFIG_FILE]") + r.Set(cfg.GoogleServiceAccount, "[DATABRICKS_GOOGLE_SERVICE_ACCOUNT]") + r.Set(cfg.AzureResourceID, "[DATABRICKS_AZURE_RESOURCE_ID]") + r.Set(cfg.AzureClientID, testerName) + r.Set(cfg.AzureTenantID, "[ARM_TENANT_ID]") + r.Set(cfg.AzureEnvironment, "[ARM_ENVIRONMENT]") + r.Set(cfg.ClientID, "[DATABRICKS_CLIENT_ID]") + r.SetPath(cfg.DatabricksCliPath, "[DATABRICKS_CLI_PATH]") } func PrepareReplacementsUser(t testutil.TestingT, r *ReplacementsContext, u iam.User) { From 34e5159c44882df5259e079b6e5fd5c485f0ad72 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 24 Apr 2025 00:59:09 +0200 Subject: [PATCH 10/57] got this to a working state after refactor; TODO: cleanup and more testing of the feature itself --- acceptance/acceptance_test.go | 18 ++- acceptance/internal/cmd_server.go | 4 +- acceptance/internal/handlers.go | 2 +- acceptance/internal/prepare_server.go | 130 +++++++++------ acceptance/selftest/record_cloud/script | 2 + libs/telemetry/logger_test.go | 6 +- libs/testserver/local_server.go | 181 +++++++++++++++++++++ libs/testserver/proxy_server.go | 128 +++++++++++++++ libs/testserver/server.go | 204 +----------------------- 9 files changed, 419 insertions(+), 256 deletions(-) create mode 100644 libs/testserver/local_server.go create mode 100644 libs/testserver/proxy_server.go diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index e3d3b0f334..7a23c1146b 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -110,8 +110,9 @@ func loadDebugEnvIfRunFromIDE(t testutil.TestingT, key string) { } func TestAccept(t *testing.T) { + // TODO: Clean this up. // loadDebugEnvIfRunFromIDE(t, "workspace") - testAccept(t, InprocessMode, "selftest/record_cloud") + testAccept(t, InprocessMode, "") } func TestInprocessMode(t *testing.T) { @@ -253,7 +254,7 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int { if len(expanded) == 1 { // env vars aren't part of the test case name, so log them for debugging t.Logf("Running test with env %v", expanded[0]) - runTest(t, dir, coverDir, repls.Clone(), config, configPath, expanded[0]) + runTest(t, dir, coverDir, repls.Clone(), config, configPath, expanded[0], InprocessMode) } else { for _, envset := range expanded { envname := strings.Join(envset, "/") @@ -261,7 +262,7 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int { if !InprocessMode { t.Parallel() } - runTest(t, dir, coverDir, repls.Clone(), config, configPath, envset) + runTest(t, dir, coverDir, repls.Clone(), config, configPath, envset, InprocessMode) }) } } @@ -353,7 +354,14 @@ func getSkipReason(config *internal.TestConfig, configPath string) string { return "" } -func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsContext, config internal.TestConfig, configPath string, customEnv []string) { +func runTest(t *testing.T, + dir, coverDir string, + repls testdiff.ReplacementsContext, + config internal.TestConfig, + configPath string, + customEnv []string, + inprocessMode bool, +) { tailOutput := Tail cloudEnv := os.Getenv("CLOUD_ENV") isRunningOnCloud := cloudEnv != "" @@ -393,7 +401,7 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont testdiff.PrepareReplacementsUser(t, &repls, user) testdiff.PrepareReplacementsWorkspaceConfig(t, &repls, cfg) - if env != nil && InprocessMode { + if env != nil && inprocessMode { testutil.NullEnvironment(t) for _, kv := range env { parts := strings.SplitN(kv, "=", 2) diff --git a/acceptance/internal/cmd_server.go b/acceptance/internal/cmd_server.go index ba0083fd0e..c776a668f5 100644 --- a/acceptance/internal/cmd_server.go +++ b/acceptance/internal/cmd_server.go @@ -12,8 +12,8 @@ import ( "github.com/stretchr/testify/require" ) -func StartCmdServer(t *testing.T) *testserver.Server { - server := testserver.New(t, false) +func StartCmdServer(t *testing.T) *testserver.LocalServer { + server := testserver.NewLocalServer(t) server.Handle("GET", "/", func(r testserver.Request) any { q := r.URL.Query() args := strings.Split(q.Get("args"), " ") diff --git a/acceptance/internal/handlers.go b/acceptance/internal/handlers.go index 4f1665b027..64767ac6f8 100644 --- a/acceptance/internal/handlers.go +++ b/acceptance/internal/handlers.go @@ -28,7 +28,7 @@ var TestMetastore = catalog.MetastoreAssignment{ WorkspaceId: 470123456789500, } -func addDefaultHandlers(server *testserver.Server) { +func addDefaultHandlers(server *testserver.LocalServer) { server.Handle("GET", "/api/2.0/policies/clusters/list", func(req testserver.Request) any { return compute.ListPoliciesResponse{ Policies: []compute.Policy{ diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 3e18f62f48..6e1262b115 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -17,6 +17,7 @@ import ( "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" sdkconfig "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/google/uuid" @@ -26,7 +27,7 @@ import ( // TODO: Make CLI auth for with this. func StartDefaultServer(t *testing.T) { - s := testserver.New(t, false) + s := testserver.NewLocalServer(t) addDefaultHandlers(s) t.Setenv("DATABRICKS_DEFAULT_HOST", s.URL) @@ -40,25 +41,14 @@ func isTruePtr(value *bool) bool { return value != nil && *value } +// TODO: Clean up the env story here. func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User, []string) { cloudEnv := os.Getenv("CLOUD_ENV") recordRequests := isTruePtr(config.RecordRequests) - tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") - token := "dbapi" + tokenSuffix - // If we are running in a cloud environment AND we are recording requests, // start a dedicated server to act as a reverse proxy and record requests. if cloudEnv != "" && recordRequests { - host := startDedicatedServer(t, config.Server, recordRequests, - logRequests, config.IncludeRequestHeaders, outputDir, - true) - - cfg := &sdkconfig.Config{ - Host: host, - Token: token, - } - // Use a non-proxy client to fetch user info so that this API call is not recorded // in out.requests.txt w, err := databricks.NewWorkspaceClient() @@ -67,6 +57,9 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o user, err := w.CurrentUser.Me(context.Background()) require.NoError(t, err) + // Start a proxy server that sits in front of of a real Databricks workspace. + cfg := startProxyServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + // TODO: I only need to set these environment variables in hte inprocess mode. // Support both child process and inprocess mode. env := auth.ProcessEnv(cfg) @@ -86,6 +79,11 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o // If we are not recording requests, and no custom server server stubs are configured, // use the default shared server. if len(config.Server) == 0 && !recordRequests { + // Use a unique token for each test. This allows us to maintain + // separate state for each test in fake workspaces. + tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") + token := "dbapi" + tokenSuffix + cfg := &sdkconfig.Config{ Host: os.Getenv("DATABRICKS_DEFAULT_HOST"), Token: token, @@ -97,54 +95,57 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o // We want later stubs takes precedence, because then leaf configs take precedence over parent directory configs // In gorilla/mux earlier handlers take precedence, so we need to reverse the order slices.Reverse(config.Server) - host := startDedicatedServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir, false) - - cfg := &sdkconfig.Config{ - Host: host, - Token: token, - } - err := cfg.EnsureResolved() - require.NoError(t, err) + cfg := startLocalServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) // For the purposes of replacements, use testUser for local runs. // Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement: return cfg, TestUser, nil } -func startDedicatedServer(t *testing.T, +func recordRequestsCallback(t *testing.T, includeHeaders []string, outputDir string) func(request *testserver.Request) { + return func(request *testserver.Request) { + req := getLoggedRequest(request, includeHeaders) + reqJson, err := json.MarshalIndent(req, "", " ") + assert.NoErrorf(t, err, "Failed to json-encode: %#v", req) + + requestsPath := filepath.Join(outputDir, "out.requests.txt") + f, err := os.OpenFile(requestsPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) + assert.NoError(t, err) + defer f.Close() + + _, err = f.WriteString(string(reqJson) + "\n") + assert.NoError(t, err) + } +} + +func logResponseCallback(t *testing.T) func(request *testserver.Request, response *testserver.EncodedResponse) { + return func(request *testserver.Request, response *testserver.EncodedResponse) { + t.Logf("%d %s %s\n%s\n%s", + response.StatusCode, request.Method, request.URL, + formatHeadersAndBody("> ", request.Headers, request.Body), + formatHeadersAndBody("# ", response.Headers, response.Body), + ) + } +} + +func startLocalServer(t *testing.T, stubs []ServerStub, recordRequests bool, logRequests bool, includeHeaders []string, outputDir string, - cloud bool, -) string { - s := testserver.New(t, cloud) +) *config.Config { + s := testserver.NewLocalServer(t) + // Record API requests in out.requests.txt if RecordRequests is true + // in test.toml if recordRequests { - requestsPath := filepath.Join(outputDir, "out.requests.txt") - s.RequestCallback = func(request *testserver.Request) { - req := getLoggedRequest(request, includeHeaders) - reqJson, err := json.MarshalIndent(req, "", " ") - assert.NoErrorf(t, err, "Failed to json-encode: %#v", req) - - f, err := os.OpenFile(requestsPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) - assert.NoError(t, err) - defer f.Close() - - _, err = f.WriteString(string(reqJson) + "\n") - assert.NoError(t, err) - } + s.SetRequestCallback(recordRequestsCallback(t, includeHeaders, outputDir)) } + // Log API responses if the -logrequests flag is set. if logRequests { - s.ResponseCallback = func(request *testserver.Request, response *testserver.EncodedResponse) { - t.Logf("%d %s %s\n%s\n%s", - response.StatusCode, request.Method, request.URL, - formatHeadersAndBody("> ", request.Headers, request.Body), - formatHeadersAndBody("# ", response.Headers, response.Body), - ) - } + s.SetResponseCallback(logResponseCallback(t)) } for _, stub := range stubs { @@ -158,11 +159,44 @@ func startDedicatedServer(t *testing.T, } // The earliest handlers take precedence, add default handlers last - if !cloud { - addDefaultHandlers(s) + addDefaultHandlers(s) + + cfg := &sdkconfig.Config{ + Host: s.URL, + Token: "dbapi123", + } + err := cfg.EnsureResolved() + require.NoError(t, err) + + return cfg +} + +func startProxyServer(t *testing.T, + stubs []ServerStub, + recordRequests bool, + logRequests bool, + includeHeaders []string, + outputDir string, +) *config.Config { + s := testserver.NewProxyServer(t) + + // Record API requests in out.requests.txt if RecordRequests is true + // in test.toml + if recordRequests { + s.SetRequestCallback(recordRequestsCallback(t, includeHeaders, outputDir)) + } + + // Log API responses if the -logrequests flag is set. + if logRequests { + s.SetResponseCallback(logResponseCallback(t)) + } + + cfg := &sdkconfig.Config{ + Host: s.URL, + Token: "dbapi1234", } - return s.URL + return cfg } type LoggedRequest struct { diff --git a/acceptance/selftest/record_cloud/script b/acceptance/selftest/record_cloud/script index 7e1c296a0d..6cabcd9f84 100644 --- a/acceptance/selftest/record_cloud/script +++ b/acceptance/selftest/record_cloud/script @@ -1,3 +1,5 @@ +# trace env | sort + trace $CLI current-user me | jq .name # TODO CONTINUE(v4): I think it works fine for debugging (with the right env vars set in the currnet process). diff --git a/libs/telemetry/logger_test.go b/libs/telemetry/logger_test.go index d71350e3be..cda3f7724b 100644 --- a/libs/telemetry/logger_test.go +++ b/libs/telemetry/logger_test.go @@ -14,7 +14,7 @@ import ( ) func TestTelemetryUploadRetriesOnPartialSuccess(t *testing.T) { - server := testserver.New(t, false) + server := testserver.NewLocalServer(t) t.Cleanup(server.Close) count := 0 @@ -57,7 +57,7 @@ func TestTelemetryUploadRetriesOnPartialSuccess(t *testing.T) { } func uploadRetriesFor(t *testing.T, statusCode int) { - server := testserver.New(t, false) + server := testserver.NewLocalServer(t) t.Cleanup(server.Close) count := 0 @@ -118,7 +118,7 @@ func TestTelemetryUploadRetriesForStatusCodes(t *testing.T) { } func TestTelemetryUploadMaxRetries(t *testing.T) { - server := testserver.New(t, false) + server := testserver.NewLocalServer(t) t.Cleanup(server.Close) count := 0 diff --git a/libs/testserver/local_server.go b/libs/testserver/local_server.go new file mode 100644 index 0000000000..857751293a --- /dev/null +++ b/libs/testserver/local_server.go @@ -0,0 +1,181 @@ +package testserver + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "sync" + + "github.com/gorilla/mux" + + "github.com/databricks/cli/internal/testutil" + "github.com/databricks/databricks-sdk-go/apierr" +) + +type LocalServer struct { + *httptest.Server + + t testutil.TestingT + mu *sync.Mutex + router *mux.Router + + fakeWorkspaces map[string]*FakeWorkspace + requestCallback func(request *Request) + responseCallback func(request *Request, response *EncodedResponse) +} + +func NewLocalServer(t testutil.TestingT) *LocalServer { + router := mux.NewRouter() + server := httptest.NewServer(router) + t.Cleanup(server.Close) + + s := &LocalServer{ + Server: server, + router: router, + t: t, + mu: &sync.Mutex{}, + fakeWorkspaces: map[string]*FakeWorkspace{}, + } + + // Set up the not found handler as fallback. + s.router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + pattern := r.Method + " " + r.URL.Path + bodyBytes, err := io.ReadAll(r.Body) + var body string + if err != nil { + body = fmt.Sprintf("failed to read the body: %s", err) + } else { + body = fmt.Sprintf("[%d bytes] %s", len(bodyBytes), bodyBytes) + } + + t.Errorf(`No handler for URL: %s +Body: %s + +For acceptance tests, add this to test.toml: +[[Server]] +Pattern = %q +Response.Body = '' +# Response.StatusCode = +`, r.URL, body, pattern) + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotImplemented) + + resp := apierr.APIError{ + Message: "No stub found for pattern: " + pattern, + } + + respBytes, err := json.Marshal(resp) + if err != nil { + t.Errorf("JSON encoding error: %s", err) + respBytes = []byte("{\"message\": \"JSON encoding error\"}") + } + + if _, err := w.Write(respBytes); err != nil { + t.Errorf("Response write error: %s", err) + } + }) + + return s +} + +func (s *LocalServer) Handle(method, path string, handler HandlerFunc) { + s.router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + // For simplicity we process requests sequentially. It's fast enough because + // we don't do any IO except reading and writing request/response bodies. + s.mu.Lock() + defer s.mu.Unlock() + + // Each test uses unique DATABRICKS_TOKEN, we simulate each token having + // it's own fake fakeWorkspace to avoid interference between tests. + var fakeWorkspace *FakeWorkspace = nil + token := getToken(r) + if token != "" { + if _, ok := s.fakeWorkspaces[token]; !ok { + s.fakeWorkspaces[token] = NewFakeWorkspace() + } + + fakeWorkspace = s.fakeWorkspaces[token] + } + + request := NewRequest(s.t, r, fakeWorkspace) + + if s.requestCallback != nil { + s.requestCallback(&request) + } + + respAny := handler(request) + resp := normalizeResponse(s.t, respAny) + + for k, v := range resp.Headers { + w.Header()[k] = v + } + + w.WriteHeader(resp.StatusCode) + + if s.responseCallback != nil { + s.responseCallback(&request, &resp) + } + + if _, err := w.Write(resp.Body); err != nil { + s.t.Errorf("Failed to write response: %s", err) + return + } + }).Methods(method) +} + +func (s *LocalServer) SetRequestCallback(callback func(request *Request)) { + s.requestCallback = callback +} + +func (s *LocalServer) SetResponseCallback(callback func(request *Request, response *EncodedResponse)) { + s.responseCallback = callback +} + +// TODO: Remove this. +// func (s *LocalServer) ProxyToCloud(w http.ResponseWriter, r *http.Request) { +// request := NewRequest(s.t, r, nil) +// s.RequestCallback(&request) + +// headers := make(map[string]string) +// for k, v := range r.Header { +// // Authorization headers will be set by the SDK. No need to pass them along here. +// if k == "Authorization" { +// continue +// } +// if k == "Accept-Encoding" { +// continue +// } +// headers[k] = v[0] +// } + +// queryParams := make(map[string]any) +// for k, v := range r.URL.Query() { +// queryParams[k] = v[0] +// } + +// // TODO: Since the response is always JSON, this should be specified in the header. +// respB := map[string]any{} +// err := s.proxyClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, r.Body, &respB) +// require.NoError(s.t, err) // todo remove +// if err != nil { +// // API errors from the SDK are expected to be of the type apierr.APIError. +// apiErr := apierr.APIError{} +// if errors.As(err, &apiErr) { +// w.WriteHeader(apiErr.StatusCode) +// w.Write(respB["message"].([]byte)) +// } else { +// // Something else went wrong. +// w.WriteHeader(http.StatusInternalServerError) +// w.Write([]byte(err.Error())) +// } +// } + +// // Successful response +// w.WriteHeader(200) +// b, err := json.Marshal(respB) +// require.NoError(s.t, err) +// w.Write(b) +// } diff --git a/libs/testserver/proxy_server.go b/libs/testserver/proxy_server.go new file mode 100644 index 0000000000..b31ec0c817 --- /dev/null +++ b/libs/testserver/proxy_server.go @@ -0,0 +1,128 @@ +package testserver + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "sync" + + "github.com/databricks/cli/internal/testutil" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/config" + "github.com/gorilla/mux" + "github.com/stretchr/testify/require" +) + +type ProxyServer struct { + *httptest.Server + + t testutil.TestingT + mu *sync.Mutex + router *mux.Router + + apiClient *client.DatabricksClient + requestCallback func(request *Request) + responseCallback func(request *Request, response *EncodedResponse) +} + +func NewProxyServer(t testutil.TestingT) *ProxyServer { + router := mux.NewRouter() + server := httptest.NewServer(router) + t.Cleanup(server.Close) + + s := &ProxyServer{ + Server: server, + t: t, + router: router, + mu: &sync.Mutex{}, + } + + // Create an API client using the current authentication context. + // In CI test environments this would read the appropriate environment + // variables. + var err error + s.apiClient, err = client.New(&config.Config{}) + require.NoError(t, err) + + router.NotFoundHandler = http.HandlerFunc(s.proxyToCloud) + return s +} + +// TODO: Iterate once on this function. +func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { + // Process requests sequentially. It's slower but is easier to reason about. + // For example, requestCallback and responseCallback functions do not need + // to be thread-safe. + s.mu.Lock() + defer s.mu.Unlock() + + request := NewRequest(s.t, r, nil) + if s.requestCallback != nil { + s.requestCallback(&request) + } + + headers := make(map[string]string) + for k, v := range r.Header { + // Authorization headers will be set by the SDK. No need to pass them along here. + if k == "Authorization" { + continue + } + if k == "Accept-Encoding" { + continue + } + headers[k] = v[0] + } + + queryParams := make(map[string]any) + for k, v := range r.URL.Query() { + queryParams[k] = v[0] + } + + // TODO: Since the response is always JSON, this should be specified in the header. + respB := map[string]any{} + err := s.apiClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, r.Body, &respB) + require.NoError(s.t, err) // todo remove + if err != nil { + // API errors from the SDK are expected to be of the type apierr.APIError. + apiErr := &apierr.APIError{} + if errors.As(err, &apiErr) { + w.WriteHeader(apiErr.StatusCode) + w.Write(respB["message"].([]byte)) + } else { + // Something else went wrong. + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(err.Error())) + } + } + + // Successful response + w.WriteHeader(200) + b, err := json.Marshal(respB) + require.NoError(s.t, err) + + if s.responseCallback != nil { + s.responseCallback(&request, &EncodedResponse{ + StatusCode: 200, + Body: b, + }) + } + + w.Write(b) +} + +// Eventually we can implement this function to allow for per-test overrides +// even in integration tests. +func (s *ProxyServer) Handle(method, path string, handler HandlerFunc) { + require.FailNow(s.t, "Not implemented") +} + +func (s *ProxyServer) SetRequestCallback(callback func(request *Request)) { + s.requestCallback = callback +} + +func (s *ProxyServer) SetResponseCallback(callback func(request *Request, response *EncodedResponse)) { + s.responseCallback = callback +} diff --git a/libs/testserver/server.go b/libs/testserver/server.go index b4b6f4da9a..d864ff4ec8 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -1,44 +1,25 @@ package testserver import ( - "context" "encoding/json" - "errors" - "fmt" "io" "net/http" - "net/http/httptest" "net/url" - "os" "reflect" "strings" - "sync" - - "github.com/gorilla/mux" "github.com/databricks/cli/internal/testutil" - "github.com/databricks/databricks-sdk-go/apierr" - "github.com/databricks/databricks-sdk-go/client" - "github.com/databricks/databricks-sdk-go/config" - "github.com/stretchr/testify/require" + "github.com/gorilla/mux" ) -type Server struct { - *httptest.Server - Router *mux.Router - - t testutil.TestingT - - workspaceUrl string - proxyClient *client.DatabricksClient - - fakeWorkspaces map[string]*FakeWorkspace - mu *sync.Mutex - - RequestCallback func(request *Request) - ResponseCallback func(request *Request, response *EncodedResponse) +type Server interface { + Handle(method, path string, handler HandlerFunc) + SetRequestCallback(func(request *Request)) + SetResponseCallback(func(request *Request, response *EncodedResponse)) } +type HandlerFunc func(req Request) any + type Request struct { Method string URL *url.URL @@ -186,177 +167,6 @@ func getHeaders(value []byte) http.Header { } } -func New(t testutil.TestingT, cloud bool) *Server { - router := mux.NewRouter() - server := httptest.NewServer(router) - t.Cleanup(server.Close) - - s := &Server{ - Server: server, - Router: router, - t: t, - mu: &sync.Mutex{}, - fakeWorkspaces: map[string]*FakeWorkspace{}, - } - - // If cloud is true, we proxy all requests to the configured Databricks workspace - // instead of using local server stubs. - var err error - if cloud { - s.workspaceUrl = os.Getenv("DATABRICKS_HOST") - s.proxyClient, err = client.New(&config.Config{}) - require.NoError(t, err) - } - - // Set up the not found handler as fallback. If a proxy URL is configured then this handler - // will proxy the request to the proxy URL. - s.Router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // For integration tests forwards all requests to the workspace instead of returning a 404. - if s.proxyClient != nil { - s.ProxyToCloud(w, r) - return - } - - // Original Not Found Logic - pattern := r.Method + " " + r.URL.Path - bodyBytes, err := io.ReadAll(r.Body) - var body string - if err != nil { - body = fmt.Sprintf("failed to read the body: %s", err) - } else { - body = fmt.Sprintf("[%d bytes] %s", len(bodyBytes), bodyBytes) - } - - t.Errorf(`No handler for URL: %s -Body: %s - -For acceptance tests, add this to test.toml: -[[Server]] -Pattern = %q -Response.Body = '' -# Response.StatusCode = -`, r.URL, body, pattern) - - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusNotImplemented) - - resp := apierr.APIError{ - Message: "No stub found for pattern: " + pattern, - } - - respBytes, err := json.Marshal(resp) - if err != nil { - t.Errorf("JSON encoding error: %s", err) - respBytes = []byte("{\"message\": \"JSON encoding error\"}") - } - - if _, err := w.Write(respBytes); err != nil { - t.Errorf("Response write error: %s", err) - } - }) - - return s -} - -type HandlerFunc func(req Request) any - -func (s *Server) ProxyToCloud(w http.ResponseWriter, r *http.Request) { - request := NewRequest(s.t, r, nil) - s.RequestCallback(&request) - - headers := make(map[string]string) - for k, v := range r.Header { - // Authorization headers will be set by the SDK. No need to pass them along here. - if k == "Authorization" { - continue - } - if k == "Accept-Encoding" { - continue - } - headers[k] = v[0] - } - - queryParams := make(map[string]any) - for k, v := range r.URL.Query() { - queryParams[k] = v[0] - } - - // TODO: Since the response is always JSON, this should be specified in the header. - respB := map[string]any{} - err := s.proxyClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, r.Body, &respB) - require.NoError(s.t, err) // todo remove - if err != nil { - // API errors from the SDK are expected to be of the type apierr.APIError. - apiErr := apierr.APIError{} - if errors.As(err, &apiErr) { - w.WriteHeader(apiErr.StatusCode) - w.Write(respB["message"].([]byte)) - } else { - // Something else went wrong. - w.WriteHeader(http.StatusInternalServerError) - w.Write([]byte(err.Error())) - } - } - - // Successful response - w.WriteHeader(200) - b, err := json.Marshal(respB) - require.NoError(s.t, err) - w.Write(b) -} - -func (s *Server) Handle(method, path string, handler HandlerFunc) { - s.Router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { - // Overriding the proxy URL is not supported. We can add a configuration option - // to disable the proxy and use the test.toml stub if and when such a use case arises. - if s.proxyClient != nil { - s.ProxyToCloud(w, r) - return - } - - // For simplicity we process requests sequentially. It's fast enough because - // we don't do any IO except reading and writing request/response bodies. - s.mu.Lock() - defer s.mu.Unlock() - - // Each test uses unique DATABRICKS_TOKEN, we simulate each token having - // it's own fake fakeWorkspace to avoid interference between tests. - var fakeWorkspace *FakeWorkspace = nil - token := getToken(r) - if token != "" { - if _, ok := s.fakeWorkspaces[token]; !ok { - s.fakeWorkspaces[token] = NewFakeWorkspace() - } - - fakeWorkspace = s.fakeWorkspaces[token] - } - - request := NewRequest(s.t, r, fakeWorkspace) - - if s.RequestCallback != nil { - s.RequestCallback(&request) - } - - respAny := handler(request) - resp := normalizeResponse(s.t, respAny) - - for k, v := range resp.Headers { - w.Header()[k] = v - } - - w.WriteHeader(resp.StatusCode) - - if s.ResponseCallback != nil { - s.ResponseCallback(&request, &resp) - } - - if _, err := w.Write(resp.Body); err != nil { - s.t.Errorf("Failed to write response: %s", err) - return - } - }).Methods(method) -} - func getToken(r *http.Request) string { header := r.Header.Get("Authorization") prefix := "Bearer " From a030b3b879e47f3ff0e3386030ea20641acda58a Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 24 Apr 2025 01:05:57 +0200 Subject: [PATCH 11/57] lint --- acceptance/internal/prepare_server.go | 5 ++--- libs/testserver/proxy_server.go | 16 ++++++++++------ 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 6e1262b115..cd7855951d 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -17,7 +17,6 @@ import ( "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" - "github.com/databricks/databricks-sdk-go/config" sdkconfig "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/google/uuid" @@ -134,7 +133,7 @@ func startLocalServer(t *testing.T, logRequests bool, includeHeaders []string, outputDir string, -) *config.Config { +) *sdkconfig.Config { s := testserver.NewLocalServer(t) // Record API requests in out.requests.txt if RecordRequests is true @@ -177,7 +176,7 @@ func startProxyServer(t *testing.T, logRequests bool, includeHeaders []string, outputDir string, -) *config.Config { +) *sdkconfig.Config { s := testserver.NewProxyServer(t) // Record API requests in out.requests.txt if RecordRequests is true diff --git a/libs/testserver/proxy_server.go b/libs/testserver/proxy_server.go index b31ec0c817..c937040953 100644 --- a/libs/testserver/proxy_server.go +++ b/libs/testserver/proxy_server.go @@ -13,6 +13,7 @@ import ( "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/config" "github.com/gorilla/mux" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -84,24 +85,26 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { // TODO: Since the response is always JSON, this should be specified in the header. respB := map[string]any{} err := s.apiClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, r.Body, &respB) - require.NoError(s.t, err) // todo remove + assert.NoError(s.t, err) if err != nil { // API errors from the SDK are expected to be of the type apierr.APIError. apiErr := &apierr.APIError{} if errors.As(err, &apiErr) { w.WriteHeader(apiErr.StatusCode) - w.Write(respB["message"].([]byte)) + _, err := w.Write(respB["message"].([]byte)) + assert.NoError(s.t, err) } else { // Something else went wrong. w.WriteHeader(http.StatusInternalServerError) - w.Write([]byte(err.Error())) + _, err := w.Write([]byte(err.Error())) + assert.NoError(s.t, err) } } // Successful response w.WriteHeader(200) b, err := json.Marshal(respB) - require.NoError(s.t, err) + assert.NoError(s.t, err) if s.responseCallback != nil { s.responseCallback(&request, &EncodedResponse{ @@ -110,13 +113,14 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { }) } - w.Write(b) + _, err = w.Write(b) + assert.NoError(s.t, err) } // Eventually we can implement this function to allow for per-test overrides // even in integration tests. func (s *ProxyServer) Handle(method, path string, handler HandlerFunc) { - require.FailNow(s.t, "Not implemented") + s.t.Fatalf("Not implemented") } func (s *ProxyServer) SetRequestCallback(callback func(request *Request)) { From cb0fde7ca66e972ead4b8c72a02ca2517f280cf2 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 24 Apr 2025 14:00:42 +0200 Subject: [PATCH 12/57] fix integration test users --- acceptance/internal/prepare_server.go | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index a75e93def3..90fc89b27f 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -54,7 +54,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o require.NoError(t, err) user, err := w.CurrentUser.Me(context.Background()) - require.NoError(t, err) + require.NoError(t, err, "Failed to get current user") // Start a proxy server that sits in front of of a real Databricks workspace. cfg := startProxyServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) @@ -68,11 +68,15 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o // If we are running on a cloud environment, use the host configured in the // environment. if cloudEnv != "" { - cfg := &sdkconfig.Config{} - err := cfg.EnsureResolved() + // Use a non-proxy client to fetch user info so that this API call is not recorded + // in out.requests.txt + w, err := databricks.NewWorkspaceClient() require.NoError(t, err) - return cfg, TestUser, nil + user, err := w.CurrentUser.Me(context.Background()) + require.NoError(t, err, "Failed to get current user") + + return w.Config, *user, nil } // If we are not recording requests, and no custom server server stubs are configured, From 12ed68365b303b21f404fa323eeb48269f397924 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 24 Apr 2025 14:32:17 +0200 Subject: [PATCH 13/57] - --- acceptance/acceptance_test.go | 68 ++++++++++----------- acceptance/selftest/record_cloud/output.txt | 10 +-- acceptance/selftest/record_cloud/script | 12 +--- 3 files changed, 35 insertions(+), 55 deletions(-) diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 750622a428..fca590e073 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -11,7 +11,6 @@ import ( "net/http" "os" "os/exec" - "path" "path/filepath" "regexp" "runtime" @@ -78,50 +77,49 @@ var Ignored = map[string]bool{ ReplsFile: true, } -// Detects if test is run from "debug test" feature in VS Code. -func IsInDebug() bool { - ex, _ := os.Executable() - return strings.HasPrefix(path.Base(ex), "__debug_bin") -} - -// Loads debug environment from ~/.databricks/debug-env.json. -func loadDebugEnvIfRunFromIDE(t testutil.TestingT, key string) { - if !IsInDebug() { - return - } - home, err := os.UserHomeDir() - if err != nil { - t.Fatalf("cannot find user home: %s", err) - } - raw, err := os.ReadFile(filepath.Join(home, ".databricks/debug-env.json")) - if err != nil { - t.Fatalf("cannot load ~/.databricks/debug-env.json: %s", err) - } - var conf map[string]map[string]string - err = json.Unmarshal(raw, &conf) - if err != nil { - t.Fatalf("cannot parse ~/.databricks/debug-env.json: %s", err) - } - vars, ok := conf[key] - if !ok { - t.Fatalf("~/.databricks/debug-env.json#%s not configured", key) - } - for k, v := range vars { - os.Setenv(k, v) - } -} +// // Detects if test is run from "debug test" feature in VS Code. +// func IsInDebug() bool { +// ex, _ := os.Executable() +// return strings.HasPrefix(path.Base(ex), "__debug_bin") +// } + +// // Loads debug environment from ~/.databricks/debug-env.json. +// func loadDebugEnvIfRunFromIDE(t testutil.TestingT, key string) { +// if !IsInDebug() { +// return +// } +// home, err := os.UserHomeDir() +// if err != nil { +// t.Fatalf("cannot find user home: %s", err) +// } +// raw, err := os.ReadFile(filepath.Join(home, ".databricks/debug-env.json")) +// if err != nil { +// t.Fatalf("cannot load ~/.databricks/debug-env.json: %s", err) +// } +// var conf map[string]map[string]string +// err = json.Unmarshal(raw, &conf) +// if err != nil { +// t.Fatalf("cannot parse ~/.databricks/debug-env.json: %s", err) +// } +// vars, ok := conf[key] +// if !ok { +// t.Fatalf("~/.databricks/debug-env.json#%s not configured", key) +// } +// for k, v := range vars { +// os.Setenv(k, v) +// } +// } func TestAccept(t *testing.T) { // TODO: Clean this up. // loadDebugEnvIfRunFromIDE(t, "workspace") - testAccept(t, InprocessMode, "") + testAccept(t, InprocessMode, "selftest/record_cloud") } func TestInprocessMode(t *testing.T) { if InprocessMode && !Forcerun { t.Skip("Already tested by TestAccept") } - loadDebugEnvIfRunFromIDE(t, "CLOUD_ENV") require.Equal(t, 1, testAccept(t, true, "selftest/record_cloud")) require.Equal(t, 1, testAccept(t, true, "selftest/server")) } diff --git a/acceptance/selftest/record_cloud/output.txt b/acceptance/selftest/record_cloud/output.txt index 2e339b53a5..7f0b941262 100644 --- a/acceptance/selftest/record_cloud/output.txt +++ b/acceptance/selftest/record_cloud/output.txt @@ -1,11 +1,3 @@ >>> [CLI] current-user me -{ - "givenName": "[USERNAME]" -} - ->>> cat out.requests.txt -{ - "method": "GET", - "path": "/api/2.0/preview/scim/v2/Me" -} +"[USERNAME]" diff --git a/acceptance/selftest/record_cloud/script b/acceptance/selftest/record_cloud/script index 6cabcd9f84..8c6c5d9be7 100644 --- a/acceptance/selftest/record_cloud/script +++ b/acceptance/selftest/record_cloud/script @@ -1,11 +1 @@ -# trace env | sort - -trace $CLI current-user me | jq .name - -# TODO CONTINUE(v4): I think it works fine for debugging (with the right env vars set in the currnet process). -# does not seem to work for the child process. -trace cat out.requests.txt | jq 'select(has("path") and .path == "/api/2.0/preview/scim/v2/Me")' - - -# TODO: Cleanup this test. -# rm out.requests.txt +trace $CLI current-user me | jq .name.givenName From d1c13dd72aecb5d862793d2f86b12683675d8f4e Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 24 Apr 2025 15:23:33 +0200 Subject: [PATCH 14/57] cleanup --- acceptance/acceptance_test.go | 2 +- acceptance/internal/prepare_server.go | 2 -- libs/testserver/local_server.go | 46 --------------------------- 3 files changed, 1 insertion(+), 49 deletions(-) diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index fca590e073..04488f99af 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -120,7 +120,7 @@ func TestInprocessMode(t *testing.T) { if InprocessMode && !Forcerun { t.Skip("Already tested by TestAccept") } - require.Equal(t, 1, testAccept(t, true, "selftest/record_cloud")) + require.Equal(t, 1, testAccept(t, true, "selftest/basic")) require.Equal(t, 1, testAccept(t, true, "selftest/server")) } diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 90fc89b27f..4b7cf1e7ed 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -226,8 +226,6 @@ func getLoggedRequest(req *testserver.Request, includedHeaders []string) LoggedR return result } -// TODO CONTINUE: The proxy is working on aws-prod-ucws. Unblock other environments -// and especially azure-prod-ucws which seems to be setting func filterHeaders(h http.Header, includedHeaders []string) http.Header { headers := make(http.Header) for k, v := range h { diff --git a/libs/testserver/local_server.go b/libs/testserver/local_server.go index 857751293a..a429a8c4fc 100644 --- a/libs/testserver/local_server.go +++ b/libs/testserver/local_server.go @@ -133,49 +133,3 @@ func (s *LocalServer) SetRequestCallback(callback func(request *Request)) { func (s *LocalServer) SetResponseCallback(callback func(request *Request, response *EncodedResponse)) { s.responseCallback = callback } - -// TODO: Remove this. -// func (s *LocalServer) ProxyToCloud(w http.ResponseWriter, r *http.Request) { -// request := NewRequest(s.t, r, nil) -// s.RequestCallback(&request) - -// headers := make(map[string]string) -// for k, v := range r.Header { -// // Authorization headers will be set by the SDK. No need to pass them along here. -// if k == "Authorization" { -// continue -// } -// if k == "Accept-Encoding" { -// continue -// } -// headers[k] = v[0] -// } - -// queryParams := make(map[string]any) -// for k, v := range r.URL.Query() { -// queryParams[k] = v[0] -// } - -// // TODO: Since the response is always JSON, this should be specified in the header. -// respB := map[string]any{} -// err := s.proxyClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, r.Body, &respB) -// require.NoError(s.t, err) // todo remove -// if err != nil { -// // API errors from the SDK are expected to be of the type apierr.APIError. -// apiErr := apierr.APIError{} -// if errors.As(err, &apiErr) { -// w.WriteHeader(apiErr.StatusCode) -// w.Write(respB["message"].([]byte)) -// } else { -// // Something else went wrong. -// w.WriteHeader(http.StatusInternalServerError) -// w.Write([]byte(err.Error())) -// } -// } - -// // Successful response -// w.WriteHeader(200) -// b, err := json.Marshal(respB) -// require.NoError(s.t, err) -// w.Write(b) -// } From 7ddb576a7dfed5eb0c537d7e317c3e88ca5023fc Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 24 Apr 2025 19:02:15 +0200 Subject: [PATCH 15/57] add crud test --- .../selftest/record_cloud/out.requests.txt | 51 +++++++++ acceptance/selftest/record_cloud/output.txt | 27 +++++ .../selftest/record_cloud/pipeline1.json | 11 ++ .../selftest/record_cloud/pipeline2.json | 11 ++ acceptance/selftest/record_cloud/script | 24 ++++ libs/testserver/proxy_server.go | 105 ++++++++++++++---- 6 files changed, 210 insertions(+), 19 deletions(-) create mode 100644 acceptance/selftest/record_cloud/pipeline1.json create mode 100644 acceptance/selftest/record_cloud/pipeline2.json diff --git a/acceptance/selftest/record_cloud/out.requests.txt b/acceptance/selftest/record_cloud/out.requests.txt index ec41b6c5f6..ac245e46ac 100644 --- a/acceptance/selftest/record_cloud/out.requests.txt +++ b/acceptance/selftest/record_cloud/out.requests.txt @@ -2,3 +2,54 @@ "method": "GET", "path": "/api/2.0/preview/scim/v2/Me" } +{ + "method": "GET", + "path": "/api/2.2/jobs/get" +} +{ + "method": "POST", + "path": "/api/2.0/pipelines", + "body": { + "allow_duplicate_names": true, + "libraries": [ + { + "file": { + "path": "/whatever.py" + } + } + ], + "name": "test-pipeline-1" + } +} +{ + "method": "GET", + "path": "/api/2.0/pipelines/[UUID]" +} +{ + "method": "PUT", + "path": "/api/2.0/pipelines/[UUID]", + "body": { + "allow_duplicate_names": true, + "libraries": [ + { + "file": { + "path": "/whatever.py" + } + } + ], + "name": "test-pipeline-2", + "pipeline_id": "[UUID]" + } +} +{ + "method": "GET", + "path": "/api/2.0/pipelines/[UUID]" +} +{ + "method": "DELETE", + "path": "/api/2.0/pipelines/[UUID]" +} +{ + "method": "GET", + "path": "/api/2.0/pipelines/[UUID]" +} diff --git a/acceptance/selftest/record_cloud/output.txt b/acceptance/selftest/record_cloud/output.txt index 7f0b941262..d80033ce81 100644 --- a/acceptance/selftest/record_cloud/output.txt +++ b/acceptance/selftest/record_cloud/output.txt @@ -1,3 +1,30 @@ >>> [CLI] current-user me "[USERNAME]" + +>>> [CLI] jobs get 1234 +Error: Job 1234 does not exist. + +Exit code: 1 + +=== Create a pipeline + +=== Get the pipeline +>>> [CLI] pipelines get [UUID] +"test-pipeline-1" + +=== Update the pipeline +>>> [CLI] pipelines update [UUID] --json @pipeline2.json + +=== Verify the update +>>> [CLI] pipelines get [UUID] +"test-pipeline-2" + +=== Delete the pipeline +>>> [CLI] pipelines delete [UUID] + +=== Verify the deletion +>>> [CLI] pipelines get [UUID] +Error: The specified pipeline [UUID] was not found. + +Exit code: 1 diff --git a/acceptance/selftest/record_cloud/pipeline1.json b/acceptance/selftest/record_cloud/pipeline1.json new file mode 100644 index 0000000000..dae02d2fd5 --- /dev/null +++ b/acceptance/selftest/record_cloud/pipeline1.json @@ -0,0 +1,11 @@ +{ + "name": "test-pipeline-1", + "allow_duplicate_names": true, + "libraries": [ + { + "file": { + "path": "/whatever.py" + } + } + ] +} diff --git a/acceptance/selftest/record_cloud/pipeline2.json b/acceptance/selftest/record_cloud/pipeline2.json new file mode 100644 index 0000000000..61f87b1f7a --- /dev/null +++ b/acceptance/selftest/record_cloud/pipeline2.json @@ -0,0 +1,11 @@ +{ + "name": "test-pipeline-2", + "allow_duplicate_names": true, + "libraries": [ + { + "file": { + "path": "/whatever.py" + } + } + ] +} diff --git a/acceptance/selftest/record_cloud/script b/acceptance/selftest/record_cloud/script index 8c6c5d9be7..a8c3c541f6 100644 --- a/acceptance/selftest/record_cloud/script +++ b/acceptance/selftest/record_cloud/script @@ -1 +1,25 @@ +# Proxy server successfully records a requests and returns a response. trace $CLI current-user me | jq .name.givenName + +# Proxy server should successfully return non 200 responses. +errcode trace $CLI jobs get 1234 + +# Verify that the entire crud lifecycle of a pipeline works. +title "Create a pipeline" +export pipeline_id=$($CLI pipelines create --json @pipeline1.json | jq -r .pipeline_id) +echo "" + +title "Get the pipeline" +trace $CLI pipelines get $pipeline_id | jq .name + +title "Update the pipeline" +trace $CLI pipelines update $pipeline_id --json @pipeline2.json + +title "Verify the update" +trace $CLI pipelines get $pipeline_id | jq .name + +title "Delete the pipeline" +trace $CLI pipelines delete $pipeline_id + +title "Verify the deletion" +errcode trace $CLI pipelines get $pipeline_id diff --git a/libs/testserver/proxy_server.go b/libs/testserver/proxy_server.go index c937040953..c78b0cf192 100644 --- a/libs/testserver/proxy_server.go +++ b/libs/testserver/proxy_server.go @@ -29,6 +29,17 @@ type ProxyServer struct { responseCallback func(request *Request, response *EncodedResponse) } +// This creates a reverse proxy server that sits in front of a real Databricks +// workspace. This is useful for recording API requests and responses in +// integration tests. +// +// Note: We cannot simply proxy the request from a localhost URL to a real +// workspace. This is because auth resolution in the Databricks SDK relies +// what the URL actually looks like to determine the auth method to use. +// For example, in OAuth flows, the SDK can make requests to different Microsoft +// OAuth endpoints based on the nature of the URL. +// For reference, see: +// https://github.com/databricks/databricks-sdk-go/blob/79e4b3a6e9b0b7dcb1af9ad4025deb447b01d933/common/environment/environments.go#L57 func NewProxyServer(t testutil.TestingT) *ProxyServer { router := mux.NewRouter() server := httptest.NewServer(router) @@ -48,11 +59,35 @@ func NewProxyServer(t testutil.TestingT) *ProxyServer { s.apiClient, err = client.New(&config.Config{}) require.NoError(t, err) + // Set up the proxy handler as the default handler for all requests. router.NotFoundHandler = http.HandlerFunc(s.proxyToCloud) return s } -// TODO: Iterate once on this function. +func (s *ProxyServer) reqBody(r Request) any { + // The SDK expects the query parameters to be specified in the "request body" + // argument for GET, DELETE, and HEAD requests in the .Do method. + if r.Method == "GET" || r.Method == "DELETE" || r.Method == "HEAD" { + queryParams := make(map[string]any) + for k, v := range r.URL.Query() { + queryParams[k] = v[0] + } + return queryParams + } + + // The SDK does not support directly passing a JSON serialized request + // body. So we convert it to a map if the body is a JSON object. + if json.Valid(r.Body) { + m := make(map[string]any) + err := json.Unmarshal(r.Body, &m) + assert.NoError(s.t, err) + return m + } + + // Otherwise, return the raw body. + return r.Body +} + func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { // Process requests sequentially. It's slower but is easier to reason about. // For example, requestCallback and responseCallback functions do not need @@ -67,10 +102,13 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { headers := make(map[string]string) for k, v := range r.Header { - // Authorization headers will be set by the SDK. No need to pass them along here. + // Authorization headers will be set by the SDK. Do not pass them along here. if k == "Authorization" { continue } + // The default HTTP client in Go sets the Accept-Encoding header to + // "gzip". Since it's meant for the server and will again be set by + // the SDK, do not pass it along here. if k == "Accept-Encoding" { continue } @@ -82,23 +120,52 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { queryParams[k] = v[0] } - // TODO: Since the response is always JSON, this should be specified in the header. + // Note: This only works for JSON responses. Eventually we should add support for other types + // of responses as and when needed. respB := map[string]any{} - err := s.apiClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, r.Body, &respB) - assert.NoError(s.t, err) + reqBody := s.reqBody(request) + err := s.apiClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, reqBody, &respB) + + // API errors from the SDK are expected to be of the type [apierr.APIError]. If we + // get an API error then parse the error and forward it in an appropriate format. + apiErr := &apierr.APIError{} + if errors.As(err, &apiErr) { + body := map[string]string{ + "error_code": apiErr.ErrorCode, + "message": apiErr.Message, + } + + b, err := json.Marshal(body) + assert.NoError(s.t, err) + + w.WriteHeader(apiErr.StatusCode) + _, err = w.Write(b) + assert.NoError(s.t, err) + + if s.responseCallback != nil { + s.responseCallback(&request, &EncodedResponse{ + StatusCode: apiErr.StatusCode, + Body: []byte(apiErr.Message), + }) + } + + return + } + + // Something else went wrong. if err != nil { - // API errors from the SDK are expected to be of the type apierr.APIError. - apiErr := &apierr.APIError{} - if errors.As(err, &apiErr) { - w.WriteHeader(apiErr.StatusCode) - _, err := w.Write(respB["message"].([]byte)) - assert.NoError(s.t, err) - } else { - // Something else went wrong. - w.WriteHeader(http.StatusInternalServerError) - _, err := w.Write([]byte(err.Error())) - assert.NoError(s.t, err) + w.WriteHeader(http.StatusInternalServerError) + _, err := w.Write([]byte(err.Error())) + assert.NoError(s.t, err) + + if s.responseCallback != nil { + s.responseCallback(&request, &EncodedResponse{ + StatusCode: 500, + Body: []byte(err.Error()), + }) } + + return } // Successful response @@ -106,15 +173,15 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { b, err := json.Marshal(respB) assert.NoError(s.t, err) + _, err = w.Write(b) + assert.NoError(s.t, err) + if s.responseCallback != nil { s.responseCallback(&request, &EncodedResponse{ StatusCode: 200, Body: b, }) } - - _, err = w.Write(b) - assert.NoError(s.t, err) } // Eventually we can implement this function to allow for per-test overrides From a2eb106738bcf7a63bc1ff605fde13e431d1c8d3 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 24 Apr 2025 19:33:14 +0200 Subject: [PATCH 16/57] split tests and add workspace-io test --- acceptance/acceptance_test.go | 2 +- .../record_cloud/basic/out.requests.txt | 4 ++ .../selftest/record_cloud/basic/output.txt | 3 + acceptance/selftest/record_cloud/basic/script | 2 + .../record_cloud/error/out.requests.txt | 4 ++ .../selftest/record_cloud/error/output.txt | 5 ++ acceptance/selftest/record_cloud/error/script | 2 + .../{ => pipeline-crud}/out.requests.txt | 8 --- .../{ => pipeline-crud}/output.txt | 8 --- .../{ => pipeline-crud}/pipeline1.json | 0 .../{ => pipeline-crud}/pipeline2.json | 0 .../record_cloud/{ => pipeline-crud}/script | 6 -- .../record_cloud/workspace-file-io/hello.txt | 1 + .../workspace-file-io/out.requests.txt | 58 +++++++++++++++++++ .../record_cloud/workspace-file-io/output.txt | 48 +++++++++++++++ .../record_cloud/workspace-file-io/script | 34 +++++++++++ 16 files changed, 162 insertions(+), 23 deletions(-) create mode 100644 acceptance/selftest/record_cloud/basic/out.requests.txt create mode 100644 acceptance/selftest/record_cloud/basic/output.txt create mode 100644 acceptance/selftest/record_cloud/basic/script create mode 100644 acceptance/selftest/record_cloud/error/out.requests.txt create mode 100644 acceptance/selftest/record_cloud/error/output.txt create mode 100644 acceptance/selftest/record_cloud/error/script rename acceptance/selftest/record_cloud/{ => pipeline-crud}/out.requests.txt (86%) rename acceptance/selftest/record_cloud/{ => pipeline-crud}/output.txt (79%) rename acceptance/selftest/record_cloud/{ => pipeline-crud}/pipeline1.json (100%) rename acceptance/selftest/record_cloud/{ => pipeline-crud}/pipeline2.json (100%) rename acceptance/selftest/record_cloud/{ => pipeline-crud}/script (72%) create mode 100644 acceptance/selftest/record_cloud/workspace-file-io/hello.txt create mode 100644 acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt create mode 100644 acceptance/selftest/record_cloud/workspace-file-io/output.txt create mode 100644 acceptance/selftest/record_cloud/workspace-file-io/script diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 04488f99af..494d3c8a7b 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -113,7 +113,7 @@ var Ignored = map[string]bool{ func TestAccept(t *testing.T) { // TODO: Clean this up. // loadDebugEnvIfRunFromIDE(t, "workspace") - testAccept(t, InprocessMode, "selftest/record_cloud") + testAccept(t, InprocessMode, "") } func TestInprocessMode(t *testing.T) { diff --git a/acceptance/selftest/record_cloud/basic/out.requests.txt b/acceptance/selftest/record_cloud/basic/out.requests.txt new file mode 100644 index 0000000000..ec41b6c5f6 --- /dev/null +++ b/acceptance/selftest/record_cloud/basic/out.requests.txt @@ -0,0 +1,4 @@ +{ + "method": "GET", + "path": "/api/2.0/preview/scim/v2/Me" +} diff --git a/acceptance/selftest/record_cloud/basic/output.txt b/acceptance/selftest/record_cloud/basic/output.txt new file mode 100644 index 0000000000..7f0b941262 --- /dev/null +++ b/acceptance/selftest/record_cloud/basic/output.txt @@ -0,0 +1,3 @@ + +>>> [CLI] current-user me +"[USERNAME]" diff --git a/acceptance/selftest/record_cloud/basic/script b/acceptance/selftest/record_cloud/basic/script new file mode 100644 index 0000000000..5d8c2eac19 --- /dev/null +++ b/acceptance/selftest/record_cloud/basic/script @@ -0,0 +1,2 @@ +# Proxy server successfully records a requests and returns a response. +trace $CLI current-user me | jq .name.givenName diff --git a/acceptance/selftest/record_cloud/error/out.requests.txt b/acceptance/selftest/record_cloud/error/out.requests.txt new file mode 100644 index 0000000000..30a1cb329c --- /dev/null +++ b/acceptance/selftest/record_cloud/error/out.requests.txt @@ -0,0 +1,4 @@ +{ + "method": "GET", + "path": "/api/2.2/jobs/get" +} diff --git a/acceptance/selftest/record_cloud/error/output.txt b/acceptance/selftest/record_cloud/error/output.txt new file mode 100644 index 0000000000..f911ca508f --- /dev/null +++ b/acceptance/selftest/record_cloud/error/output.txt @@ -0,0 +1,5 @@ + +>>> [CLI] jobs get 1234 +Error: Job 1234 does not exist. + +Exit code: 1 diff --git a/acceptance/selftest/record_cloud/error/script b/acceptance/selftest/record_cloud/error/script new file mode 100644 index 0000000000..662a0732cc --- /dev/null +++ b/acceptance/selftest/record_cloud/error/script @@ -0,0 +1,2 @@ +# Proxy server should successfully return non 200 responses. +errcode trace $CLI jobs get 1234 diff --git a/acceptance/selftest/record_cloud/out.requests.txt b/acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt similarity index 86% rename from acceptance/selftest/record_cloud/out.requests.txt rename to acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt index ac245e46ac..c5c203e309 100644 --- a/acceptance/selftest/record_cloud/out.requests.txt +++ b/acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt @@ -1,11 +1,3 @@ -{ - "method": "GET", - "path": "/api/2.0/preview/scim/v2/Me" -} -{ - "method": "GET", - "path": "/api/2.2/jobs/get" -} { "method": "POST", "path": "/api/2.0/pipelines", diff --git a/acceptance/selftest/record_cloud/output.txt b/acceptance/selftest/record_cloud/pipeline-crud/output.txt similarity index 79% rename from acceptance/selftest/record_cloud/output.txt rename to acceptance/selftest/record_cloud/pipeline-crud/output.txt index d80033ce81..cceca41f90 100644 --- a/acceptance/selftest/record_cloud/output.txt +++ b/acceptance/selftest/record_cloud/pipeline-crud/output.txt @@ -1,12 +1,4 @@ ->>> [CLI] current-user me -"[USERNAME]" - ->>> [CLI] jobs get 1234 -Error: Job 1234 does not exist. - -Exit code: 1 - === Create a pipeline === Get the pipeline diff --git a/acceptance/selftest/record_cloud/pipeline1.json b/acceptance/selftest/record_cloud/pipeline-crud/pipeline1.json similarity index 100% rename from acceptance/selftest/record_cloud/pipeline1.json rename to acceptance/selftest/record_cloud/pipeline-crud/pipeline1.json diff --git a/acceptance/selftest/record_cloud/pipeline2.json b/acceptance/selftest/record_cloud/pipeline-crud/pipeline2.json similarity index 100% rename from acceptance/selftest/record_cloud/pipeline2.json rename to acceptance/selftest/record_cloud/pipeline-crud/pipeline2.json diff --git a/acceptance/selftest/record_cloud/script b/acceptance/selftest/record_cloud/pipeline-crud/script similarity index 72% rename from acceptance/selftest/record_cloud/script rename to acceptance/selftest/record_cloud/pipeline-crud/script index a8c3c541f6..965ffe0e4d 100644 --- a/acceptance/selftest/record_cloud/script +++ b/acceptance/selftest/record_cloud/pipeline-crud/script @@ -1,9 +1,3 @@ -# Proxy server successfully records a requests and returns a response. -trace $CLI current-user me | jq .name.givenName - -# Proxy server should successfully return non 200 responses. -errcode trace $CLI jobs get 1234 - # Verify that the entire crud lifecycle of a pipeline works. title "Create a pipeline" export pipeline_id=$($CLI pipelines create --json @pipeline1.json | jq -r .pipeline_id) diff --git a/acceptance/selftest/record_cloud/workspace-file-io/hello.txt b/acceptance/selftest/record_cloud/workspace-file-io/hello.txt new file mode 100644 index 0000000000..4b5fa63702 --- /dev/null +++ b/acceptance/selftest/record_cloud/workspace-file-io/hello.txt @@ -0,0 +1 @@ +hello, world diff --git a/acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt b/acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt new file mode 100644 index 0000000000..1612113586 --- /dev/null +++ b/acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt @@ -0,0 +1,58 @@ +{ + "method": "GET", + "path": "/api/2.0/preview/scim/v2/Me" +} +{ + "method": "POST", + "path": "/api/2.0/workspace/mkdirs", + "body": { + "path": "/Users/[USERNAME]/[UNIQUE_NAME]" + } +} +{ + "method": "POST", + "path": "/api/2.0/workspace/import", + "body": { + "content": "aGVsbG8sIHdvcmxkCg==", + "format": "AUTO", + "path": "/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt" + } +} +{ + "method": "GET", + "path": "/api/2.0/workspace/list" +} +{ + "method": "GET", + "path": "/api/2.0/workspace/get-status" +} +{ + "method": "GET", + "path": "/api/2.0/workspace/export" +} +{ + "method": "POST", + "path": "/api/2.0/workspace/delete", + "body": { + "path": "/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt" + } +} +{ + "method": "GET", + "path": "/api/2.0/workspace/get-status" +} +{ + "method": "GET", + "path": "/api/2.0/workspace/list" +} +{ + "method": "POST", + "path": "/api/2.0/workspace/delete", + "body": { + "path": "/Users/[USERNAME]/[UNIQUE_NAME]" + } +} +{ + "method": "GET", + "path": "/api/2.0/workspace/list" +} diff --git a/acceptance/selftest/record_cloud/workspace-file-io/output.txt b/acceptance/selftest/record_cloud/workspace-file-io/output.txt new file mode 100644 index 0000000000..1fb61c716b --- /dev/null +++ b/acceptance/selftest/record_cloud/workspace-file-io/output.txt @@ -0,0 +1,48 @@ + +=== create a folder +>>> [CLI] workspace mkdirs /Users/[USERNAME]/[UNIQUE_NAME] + +=== upload a file +>>> [CLI] workspace import /Users/[USERNAME]/[UNIQUE_NAME]/hello.txt --format AUTO --file ./hello.txt + +=== list the folder +>>> [CLI] workspace list /Users/[USERNAME]/[UNIQUE_NAME] --output json +{ + "path": "/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt", + "object_type": "FILE" +} + +=== stat the file +>>> [CLI] workspace get-status /Users/[USERNAME]/[UNIQUE_NAME]/hello.txt +{ + "path": "/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt", + "object_type": "FILE" +} + +=== download the file +>>> [CLI] workspace export /Users/[USERNAME]/[UNIQUE_NAME]/hello.txt --format AUTO --file ./hello2.txt + +>>> cat hello2.txt +hello, world + +=== delete the file +>>> [CLI] workspace delete /Users/[USERNAME]/[UNIQUE_NAME]/hello.txt + +=== stat the file again +>>> [CLI] workspace get-status /Users/[USERNAME]/[UNIQUE_NAME]/hello.txt +Error: Path (/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt) doesn't exist. + +Exit code: 1 + +=== list the folder +>>> [CLI] workspace list /Users/[USERNAME]/[UNIQUE_NAME] +ID Type Language Path + +=== delete the folder +>>> [CLI] workspace delete /Users/[USERNAME]/[UNIQUE_NAME] + +=== list the folder again +>>> [CLI] workspace list /Users/[USERNAME]/[UNIQUE_NAME] +Error: Path (/Users/[USERNAME]/[UNIQUE_NAME]) doesn't exist. + +Exit code: 1 diff --git a/acceptance/selftest/record_cloud/workspace-file-io/script b/acceptance/selftest/record_cloud/workspace-file-io/script new file mode 100644 index 0000000000..5924418d9a --- /dev/null +++ b/acceptance/selftest/record_cloud/workspace-file-io/script @@ -0,0 +1,34 @@ +export username=$($CLI current-user me | jq -r .userName) + +title "create a folder" +trace $CLI workspace mkdirs /Users/$username/$UNIQUE_NAME + +title "upload a file" +trace $CLI workspace import /Users/$username/$UNIQUE_NAME/hello.txt --format AUTO --file ./hello.txt + +title "list the folder" +trace $CLI workspace list /Users/$username/$UNIQUE_NAME --output json | jq '.[] | {path, object_type}' + +title "stat the file" +trace $CLI workspace get-status /Users/$username/$UNIQUE_NAME/hello.txt | jq '{path, object_type}' + +title "download the file" +trace $CLI workspace export /Users/$username/$UNIQUE_NAME/hello.txt --format AUTO --file ./hello2.txt + +trace cat hello2.txt +rm hello2.txt + +title "delete the file" +trace $CLI workspace delete /Users/$username/$UNIQUE_NAME/hello.txt + +title "stat the file again" +errcode trace $CLI workspace get-status /Users/$username/$UNIQUE_NAME/hello.txt + +title "list the folder" +trace $CLI workspace list /Users/$username/$UNIQUE_NAME + +title "delete the folder" +trace $CLI workspace delete /Users/$username/$UNIQUE_NAME + +title "list the folder again" +errcode trace $CLI workspace list /Users/$username/$UNIQUE_NAME From b32da964f96192ea2fd4d6350d33a611ea3ff097 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 24 Apr 2025 19:56:54 +0200 Subject: [PATCH 17/57] add volume io testcase --- .../selftest/record_cloud/volume-io/hello.txt | 1 + .../record_cloud/volume-io/out.requests.txt | 79 +++++++++++++++++++ .../record_cloud/volume-io/output.txt | 51 ++++++++++++ .../selftest/record_cloud/volume-io/script | 25 ++++++ .../selftest/record_cloud/volume-io/test.toml | 1 + libs/testserver/proxy_server.go | 10 +-- 6 files changed, 161 insertions(+), 6 deletions(-) create mode 100644 acceptance/selftest/record_cloud/volume-io/hello.txt create mode 100644 acceptance/selftest/record_cloud/volume-io/out.requests.txt create mode 100644 acceptance/selftest/record_cloud/volume-io/output.txt create mode 100644 acceptance/selftest/record_cloud/volume-io/script create mode 100644 acceptance/selftest/record_cloud/volume-io/test.toml diff --git a/acceptance/selftest/record_cloud/volume-io/hello.txt b/acceptance/selftest/record_cloud/volume-io/hello.txt new file mode 100644 index 0000000000..4b5fa63702 --- /dev/null +++ b/acceptance/selftest/record_cloud/volume-io/hello.txt @@ -0,0 +1 @@ +hello, world diff --git a/acceptance/selftest/record_cloud/volume-io/out.requests.txt b/acceptance/selftest/record_cloud/volume-io/out.requests.txt new file mode 100644 index 0000000000..6b72b16797 --- /dev/null +++ b/acceptance/selftest/record_cloud/volume-io/out.requests.txt @@ -0,0 +1,79 @@ +{ + "method": "POST", + "path": "/api/2.1/unity-catalog/schemas", + "body": { + "catalog_name": "main", + "name": "schema-[UNIQUE_NAME]" + } +} +{ + "method": "GET", + "path": "/api/2.1/unity-catalog/schemas/main.schema-[UNIQUE_NAME]" +} +{ + "method": "POST", + "path": "/api/2.1/unity-catalog/volumes", + "body": { + "catalog_name": "main", + "name": "volume-[UNIQUE_NAME]", + "schema_name": "schema-[UNIQUE_NAME]", + "volume_type": "MANAGED" + } +} +{ + "method": "GET", + "path": "/api/2.1/unity-catalog/volumes/main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]" +} +{ + "method": "HEAD", + "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" +} +{ + "method": "HEAD", + "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" +} +{ + "method": "PUT", + "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt", + "raw_body": "hello, world\n" +} +{ + "method": "GET", + "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" +} +{ + "method": "GET", + "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" +} +{ + "method": "HEAD", + "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" +} +{ + "method": "HEAD", + "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" +} +{ + "method": "DELETE", + "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" +} +{ + "method": "GET", + "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" +} +{ + "method": "DELETE", + "path": "/api/2.1/unity-catalog/volumes/main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]" +} +{ + "method": "GET", + "path": "/api/2.1/unity-catalog/volumes/main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]" +} +{ + "method": "DELETE", + "path": "/api/2.1/unity-catalog/schemas/main.schema-[UNIQUE_NAME]" +} +{ + "method": "GET", + "path": "/api/2.1/unity-catalog/schemas/main.schema-[UNIQUE_NAME]" +} diff --git a/acceptance/selftest/record_cloud/volume-io/output.txt b/acceptance/selftest/record_cloud/volume-io/output.txt new file mode 100644 index 0000000000..a6e2673ff4 --- /dev/null +++ b/acceptance/selftest/record_cloud/volume-io/output.txt @@ -0,0 +1,51 @@ + +>>> [CLI] schemas create schema-[UNIQUE_NAME] main +{ + "full_name": "main.schema-[UNIQUE_NAME]", + "owner": "[USERNAME]" +} + +>>> [CLI] schemas get main.schema-[UNIQUE_NAME] +{ + "full_name": "main.schema-[UNIQUE_NAME]", + "owner": "[USERNAME]" +} + +>>> [CLI] volumes create main schema-[UNIQUE_NAME] volume-[UNIQUE_NAME] MANAGED +{ + "full_name": "main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]", + "owner": "[USERNAME]" +} + +>>> [CLI] volumes read main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME] +{ + "full_name": "main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]", + "owner": "[USERNAME]" +} + +>>> [CLI] fs cp ./hello.txt dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME] +./hello.txt -> dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt + +>>> [CLI] fs ls dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME] +hello.txt + +>>> [CLI] fs cat dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt +hello, world + +>>> [CLI] fs rm dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt + +>>> [CLI] fs ls dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME] + +>>> [CLI] volumes delete main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME] + +>>> [CLI] volumes read main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME] +Error: Volume 'main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]' does not exist. + +Exit code: 1 + +>>> [CLI] schemas delete main.schema-[UNIQUE_NAME] + +>>> [CLI] schemas get main.schema-[UNIQUE_NAME] +Error: Schema 'main.schema-[UNIQUE_NAME]' does not exist. + +Exit code: 1 diff --git a/acceptance/selftest/record_cloud/volume-io/script b/acceptance/selftest/record_cloud/volume-io/script new file mode 100644 index 0000000000..04e297d6b9 --- /dev/null +++ b/acceptance/selftest/record_cloud/volume-io/script @@ -0,0 +1,25 @@ +trace $CLI schemas create schema-$UNIQUE_NAME main | jq '{full_name, owner}' + +trace $CLI schemas get main.schema-$UNIQUE_NAME | jq '{full_name, owner}' + +trace $CLI volumes create main schema-$UNIQUE_NAME volume-$UNIQUE_NAME MANAGED | jq '{full_name, owner}' + +trace $CLI volumes read main.schema-$UNIQUE_NAME.volume-$UNIQUE_NAME | jq '{full_name, owner}' + +trace $CLI fs cp ./hello.txt dbfs:/Volumes/main/schema-$UNIQUE_NAME/volume-$UNIQUE_NAME + +trace $CLI fs ls dbfs:/Volumes/main/schema-$UNIQUE_NAME/volume-$UNIQUE_NAME + +trace $CLI fs cat dbfs:/Volumes/main/schema-$UNIQUE_NAME/volume-$UNIQUE_NAME/hello.txt + +trace $CLI fs rm dbfs:/Volumes/main/schema-$UNIQUE_NAME/volume-$UNIQUE_NAME/hello.txt + +trace $CLI fs ls dbfs:/Volumes/main/schema-$UNIQUE_NAME/volume-$UNIQUE_NAME + +trace $CLI volumes delete main.schema-$UNIQUE_NAME.volume-$UNIQUE_NAME + +errcode trace $CLI volumes read main.schema-$UNIQUE_NAME.volume-$UNIQUE_NAME + +trace $CLI schemas delete main.schema-$UNIQUE_NAME + +errcode trace $CLI schemas get main.schema-$UNIQUE_NAME diff --git a/acceptance/selftest/record_cloud/volume-io/test.toml b/acceptance/selftest/record_cloud/volume-io/test.toml new file mode 100644 index 0000000000..9f65b59bbd --- /dev/null +++ b/acceptance/selftest/record_cloud/volume-io/test.toml @@ -0,0 +1 @@ +RequiresUnityCatalog = true diff --git a/libs/testserver/proxy_server.go b/libs/testserver/proxy_server.go index c78b0cf192..451f5f6ccc 100644 --- a/libs/testserver/proxy_server.go +++ b/libs/testserver/proxy_server.go @@ -1,6 +1,7 @@ package testserver import ( + "bytes" "context" "encoding/json" "errors" @@ -120,11 +121,9 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { queryParams[k] = v[0] } - // Note: This only works for JSON responses. Eventually we should add support for other types - // of responses as and when needed. - respB := map[string]any{} reqBody := s.reqBody(request) - err := s.apiClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, reqBody, &respB) + respBody := bytes.Buffer{} + err := s.apiClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, reqBody, &respBody) // API errors from the SDK are expected to be of the type [apierr.APIError]. If we // get an API error then parse the error and forward it in an appropriate format. @@ -170,8 +169,7 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { // Successful response w.WriteHeader(200) - b, err := json.Marshal(respB) - assert.NoError(s.t, err) + b := respBody.Bytes() _, err = w.Write(b) assert.NoError(s.t, err) From 9b36628e59ffe61c9eba9c2da424cdce261dbfe1 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 24 Apr 2025 20:04:12 +0200 Subject: [PATCH 18/57] - --- acceptance/selftest/record_cloud/workspace-file-io/script | 3 +++ 1 file changed, 3 insertions(+) diff --git a/acceptance/selftest/record_cloud/workspace-file-io/script b/acceptance/selftest/record_cloud/workspace-file-io/script index 5924418d9a..ddc038f3a4 100644 --- a/acceptance/selftest/record_cloud/workspace-file-io/script +++ b/acceptance/selftest/record_cloud/workspace-file-io/script @@ -1,5 +1,8 @@ export username=$($CLI current-user me | jq -r .userName) +# TODO CONTINUE(v5): Figure out why this is failing on windows, and it automatically +# being recorded as as windows file without quotes. Why the difference? +# test run: https://github.com/databricks-eng/eng-dev-ecosystem/actions/runs/14647936035/job/41106657293 title "create a folder" trace $CLI workspace mkdirs /Users/$username/$UNIQUE_NAME From 18f08e82d18d40309a53f0e225f4c2c121e059d9 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 00:12:52 +0200 Subject: [PATCH 19/57] fix test on windows --- acceptance/selftest/record_cloud/workspace-file-io/script | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/acceptance/selftest/record_cloud/workspace-file-io/script b/acceptance/selftest/record_cloud/workspace-file-io/script index ddc038f3a4..db9cc81cd1 100644 --- a/acceptance/selftest/record_cloud/workspace-file-io/script +++ b/acceptance/selftest/record_cloud/workspace-file-io/script @@ -1,5 +1,10 @@ export username=$($CLI current-user me | jq -r .userName) +# MSYS automatically converts absolute paths like /Users/$username/$UNIQUE_NAME to +# C:/Program Files/Git/Users/$username/UNIQUE_NAME before passing it to the CLI +# Setting this environment variable prevents that conversion on windows. +export MSYS_NO_PATHCONV=1 + # TODO CONTINUE(v5): Figure out why this is failing on windows, and it automatically # being recorded as as windows file without quotes. Why the difference? # test run: https://github.com/databricks-eng/eng-dev-ecosystem/actions/runs/14647936035/job/41106657293 From c5deda84d8ae3d2bb2728c63a1ede5a39d4006b8 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 00:15:28 +0200 Subject: [PATCH 20/57] - --- acceptance/selftest/record_cloud/workspace-file-io/script | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/acceptance/selftest/record_cloud/workspace-file-io/script b/acceptance/selftest/record_cloud/workspace-file-io/script index db9cc81cd1..0c7eeb3378 100644 --- a/acceptance/selftest/record_cloud/workspace-file-io/script +++ b/acceptance/selftest/record_cloud/workspace-file-io/script @@ -5,9 +5,8 @@ export username=$($CLI current-user me | jq -r .userName) # Setting this environment variable prevents that conversion on windows. export MSYS_NO_PATHCONV=1 -# TODO CONTINUE(v5): Figure out why this is failing on windows, and it automatically -# being recorded as as windows file without quotes. Why the difference? -# test run: https://github.com/databricks-eng/eng-dev-ecosystem/actions/runs/14647936035/job/41106657293 +# TODO CONTINUE(v6): Is fixed on windows? If so clean up the env story. Mostly cleanup left with this PR. +# TODO: Try recording requests in a bundle deploy acceptance test? title "create a folder" trace $CLI workspace mkdirs /Users/$username/$UNIQUE_NAME From cd8ebc6a7559701247558a688816c39a6ed6d8d0 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 10:00:21 +0200 Subject: [PATCH 21/57] fix windows base64 --- .../record_cloud/workspace-file-io/out.requests.txt | 2 +- .../selftest/record_cloud/workspace-file-io/test.toml | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) create mode 100644 acceptance/selftest/record_cloud/workspace-file-io/test.toml diff --git a/acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt b/acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt index 1612113586..2679e24262 100644 --- a/acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt +++ b/acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt @@ -13,7 +13,7 @@ "method": "POST", "path": "/api/2.0/workspace/import", "body": { - "content": "aGVsbG8sIHdvcmxkCg==", + "content": "[HELLO-WORLD]", "format": "AUTO", "path": "/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt" } diff --git a/acceptance/selftest/record_cloud/workspace-file-io/test.toml b/acceptance/selftest/record_cloud/workspace-file-io/test.toml new file mode 100644 index 0000000000..fa35335247 --- /dev/null +++ b/acceptance/selftest/record_cloud/workspace-file-io/test.toml @@ -0,0 +1,9 @@ +# hello, world\n base64 encoded in unix +[[Repls]] +Old = "aGVsbG8sIHdvcmxkCg==" +New = "[HELLO-WORLD]" + +# hello, world\n base64 encoded in windows +[[Repls]] +Old = "aGVsbG8sIHdvcmxkDQo=" +New = "[HELLO-WORLD]" From ef2ec73ee8230b3ccc452ba19be4bc0c4c7a7ddc Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 11:11:43 +0200 Subject: [PATCH 22/57] fix windows newline in raw body --- acceptance/selftest/record_cloud/volume-io/test.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/acceptance/selftest/record_cloud/volume-io/test.toml b/acceptance/selftest/record_cloud/volume-io/test.toml index 9f65b59bbd..45c81a9210 100644 --- a/acceptance/selftest/record_cloud/volume-io/test.toml +++ b/acceptance/selftest/record_cloud/volume-io/test.toml @@ -1 +1,5 @@ RequiresUnityCatalog = true + +[[Repls]] +Old = '\r\n' +New = '\n' From 6e5fa126a6a6a00414ad40d6920cb0c08c1aceba Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 12:07:33 +0200 Subject: [PATCH 23/57] lint --- acceptance/selftest/record_cloud/volume-io/test.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/acceptance/selftest/record_cloud/volume-io/test.toml b/acceptance/selftest/record_cloud/volume-io/test.toml index 45c81a9210..afc735e548 100644 --- a/acceptance/selftest/record_cloud/volume-io/test.toml +++ b/acceptance/selftest/record_cloud/volume-io/test.toml @@ -2,4 +2,4 @@ RequiresUnityCatalog = true [[Repls]] Old = '\r\n' -New = '\n' +New = '\n' From ee83d69fcf442ca5daa4adcb4a99ef0206c447fd Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 14:34:51 +0200 Subject: [PATCH 24/57] fix windows newline in raw body again --- .../selftest/record_cloud/volume-io/out.requests.txt | 2 +- acceptance/selftest/record_cloud/volume-io/test.toml | 10 ++++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/acceptance/selftest/record_cloud/volume-io/out.requests.txt b/acceptance/selftest/record_cloud/volume-io/out.requests.txt index 6b72b16797..06834706d7 100644 --- a/acceptance/selftest/record_cloud/volume-io/out.requests.txt +++ b/acceptance/selftest/record_cloud/volume-io/out.requests.txt @@ -35,7 +35,7 @@ { "method": "PUT", "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt", - "raw_body": "hello, world\n" + "raw_body": "hello, world" } { "method": "GET", diff --git a/acceptance/selftest/record_cloud/volume-io/test.toml b/acceptance/selftest/record_cloud/volume-io/test.toml index afc735e548..95af68200f 100644 --- a/acceptance/selftest/record_cloud/volume-io/test.toml +++ b/acceptance/selftest/record_cloud/volume-io/test.toml @@ -1,5 +1,11 @@ RequiresUnityCatalog = true +# Windows and unix send different EOL characters in the upload request. +# Normalize both to a single pattern for assertions. [[Repls]] -Old = '\r\n' -New = '\n' +Old = '\\n' +New = '' + +[[Repls]] +Old = '\\r\\n' +New = '' From c95c2b8e35a7fa05316ad4d09349f052a8b80276 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 14:52:31 +0200 Subject: [PATCH 25/57] pass process env to child process --- acceptance/acceptance_test.go | 29 ++++++++++++++------------- acceptance/internal/prepare_server.go | 19 +++++++----------- 2 files changed, 22 insertions(+), 26 deletions(-) diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 494d3c8a7b..7982d11f31 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -25,6 +25,7 @@ import ( "github.com/databricks/cli/acceptance/internal" "github.com/databricks/cli/internal/testutil" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/testdiff" "github.com/stretchr/testify/require" ) @@ -404,13 +405,23 @@ func runTest(t *testing.T, err = CopyDir(dir, tmpDir, inputs, outputs) require.NoError(t, err) - cfg, user, env := internal.PrepareServerAndClient(t, config, LogRequests, tmpDir) + cfg, user := internal.PrepareServerAndClient(t, config, LogRequests, tmpDir) testdiff.PrepareReplacementsUser(t, &repls, user) testdiff.PrepareReplacementsWorkspaceConfig(t, &repls, cfg) - if env != nil && inprocessMode { + // In inprocess mode, the "script" is executed in the same process as the test runner. + // Thus we need to modify the environment of the test runner, so that the script sees the + // appropriate environment variables. + // + // This is important for when a reverse proxy sits between the script and a real databricks workspace. + // In that case we need to modify the current process environment so that the script communicates with + // the reverse proxy. + + // TODO: Add comment here. Remove old comment + processEnv := auth.ProcessEnv(cfg) + if inprocessMode { testutil.NullEnvironment(t) - for _, kv := range env { + for _, kv := range processEnv { parts := strings.SplitN(kv, "=", 2) require.Len(t, parts, 2) t.Setenv(parts[0], parts[1]) @@ -419,19 +430,9 @@ func runTest(t *testing.T, args := []string{"bash", "-euo", "pipefail", EntryPointScript} cmd := exec.Command(args[0], args[1:]...) - cmd.Env = os.Environ() - if env != nil { - cmd.Env = env - } + cmd.Env = processEnv cmd.Env = append(cmd.Env, "UNIQUE_NAME="+uniqueName) cmd.Env = append(cmd.Env, "TEST_TMP_DIR="+tmpDir) - - // Configure resolved credentials in the environment. - cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+cfg.Host) - if cfg.Token != "" { - cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+cfg.Token) - } - // Must be added PrepareReplacementsUser, otherwise conflicts with [USERNAME] testdiff.PrepareReplacementsUUID(t, &repls) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 4b7cf1e7ed..43375e960c 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -13,7 +13,6 @@ import ( "time" "unicode/utf8" - "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" @@ -24,7 +23,6 @@ import ( "github.com/stretchr/testify/require" ) -// TODO: Make CLI auth for with this. func StartDefaultServer(t *testing.T) { s := testserver.NewLocalServer(t) addDefaultHandlers(s) @@ -40,8 +38,9 @@ func isTruePtr(value *bool) bool { return value != nil && *value } -// TODO: Clean up the env story here. -func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User, []string) { +// This function return a configuration that contains the auth credentials that should be used +// to run the test. +func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User) { cloudEnv := os.Getenv("CLOUD_ENV") recordRequests := isTruePtr(config.RecordRequests) @@ -58,11 +57,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o // Start a proxy server that sits in front of of a real Databricks workspace. cfg := startProxyServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) - - // TODO: I only need to set these environment variables in hte inprocess mode. - // Support both child process and inprocess mode. - env := auth.ProcessEnv(cfg) - return cfg, *user, env + return cfg, *user } // If we are running on a cloud environment, use the host configured in the @@ -76,7 +71,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o user, err := w.CurrentUser.Me(context.Background()) require.NoError(t, err, "Failed to get current user") - return w.Config, *user, nil + return w.Config, *user } // If we are not recording requests, and no custom server server stubs are configured, @@ -92,14 +87,14 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o Token: token, } - return cfg, TestUser, nil + return cfg, TestUser } cfg := startLocalServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) // For the purposes of replacements, use testUser for local runs. // Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement: - return cfg, TestUser, nil + return cfg, TestUser } func recordRequestsCallback(t *testing.T, includeHeaders []string, outputDir string) func(request *testserver.Request) { From 94aaa4a2872bae73a2ebeed01d246295c6a36a74 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 15:22:30 +0200 Subject: [PATCH 26/57] - --- acceptance/selftest/record_cloud/volume-io/hello.txt | 2 +- .../selftest/record_cloud/volume-io/out.requests.txt | 2 +- acceptance/selftest/record_cloud/volume-io/output.txt | 1 - acceptance/selftest/record_cloud/volume-io/test.toml | 10 ---------- 4 files changed, 2 insertions(+), 13 deletions(-) diff --git a/acceptance/selftest/record_cloud/volume-io/hello.txt b/acceptance/selftest/record_cloud/volume-io/hello.txt index 4b5fa63702..8c01d89ae0 100644 --- a/acceptance/selftest/record_cloud/volume-io/hello.txt +++ b/acceptance/selftest/record_cloud/volume-io/hello.txt @@ -1 +1 @@ -hello, world +hello, world \ No newline at end of file diff --git a/acceptance/selftest/record_cloud/volume-io/out.requests.txt b/acceptance/selftest/record_cloud/volume-io/out.requests.txt index 06834706d7..36425422d5 100644 --- a/acceptance/selftest/record_cloud/volume-io/out.requests.txt +++ b/acceptance/selftest/record_cloud/volume-io/out.requests.txt @@ -35,7 +35,7 @@ { "method": "PUT", "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt", - "raw_body": "hello, world" + "raw_body": "hello, world" } { "method": "GET", diff --git a/acceptance/selftest/record_cloud/volume-io/output.txt b/acceptance/selftest/record_cloud/volume-io/output.txt index a6e2673ff4..3ab411a8a4 100644 --- a/acceptance/selftest/record_cloud/volume-io/output.txt +++ b/acceptance/selftest/record_cloud/volume-io/output.txt @@ -31,7 +31,6 @@ hello.txt >>> [CLI] fs cat dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt hello, world - >>> [CLI] fs rm dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt >>> [CLI] fs ls dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME] diff --git a/acceptance/selftest/record_cloud/volume-io/test.toml b/acceptance/selftest/record_cloud/volume-io/test.toml index 95af68200f..9f65b59bbd 100644 --- a/acceptance/selftest/record_cloud/volume-io/test.toml +++ b/acceptance/selftest/record_cloud/volume-io/test.toml @@ -1,11 +1 @@ RequiresUnityCatalog = true - -# Windows and unix send different EOL characters in the upload request. -# Normalize both to a single pattern for assertions. -[[Repls]] -Old = '\\n' -New = '' - -[[Repls]] -Old = '\\r\\n' -New = '' From bdddde64081efe38b1338f334f34388b6fe22cdb Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 15:25:56 +0200 Subject: [PATCH 27/57] - --- .wsignore | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.wsignore b/.wsignore index 827ebf4332..d4831cd94e 100644 --- a/.wsignore +++ b/.wsignore @@ -11,6 +11,12 @@ experimental/python/docs/images/databricks-logo.svg **/*.zip **/*.whl +# new lines are recorded differently on windows and unix. +# In unix: "raw_body": "hello, world\n" +# In windows: "raw_body": "hello, world\r\n" +# In order to prevent that difference, hello.txt does not have a trailing newline. +acceptance/selftest/record_cloud/volume-io/hello.txt + # "bundle run" has trailing whitespace: acceptance/bundle/integration_whl/*/output.txt From 609fb7e8c87bf5d5e474f9a98d30cab6e2d89ec3 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 16:10:23 +0200 Subject: [PATCH 28/57] cleanup --- acceptance/acceptance_test.go | 47 ++----------------- .../record_cloud/workspace-file-io/script | 2 - libs/testdiff/replacement.go | 1 - 3 files changed, 3 insertions(+), 47 deletions(-) diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index c697431362..caae4a39d7 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -79,42 +79,7 @@ var Ignored = map[string]bool{ ReplsFile: true, } -// // Detects if test is run from "debug test" feature in VS Code. -// func IsInDebug() bool { -// ex, _ := os.Executable() -// return strings.HasPrefix(path.Base(ex), "__debug_bin") -// } - -// // Loads debug environment from ~/.databricks/debug-env.json. -// func loadDebugEnvIfRunFromIDE(t testutil.TestingT, key string) { -// if !IsInDebug() { -// return -// } -// home, err := os.UserHomeDir() -// if err != nil { -// t.Fatalf("cannot find user home: %s", err) -// } -// raw, err := os.ReadFile(filepath.Join(home, ".databricks/debug-env.json")) -// if err != nil { -// t.Fatalf("cannot load ~/.databricks/debug-env.json: %s", err) -// } -// var conf map[string]map[string]string -// err = json.Unmarshal(raw, &conf) -// if err != nil { -// t.Fatalf("cannot parse ~/.databricks/debug-env.json: %s", err) -// } -// vars, ok := conf[key] -// if !ok { -// t.Fatalf("~/.databricks/debug-env.json#%s not configured", key) -// } -// for k, v := range vars { -// os.Setenv(k, v) -// } -// } - func TestAccept(t *testing.T) { - // TODO: Clean this up. - // loadDebugEnvIfRunFromIDE(t, "workspace") testAccept(t, InprocessMode, "") } @@ -413,15 +378,9 @@ func runTest(t *testing.T, testdiff.PrepareReplacementsUser(t, &repls, user) testdiff.PrepareReplacementsWorkspaceConfig(t, &repls, cfg) - // In inprocess mode, the "script" is executed in the same process as the test runner. - // Thus we need to modify the environment of the test runner, so that the script sees the - // appropriate environment variables. - // - // This is important for when a reverse proxy sits between the script and a real databricks workspace. - // In that case we need to modify the current process environment so that the script communicates with - // the reverse proxy. - - // TODO: Add comment here. Remove old comment + // In inprocess mode, we need to modify the environment of the test runner, so that the script sees the + // appropriate environment variables. This is necessary because any $CLI invocations in the "script" + // will be executed in the same process as the test runner. processEnv := auth.ProcessEnv(cfg) if inprocessMode { testutil.NullEnvironment(t) diff --git a/acceptance/selftest/record_cloud/workspace-file-io/script b/acceptance/selftest/record_cloud/workspace-file-io/script index 0c7eeb3378..c44e67e7de 100644 --- a/acceptance/selftest/record_cloud/workspace-file-io/script +++ b/acceptance/selftest/record_cloud/workspace-file-io/script @@ -5,8 +5,6 @@ export username=$($CLI current-user me | jq -r .userName) # Setting this environment variable prevents that conversion on windows. export MSYS_NO_PATHCONV=1 -# TODO CONTINUE(v6): Is fixed on windows? If so clean up the env story. Mostly cleanup left with this PR. -# TODO: Try recording requests in a bundle deploy acceptance test? title "create a folder" trace $CLI workspace mkdirs /Users/$username/$UNIQUE_NAME diff --git a/libs/testdiff/replacement.go b/libs/testdiff/replacement.go index b8add9b9a8..51680afba9 100644 --- a/libs/testdiff/replacement.go +++ b/libs/testdiff/replacement.go @@ -137,7 +137,6 @@ func (r *ReplacementsContext) SetPathWithParents(old, new string) { r.SetPath(filepath.Dir(filepath.Dir(old)), new+"_GPARENT") } -// TODO (followup): remove this function. func PrepareReplacementsWorkspaceClient(t testutil.TestingT, r *ReplacementsContext, w *databricks.WorkspaceClient) { t.Helper() PrepareReplacementsWorkspaceConfig(t, r, w.Config) From b9f52cb52b9e77068d00bad4cc202254e963b1fa Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 16:13:07 +0200 Subject: [PATCH 29/57] - --- acceptance/acceptance_test.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index caae4a39d7..7d0d38a465 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -91,7 +91,7 @@ func TestInprocessMode(t *testing.T) { require.Equal(t, 1, testAccept(t, true, "selftest/server")) } -func testAccept(t *testing.T, InprocessMode bool, singleTest string) int { +func testAccept(t *testing.T, inprocessMode bool, singleTest string) int { repls := testdiff.ReplacementsContext{} cwd, err := os.Getwd() require.NoError(t, err) @@ -121,7 +121,7 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int { execPath := "" - if InprocessMode { + if inprocessMode { cmdServer := internal.StartCmdServer(t) t.Setenv("CMD_SERVER_URL", cmdServer.URL) execPath = filepath.Join(cwd, "bin", "callserver.py") @@ -210,7 +210,7 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int { t.Skip(skipReason) } - if !InprocessMode { + if !inprocessMode { t.Parallel() } @@ -225,15 +225,15 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int { if len(expanded) == 1 { // env vars aren't part of the test case name, so log them for debugging t.Logf("Running test with env %v", expanded[0]) - runTest(t, dir, coverDir, repls.Clone(), config, configPath, expanded[0], InprocessMode) + runTest(t, dir, coverDir, repls.Clone(), config, configPath, expanded[0], inprocessMode) } else { for _, envset := range expanded { envname := strings.Join(envset, "/") t.Run(envname, func(t *testing.T) { - if !InprocessMode { + if !inprocessMode { t.Parallel() } - runTest(t, dir, coverDir, repls.Clone(), config, configPath, envset, InprocessMode) + runTest(t, dir, coverDir, repls.Clone(), config, configPath, envset, inprocessMode) }) } } From f2faf044144c0f7466f2cc6432541b1f51a3db61 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 16:25:40 +0200 Subject: [PATCH 30/57] cleanup --- acceptance/internal/prepare_server.go | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 43375e960c..7b82b9b81d 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -38,14 +38,14 @@ func isTruePtr(value *bool) bool { return value != nil && *value } -// This function return a configuration that contains the auth credentials that should be used -// to run the test. +// This function returns the workspace auth configuration that should be used to authenticate +// the "script" that's executed as part of the acceptance test. func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User) { cloudEnv := os.Getenv("CLOUD_ENV") recordRequests := isTruePtr(config.RecordRequests) // If we are running in a cloud environment AND we are recording requests, - // start a dedicated server to act as a reverse proxy and record requests. + // start a dedicated server to act as a reverse proxy to a real Databricks workspace. if cloudEnv != "" && recordRequests { // Use a non-proxy client to fetch user info so that this API call is not recorded // in out.requests.txt @@ -55,16 +55,13 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o user, err := w.CurrentUser.Me(context.Background()) require.NoError(t, err, "Failed to get current user") - // Start a proxy server that sits in front of of a real Databricks workspace. - cfg := startProxyServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + cfg := startProxyServer(t, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) return cfg, *user } // If we are running on a cloud environment, use the host configured in the // environment. if cloudEnv != "" { - // Use a non-proxy client to fetch user info so that this API call is not recorded - // in out.requests.txt w, err := databricks.NewWorkspaceClient() require.NoError(t, err) @@ -74,7 +71,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o return w.Config, *user } - // If we are not recording requests, and no custom server server stubs are configured, + // If we are not recording requests, and no custom server stubs are configured, // use the default shared server. if len(config.Server) == 0 && !recordRequests { // Use a unique token for each test. This allows us to maintain @@ -90,6 +87,8 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o return cfg, TestUser } + // Default case. Start a dedicated local server for the test with the server stubs configured + // as overrides. cfg := startLocalServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) // For the purposes of replacements, use testUser for local runs. @@ -170,7 +169,6 @@ func startLocalServer(t *testing.T, } func startProxyServer(t *testing.T, - stubs []ServerStub, recordRequests bool, logRequests bool, includeHeaders []string, From c4e5aa88504071713569d4a8345658fb8cb9bb6b Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 25 Apr 2025 16:52:23 +0200 Subject: [PATCH 31/57] final cleanup --- acceptance/internal/prepare_server.go | 2 -- acceptance/selftest/record_cloud/workspace-file-io/script | 2 +- libs/testserver/proxy_server.go | 8 +++++--- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 7b82b9b81d..49ba37622a 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -162,8 +162,6 @@ func startLocalServer(t *testing.T, Host: s.URL, Token: "dbapi123", } - err := cfg.EnsureResolved() - require.NoError(t, err) return cfg } diff --git a/acceptance/selftest/record_cloud/workspace-file-io/script b/acceptance/selftest/record_cloud/workspace-file-io/script index c44e67e7de..911b6a6df2 100644 --- a/acceptance/selftest/record_cloud/workspace-file-io/script +++ b/acceptance/selftest/record_cloud/workspace-file-io/script @@ -1,6 +1,6 @@ export username=$($CLI current-user me | jq -r .userName) -# MSYS automatically converts absolute paths like /Users/$username/$UNIQUE_NAME to +# MSYS2 automatically converts absolute paths like /Users/$username/$UNIQUE_NAME to # C:/Program Files/Git/Users/$username/UNIQUE_NAME before passing it to the CLI # Setting this environment variable prevents that conversion on windows. export MSYS_NO_PATHCONV=1 diff --git a/libs/testserver/proxy_server.go b/libs/testserver/proxy_server.go index 451f5f6ccc..953a7c5b6b 100644 --- a/libs/testserver/proxy_server.go +++ b/libs/testserver/proxy_server.go @@ -108,8 +108,9 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { continue } // The default HTTP client in Go sets the Accept-Encoding header to - // "gzip". Since it's meant for the server and will again be set by - // the SDK, do not pass it along here. + // "gzip". Since it is originally meant to be read by the server and + // will be set again when the SDK makes the request to the workspace, do + // not pass it along here. if k == "Accept-Encoding" { continue } @@ -126,7 +127,8 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { err := s.apiClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, reqBody, &respBody) // API errors from the SDK are expected to be of the type [apierr.APIError]. If we - // get an API error then parse the error and forward it in an appropriate format. + // get an API error then parse the error and forward it back to the client + // in an appropriate format. apiErr := &apierr.APIError{} if errors.As(err, &apiErr) { body := map[string]string{ From c72a6d7c42342417bbaf4013df3b65c2784c2ea9 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 5 May 2025 14:52:49 +0200 Subject: [PATCH 32/57] always record requests in proxy --- acceptance/internal/prepare_server.go | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 49ba37622a..16c024df91 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -55,7 +55,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o user, err := w.CurrentUser.Me(context.Background()) require.NoError(t, err, "Failed to get current user") - cfg := startProxyServer(t, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + cfg := startProxyServer(t, logRequests, config.IncludeRequestHeaders, outputDir) return cfg, *user } @@ -167,18 +167,14 @@ func startLocalServer(t *testing.T, } func startProxyServer(t *testing.T, - recordRequests bool, logRequests bool, includeHeaders []string, outputDir string, ) *sdkconfig.Config { s := testserver.NewProxyServer(t) - // Record API requests in out.requests.txt if RecordRequests is true - // in test.toml - if recordRequests { - s.SetRequestCallback(recordRequestsCallback(t, includeHeaders, outputDir)) - } + // Always record requests for a proxy server. + s.SetRequestCallback(recordRequestsCallback(t, includeHeaders, outputDir)) // Log API responses if the -logrequests flag is set. if logRequests { From 55767f18c69795c16934b02efe334c10e985b297 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 5 May 2025 14:58:15 +0200 Subject: [PATCH 33/57] common code for current-user in PrepareServerAndClient --- acceptance/internal/prepare_server.go | 25 ++++++++----------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 16c024df91..61012c8410 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -44,31 +44,22 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o cloudEnv := os.Getenv("CLOUD_ENV") recordRequests := isTruePtr(config.RecordRequests) - // If we are running in a cloud environment AND we are recording requests, - // start a dedicated server to act as a reverse proxy to a real Databricks workspace. - if cloudEnv != "" && recordRequests { - // Use a non-proxy client to fetch user info so that this API call is not recorded - // in out.requests.txt + if cloudEnv != "" { w, err := databricks.NewWorkspaceClient() require.NoError(t, err) user, err := w.CurrentUser.Me(context.Background()) require.NoError(t, err, "Failed to get current user") - cfg := startProxyServer(t, logRequests, config.IncludeRequestHeaders, outputDir) - return cfg, *user - } + cfg := w.Config - // If we are running on a cloud environment, use the host configured in the - // environment. - if cloudEnv != "" { - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) - - user, err := w.CurrentUser.Me(context.Background()) - require.NoError(t, err, "Failed to get current user") + // If we are running in a cloud environment AND we are recording requests, + // start a dedicated server to act as a reverse proxy to a real Databricks workspace. + if recordRequests { + cfg = startProxyServer(t, logRequests, config.IncludeRequestHeaders, outputDir) + } - return w.Config, *user + return cfg, *user } // If we are not recording requests, and no custom server stubs are configured, From 573e82cb22d3dee14ca7fd74ff8b87f392976cd6 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 5 May 2025 15:02:44 +0200 Subject: [PATCH 34/57] - --- libs/testserver/server.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index d864ff4ec8..7a1adfc341 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -18,6 +18,12 @@ type Server interface { SetResponseCallback(func(request *Request, response *EncodedResponse)) } +// Both ProxyServer and LocalServer implement the server interface. +var ( + _ Server = &ProxyServer{} + _ Server = &LocalServer{} +) + type HandlerFunc func(req Request) any type Request struct { From 4bd108333ed0c1d4bf0b3d09a63c13778c22c036 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Tue, 6 May 2025 16:31:25 +0200 Subject: [PATCH 35/57] have start functions return host and token directly --- acceptance/internal/prepare_server.go | 31 ++++++++++++--------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 61012c8410..10ff4fd4bb 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -56,7 +56,11 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o // If we are running in a cloud environment AND we are recording requests, // start a dedicated server to act as a reverse proxy to a real Databricks workspace. if recordRequests { - cfg = startProxyServer(t, logRequests, config.IncludeRequestHeaders, outputDir) + host, token := startProxyServer(t, logRequests, config.IncludeRequestHeaders, outputDir) + cfg = &sdkconfig.Config{ + Host: host, + Token: token, + } } return cfg, *user @@ -80,7 +84,11 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o // Default case. Start a dedicated local server for the test with the server stubs configured // as overrides. - cfg := startLocalServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + host, token := startLocalServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + cfg := &sdkconfig.Config{ + Host: host, + Token: token, + } // For the purposes of replacements, use testUser for local runs. // Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement: @@ -119,7 +127,7 @@ func startLocalServer(t *testing.T, logRequests bool, includeHeaders []string, outputDir string, -) *sdkconfig.Config { +) (string, string) { s := testserver.NewLocalServer(t) // Record API requests in out.requests.txt if RecordRequests is true @@ -148,20 +156,14 @@ func startLocalServer(t *testing.T, // The earliest handlers take precedence, add default handlers last addDefaultHandlers(s) - - cfg := &sdkconfig.Config{ - Host: s.URL, - Token: "dbapi123", - } - - return cfg + return s.URL, "dbapi123" } func startProxyServer(t *testing.T, logRequests bool, includeHeaders []string, outputDir string, -) *sdkconfig.Config { +) (string, string) { s := testserver.NewProxyServer(t) // Always record requests for a proxy server. @@ -172,12 +174,7 @@ func startProxyServer(t *testing.T, s.SetResponseCallback(logResponseCallback(t)) } - cfg := &sdkconfig.Config{ - Host: s.URL, - Token: "dbapi1234", - } - - return cfg + return s.URL, "dbapi1234" } type LoggedRequest struct { From e94556d6cb5b78faee765e8751d2abb8916c1320 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Tue, 6 May 2025 17:56:41 +0200 Subject: [PATCH 36/57] undo local server changse --- acceptance/internal/cmd_server.go | 4 +- acceptance/internal/handlers.go | 2 +- acceptance/internal/prepare_server.go | 8 +- libs/telemetry/logger_test.go | 6 +- libs/testserver/local_server.go | 135 -------------------------- libs/testserver/server.go | 122 +++++++++++++++++++++-- 6 files changed, 123 insertions(+), 154 deletions(-) delete mode 100644 libs/testserver/local_server.go diff --git a/acceptance/internal/cmd_server.go b/acceptance/internal/cmd_server.go index c776a668f5..69d4f02ce3 100644 --- a/acceptance/internal/cmd_server.go +++ b/acceptance/internal/cmd_server.go @@ -12,8 +12,8 @@ import ( "github.com/stretchr/testify/require" ) -func StartCmdServer(t *testing.T) *testserver.LocalServer { - server := testserver.NewLocalServer(t) +func StartCmdServer(t *testing.T) *testserver.Server { + server := testserver.New(t) server.Handle("GET", "/", func(r testserver.Request) any { q := r.URL.Query() args := strings.Split(q.Get("args"), " ") diff --git a/acceptance/internal/handlers.go b/acceptance/internal/handlers.go index 16caa376f4..1d2b6ab5ad 100644 --- a/acceptance/internal/handlers.go +++ b/acceptance/internal/handlers.go @@ -28,7 +28,7 @@ var TestMetastore = catalog.MetastoreAssignment{ WorkspaceId: 470123456789500, } -func addDefaultHandlers(server *testserver.LocalServer) { +func addDefaultHandlers(server *testserver.Server) { server.Handle("GET", "/api/2.0/policies/clusters/list", func(req testserver.Request) any { return compute.ListPoliciesResponse{ Policies: []compute.Policy{ diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 10ff4fd4bb..b79ce7bb25 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -24,7 +24,7 @@ import ( ) func StartDefaultServer(t *testing.T) { - s := testserver.NewLocalServer(t) + s := testserver.New(t) addDefaultHandlers(s) t.Setenv("DATABRICKS_DEFAULT_HOST", s.URL) @@ -128,17 +128,17 @@ func startLocalServer(t *testing.T, includeHeaders []string, outputDir string, ) (string, string) { - s := testserver.NewLocalServer(t) + s := testserver.New(t) // Record API requests in out.requests.txt if RecordRequests is true // in test.toml if recordRequests { - s.SetRequestCallback(recordRequestsCallback(t, includeHeaders, outputDir)) + s.RequestCallback = recordRequestsCallback(t, includeHeaders, outputDir) } // Log API responses if the -logrequests flag is set. if logRequests { - s.SetResponseCallback(logResponseCallback(t)) + s.ResponseCallback = logResponseCallback(t) } for ind := range stubs { diff --git a/libs/telemetry/logger_test.go b/libs/telemetry/logger_test.go index cda3f7724b..f0f2a10020 100644 --- a/libs/telemetry/logger_test.go +++ b/libs/telemetry/logger_test.go @@ -14,7 +14,7 @@ import ( ) func TestTelemetryUploadRetriesOnPartialSuccess(t *testing.T) { - server := testserver.NewLocalServer(t) + server := testserver.New(t) t.Cleanup(server.Close) count := 0 @@ -57,7 +57,7 @@ func TestTelemetryUploadRetriesOnPartialSuccess(t *testing.T) { } func uploadRetriesFor(t *testing.T, statusCode int) { - server := testserver.NewLocalServer(t) + server := testserver.New(t) t.Cleanup(server.Close) count := 0 @@ -118,7 +118,7 @@ func TestTelemetryUploadRetriesForStatusCodes(t *testing.T) { } func TestTelemetryUploadMaxRetries(t *testing.T) { - server := testserver.NewLocalServer(t) + server := testserver.New(t) t.Cleanup(server.Close) count := 0 diff --git a/libs/testserver/local_server.go b/libs/testserver/local_server.go deleted file mode 100644 index a429a8c4fc..0000000000 --- a/libs/testserver/local_server.go +++ /dev/null @@ -1,135 +0,0 @@ -package testserver - -import ( - "encoding/json" - "fmt" - "io" - "net/http" - "net/http/httptest" - "sync" - - "github.com/gorilla/mux" - - "github.com/databricks/cli/internal/testutil" - "github.com/databricks/databricks-sdk-go/apierr" -) - -type LocalServer struct { - *httptest.Server - - t testutil.TestingT - mu *sync.Mutex - router *mux.Router - - fakeWorkspaces map[string]*FakeWorkspace - requestCallback func(request *Request) - responseCallback func(request *Request, response *EncodedResponse) -} - -func NewLocalServer(t testutil.TestingT) *LocalServer { - router := mux.NewRouter() - server := httptest.NewServer(router) - t.Cleanup(server.Close) - - s := &LocalServer{ - Server: server, - router: router, - t: t, - mu: &sync.Mutex{}, - fakeWorkspaces: map[string]*FakeWorkspace{}, - } - - // Set up the not found handler as fallback. - s.router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - pattern := r.Method + " " + r.URL.Path - bodyBytes, err := io.ReadAll(r.Body) - var body string - if err != nil { - body = fmt.Sprintf("failed to read the body: %s", err) - } else { - body = fmt.Sprintf("[%d bytes] %s", len(bodyBytes), bodyBytes) - } - - t.Errorf(`No handler for URL: %s -Body: %s - -For acceptance tests, add this to test.toml: -[[Server]] -Pattern = %q -Response.Body = '' -# Response.StatusCode = -`, r.URL, body, pattern) - - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusNotImplemented) - - resp := apierr.APIError{ - Message: "No stub found for pattern: " + pattern, - } - - respBytes, err := json.Marshal(resp) - if err != nil { - t.Errorf("JSON encoding error: %s", err) - respBytes = []byte("{\"message\": \"JSON encoding error\"}") - } - - if _, err := w.Write(respBytes); err != nil { - t.Errorf("Response write error: %s", err) - } - }) - - return s -} - -func (s *LocalServer) Handle(method, path string, handler HandlerFunc) { - s.router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { - // For simplicity we process requests sequentially. It's fast enough because - // we don't do any IO except reading and writing request/response bodies. - s.mu.Lock() - defer s.mu.Unlock() - - // Each test uses unique DATABRICKS_TOKEN, we simulate each token having - // it's own fake fakeWorkspace to avoid interference between tests. - var fakeWorkspace *FakeWorkspace = nil - token := getToken(r) - if token != "" { - if _, ok := s.fakeWorkspaces[token]; !ok { - s.fakeWorkspaces[token] = NewFakeWorkspace() - } - - fakeWorkspace = s.fakeWorkspaces[token] - } - - request := NewRequest(s.t, r, fakeWorkspace) - - if s.requestCallback != nil { - s.requestCallback(&request) - } - - respAny := handler(request) - resp := normalizeResponse(s.t, respAny) - - for k, v := range resp.Headers { - w.Header()[k] = v - } - - w.WriteHeader(resp.StatusCode) - - if s.responseCallback != nil { - s.responseCallback(&request, &resp) - } - - if _, err := w.Write(resp.Body); err != nil { - s.t.Errorf("Failed to write response: %s", err) - return - } - }).Methods(method) -} - -func (s *LocalServer) SetRequestCallback(callback func(request *Request)) { - s.requestCallback = callback -} - -func (s *LocalServer) SetResponseCallback(callback func(request *Request, response *EncodedResponse)) { - s.responseCallback = callback -} diff --git a/libs/testserver/server.go b/libs/testserver/server.go index 7a1adfc341..ebaaa002b8 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -2,27 +2,131 @@ package testserver import ( "encoding/json" + "fmt" "io" "net/http" + "net/http/httptest" "net/url" "reflect" "strings" + "sync" "github.com/databricks/cli/internal/testutil" + "github.com/databricks/databricks-sdk-go/apierr" "github.com/gorilla/mux" ) -type Server interface { - Handle(method, path string, handler HandlerFunc) - SetRequestCallback(func(request *Request)) - SetResponseCallback(func(request *Request, response *EncodedResponse)) +type Server struct { + *httptest.Server + + t testutil.TestingT + mu *sync.Mutex + router *mux.Router + + fakeWorkspaces map[string]*FakeWorkspace + RequestCallback func(request *Request) + ResponseCallback func(request *Request, response *EncodedResponse) } -// Both ProxyServer and LocalServer implement the server interface. -var ( - _ Server = &ProxyServer{} - _ Server = &LocalServer{} -) +func New(t testutil.TestingT) *Server { + router := mux.NewRouter() + server := httptest.NewServer(router) + t.Cleanup(server.Close) + + s := &Server{ + Server: server, + router: router, + t: t, + mu: &sync.Mutex{}, + fakeWorkspaces: map[string]*FakeWorkspace{}, + } + + // Set up the not found handler as fallback. + s.router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + pattern := r.Method + " " + r.URL.Path + bodyBytes, err := io.ReadAll(r.Body) + var body string + if err != nil { + body = fmt.Sprintf("failed to read the body: %s", err) + } else { + body = fmt.Sprintf("[%d bytes] %s", len(bodyBytes), bodyBytes) + } + + t.Errorf(`No handler for URL: %s +Body: %s + +For acceptance tests, add this to test.toml: +[[Server]] +Pattern = %q +Response.Body = '' +# Response.StatusCode = +`, r.URL, body, pattern) + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotImplemented) + + resp := apierr.APIError{ + Message: "No stub found for pattern: " + pattern, + } + + respBytes, err := json.Marshal(resp) + if err != nil { + t.Errorf("JSON encoding error: %s", err) + respBytes = []byte("{\"message\": \"JSON encoding error\"}") + } + + if _, err := w.Write(respBytes); err != nil { + t.Errorf("Response write error: %s", err) + } + }) + + return s +} + +func (s *Server) Handle(method, path string, handler HandlerFunc) { + s.router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + // For simplicity we process requests sequentially. It's fast enough because + // we don't do any IO except reading and writing request/response bodies. + s.mu.Lock() + defer s.mu.Unlock() + + // Each test uses unique DATABRICKS_TOKEN, we simulate each token having + // it's own fake fakeWorkspace to avoid interference between tests. + var fakeWorkspace *FakeWorkspace = nil + token := getToken(r) + if token != "" { + if _, ok := s.fakeWorkspaces[token]; !ok { + s.fakeWorkspaces[token] = NewFakeWorkspace() + } + + fakeWorkspace = s.fakeWorkspaces[token] + } + + request := NewRequest(s.t, r, fakeWorkspace) + + if s.RequestCallback != nil { + s.RequestCallback(&request) + } + + respAny := handler(request) + resp := normalizeResponse(s.t, respAny) + + for k, v := range resp.Headers { + w.Header()[k] = v + } + + w.WriteHeader(resp.StatusCode) + + if s.ResponseCallback != nil { + s.ResponseCallback(&request, &resp) + } + + if _, err := w.Write(resp.Body); err != nil { + s.t.Errorf("Failed to write response: %s", err) + return + } + }).Methods(method) +} type HandlerFunc func(req Request) any From e61ecc6c93e53abf8b3767d70c77b16c35662d63 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Tue, 6 May 2025 17:59:25 +0200 Subject: [PATCH 37/57] reduce diff --- libs/testserver/server.go | 34 ++++++++++++++++++---------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index ebaaa002b8..18896ffc97 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -11,23 +11,34 @@ import ( "strings" "sync" + "github.com/gorilla/mux" + "github.com/databricks/cli/internal/testutil" "github.com/databricks/databricks-sdk-go/apierr" - "github.com/gorilla/mux" ) type Server struct { *httptest.Server - - t testutil.TestingT - mu *sync.Mutex router *mux.Router - fakeWorkspaces map[string]*FakeWorkspace + t testutil.TestingT + + fakeWorkspaces map[string]*FakeWorkspace + mu *sync.Mutex + RequestCallback func(request *Request) ResponseCallback func(request *Request, response *EncodedResponse) } +type Request struct { + Method string + URL *url.URL + Headers http.Header + Body []byte + Vars map[string]string + Workspace *FakeWorkspace +} + func New(t testutil.TestingT) *Server { router := mux.NewRouter() server := httptest.NewServer(router) @@ -83,6 +94,8 @@ Response.Body = '' return s } +type HandlerFunc func(req Request) any + func (s *Server) Handle(method, path string, handler HandlerFunc) { s.router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { // For simplicity we process requests sequentially. It's fast enough because @@ -128,17 +141,6 @@ func (s *Server) Handle(method, path string, handler HandlerFunc) { }).Methods(method) } -type HandlerFunc func(req Request) any - -type Request struct { - Method string - URL *url.URL - Headers http.Header - Body []byte - Vars map[string]string - Workspace *FakeWorkspace -} - type Response struct { StatusCode int Headers http.Header From bc39366212499cdba3e5fbfa9a5f8b08f36f7570 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Tue, 6 May 2025 18:00:04 +0200 Subject: [PATCH 38/57] - --- libs/testserver/server.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index 18896ffc97..d3127ea956 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -39,6 +39,12 @@ type Request struct { Workspace *FakeWorkspace } +type Response struct { + StatusCode int + Headers http.Header + Body any +} + func New(t testutil.TestingT) *Server { router := mux.NewRouter() server := httptest.NewServer(router) @@ -141,12 +147,6 @@ func (s *Server) Handle(method, path string, handler HandlerFunc) { }).Methods(method) } -type Response struct { - StatusCode int - Headers http.Header - Body any -} - type EncodedResponse struct { StatusCode int Headers http.Header From 38e8901fe1481479b913fc376fb0d6c4cab0ce39 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Tue, 6 May 2025 18:01:40 +0200 Subject: [PATCH 39/57] - --- libs/testserver/server.go | 205 +++++++++++++++++++------------------- 1 file changed, 102 insertions(+), 103 deletions(-) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index d3127ea956..33931340a2 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -44,109 +44,6 @@ type Response struct { Headers http.Header Body any } - -func New(t testutil.TestingT) *Server { - router := mux.NewRouter() - server := httptest.NewServer(router) - t.Cleanup(server.Close) - - s := &Server{ - Server: server, - router: router, - t: t, - mu: &sync.Mutex{}, - fakeWorkspaces: map[string]*FakeWorkspace{}, - } - - // Set up the not found handler as fallback. - s.router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - pattern := r.Method + " " + r.URL.Path - bodyBytes, err := io.ReadAll(r.Body) - var body string - if err != nil { - body = fmt.Sprintf("failed to read the body: %s", err) - } else { - body = fmt.Sprintf("[%d bytes] %s", len(bodyBytes), bodyBytes) - } - - t.Errorf(`No handler for URL: %s -Body: %s - -For acceptance tests, add this to test.toml: -[[Server]] -Pattern = %q -Response.Body = '' -# Response.StatusCode = -`, r.URL, body, pattern) - - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusNotImplemented) - - resp := apierr.APIError{ - Message: "No stub found for pattern: " + pattern, - } - - respBytes, err := json.Marshal(resp) - if err != nil { - t.Errorf("JSON encoding error: %s", err) - respBytes = []byte("{\"message\": \"JSON encoding error\"}") - } - - if _, err := w.Write(respBytes); err != nil { - t.Errorf("Response write error: %s", err) - } - }) - - return s -} - -type HandlerFunc func(req Request) any - -func (s *Server) Handle(method, path string, handler HandlerFunc) { - s.router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { - // For simplicity we process requests sequentially. It's fast enough because - // we don't do any IO except reading and writing request/response bodies. - s.mu.Lock() - defer s.mu.Unlock() - - // Each test uses unique DATABRICKS_TOKEN, we simulate each token having - // it's own fake fakeWorkspace to avoid interference between tests. - var fakeWorkspace *FakeWorkspace = nil - token := getToken(r) - if token != "" { - if _, ok := s.fakeWorkspaces[token]; !ok { - s.fakeWorkspaces[token] = NewFakeWorkspace() - } - - fakeWorkspace = s.fakeWorkspaces[token] - } - - request := NewRequest(s.t, r, fakeWorkspace) - - if s.RequestCallback != nil { - s.RequestCallback(&request) - } - - respAny := handler(request) - resp := normalizeResponse(s.t, respAny) - - for k, v := range resp.Headers { - w.Header()[k] = v - } - - w.WriteHeader(resp.StatusCode) - - if s.ResponseCallback != nil { - s.ResponseCallback(&request, &resp) - } - - if _, err := w.Write(resp.Body); err != nil { - s.t.Errorf("Failed to write response: %s", err) - return - } - }).Methods(method) -} - type EncodedResponse struct { StatusCode int Headers http.Header @@ -279,6 +176,108 @@ func getHeaders(value []byte) http.Header { } } +func New(t testutil.TestingT) *Server { + router := mux.NewRouter() + server := httptest.NewServer(router) + t.Cleanup(server.Close) + + s := &Server{ + Server: server, + router: router, + t: t, + mu: &sync.Mutex{}, + fakeWorkspaces: map[string]*FakeWorkspace{}, + } + + // Set up the not found handler as fallback. + s.router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + pattern := r.Method + " " + r.URL.Path + bodyBytes, err := io.ReadAll(r.Body) + var body string + if err != nil { + body = fmt.Sprintf("failed to read the body: %s", err) + } else { + body = fmt.Sprintf("[%d bytes] %s", len(bodyBytes), bodyBytes) + } + + t.Errorf(`No handler for URL: %s +Body: %s + +For acceptance tests, add this to test.toml: +[[Server]] +Pattern = %q +Response.Body = '' +# Response.StatusCode = +`, r.URL, body, pattern) + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotImplemented) + + resp := apierr.APIError{ + Message: "No stub found for pattern: " + pattern, + } + + respBytes, err := json.Marshal(resp) + if err != nil { + t.Errorf("JSON encoding error: %s", err) + respBytes = []byte("{\"message\": \"JSON encoding error\"}") + } + + if _, err := w.Write(respBytes); err != nil { + t.Errorf("Response write error: %s", err) + } + }) + + return s +} + +type HandlerFunc func(req Request) any + +func (s *Server) Handle(method, path string, handler HandlerFunc) { + s.router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + // For simplicity we process requests sequentially. It's fast enough because + // we don't do any IO except reading and writing request/response bodies. + s.mu.Lock() + defer s.mu.Unlock() + + // Each test uses unique DATABRICKS_TOKEN, we simulate each token having + // it's own fake fakeWorkspace to avoid interference between tests. + var fakeWorkspace *FakeWorkspace = nil + token := getToken(r) + if token != "" { + if _, ok := s.fakeWorkspaces[token]; !ok { + s.fakeWorkspaces[token] = NewFakeWorkspace() + } + + fakeWorkspace = s.fakeWorkspaces[token] + } + + request := NewRequest(s.t, r, fakeWorkspace) + + if s.RequestCallback != nil { + s.RequestCallback(&request) + } + + respAny := handler(request) + resp := normalizeResponse(s.t, respAny) + + for k, v := range resp.Headers { + w.Header()[k] = v + } + + w.WriteHeader(resp.StatusCode) + + if s.ResponseCallback != nil { + s.ResponseCallback(&request, &resp) + } + + if _, err := w.Write(resp.Body); err != nil { + s.t.Errorf("Failed to write response: %s", err) + return + } + }).Methods(method) +} + func getToken(r *http.Request) string { header := r.Header.Get("Authorization") prefix := "Bearer " From 8e37d99b119c9ea76a405a14bcd1c32be973ded0 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Tue, 6 May 2025 18:07:26 +0200 Subject: [PATCH 40/57] separate package for the proxy server --- acceptance/internal/prepare_server.go | 7 +- libs/testproxy/server.go | 186 ++++++++++++++++++++++++++ 2 files changed, 190 insertions(+), 3 deletions(-) create mode 100644 libs/testproxy/server.go diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index b79ce7bb25..3ac2aec609 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -14,6 +14,7 @@ import ( "unicode/utf8" "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/testproxy" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" sdkconfig "github.com/databricks/databricks-sdk-go/config" @@ -164,14 +165,14 @@ func startProxyServer(t *testing.T, includeHeaders []string, outputDir string, ) (string, string) { - s := testserver.NewProxyServer(t) + s := testproxy.New(t) // Always record requests for a proxy server. - s.SetRequestCallback(recordRequestsCallback(t, includeHeaders, outputDir)) + s.RequestCallback = recordRequestsCallback(t, includeHeaders, outputDir) // Log API responses if the -logrequests flag is set. if logRequests { - s.SetResponseCallback(logResponseCallback(t)) + s.ResponseCallback = logResponseCallback(t) } return s.URL, "dbapi1234" diff --git a/libs/testproxy/server.go b/libs/testproxy/server.go new file mode 100644 index 0000000000..5553013eb5 --- /dev/null +++ b/libs/testproxy/server.go @@ -0,0 +1,186 @@ +package testproxy + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "sync" + + "github.com/databricks/cli/internal/testutil" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/config" + "github.com/gorilla/mux" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type ProxyServer struct { + *httptest.Server + + t testutil.TestingT + mu *sync.Mutex + router *mux.Router + + apiClient *client.DatabricksClient + RequestCallback func(request *testserver.Request) + ResponseCallback func(request *testserver.Request, response *testserver.EncodedResponse) +} + +// This creates a reverse proxy server that sits in front of a real Databricks +// workspace. This is useful for recording API requests and responses in +// integration tests. +// +// Note: We cannot simply proxy the request from a localhost URL to a real +// workspace. This is because auth resolution in the Databricks SDK relies +// what the URL actually looks like to determine the auth method to use. +// For example, in OAuth flows, the SDK can make requests to different Microsoft +// OAuth endpoints based on the nature of the URL. +// For reference, see: +// https://github.com/databricks/databricks-sdk-go/blob/79e4b3a6e9b0b7dcb1af9ad4025deb447b01d933/common/environment/environments.go#L57 +func New(t testutil.TestingT) *ProxyServer { + router := mux.NewRouter() + server := httptest.NewServer(router) + t.Cleanup(server.Close) + + s := &ProxyServer{ + Server: server, + t: t, + router: router, + mu: &sync.Mutex{}, + } + + // Create an API client using the current authentication context. + // In CI test environments this would read the appropriate environment + // variables. + var err error + s.apiClient, err = client.New(&config.Config{}) + require.NoError(t, err) + + // Set up the proxy handler as the default handler for all requests. + router.NotFoundHandler = http.HandlerFunc(s.proxyToCloud) + return s +} + +func (s *ProxyServer) reqBody(r testserver.Request) any { + // The SDK expects the query parameters to be specified in the "request body" + // argument for GET, DELETE, and HEAD requests in the .Do method. + if r.Method == "GET" || r.Method == "DELETE" || r.Method == "HEAD" { + queryParams := make(map[string]any) + for k, v := range r.URL.Query() { + queryParams[k] = v[0] + } + return queryParams + } + + // The SDK does not support directly passing a JSON serialized request + // body. So we convert it to a map if the body is a JSON object. + if json.Valid(r.Body) { + m := make(map[string]any) + err := json.Unmarshal(r.Body, &m) + assert.NoError(s.t, err) + return m + } + + // Otherwise, return the raw body. + return r.Body +} + +func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { + // Process requests sequentially. It's slower but is easier to reason about. + // For example, requestCallback and responseCallback functions do not need + // to be thread-safe. + s.mu.Lock() + defer s.mu.Unlock() + + request := testserver.NewRequest(s.t, r, nil) + if s.RequestCallback != nil { + s.RequestCallback(&request) + } + + headers := make(map[string]string) + for k, v := range r.Header { + // Authorization headers will be set by the SDK. Do not pass them along here. + if k == "Authorization" { + continue + } + // The default HTTP client in Go sets the Accept-Encoding header to + // "gzip". Since it is originally meant to be read by the server and + // will be set again when the SDK makes the request to the workspace, do + // not pass it along here. + if k == "Accept-Encoding" { + continue + } + headers[k] = v[0] + } + + queryParams := make(map[string]any) + for k, v := range r.URL.Query() { + queryParams[k] = v[0] + } + + reqBody := s.reqBody(request) + respBody := bytes.Buffer{} + err := s.apiClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, reqBody, &respBody) + + // API errors from the SDK are expected to be of the type [apierr.APIError]. If we + // get an API error then parse the error and forward it back to the client + // in an appropriate format. + apiErr := &apierr.APIError{} + if errors.As(err, &apiErr) { + body := map[string]string{ + "error_code": apiErr.ErrorCode, + "message": apiErr.Message, + } + + b, err := json.Marshal(body) + assert.NoError(s.t, err) + + w.WriteHeader(apiErr.StatusCode) + _, err = w.Write(b) + assert.NoError(s.t, err) + + if s.ResponseCallback != nil { + s.ResponseCallback(&request, &testserver.EncodedResponse{ + StatusCode: apiErr.StatusCode, + Body: []byte(apiErr.Message), + }) + } + + return + } + + // Something else went wrong. + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + _, err := w.Write([]byte(err.Error())) + assert.NoError(s.t, err) + + if s.ResponseCallback != nil { + s.ResponseCallback(&request, &testserver.EncodedResponse{ + StatusCode: 500, + Body: []byte(err.Error()), + }) + } + + return + } + + // Successful response + w.WriteHeader(200) + b := respBody.Bytes() + + _, err = w.Write(b) + assert.NoError(s.t, err) + + if s.ResponseCallback != nil { + s.ResponseCallback(&request, &testserver.EncodedResponse{ + StatusCode: 200, + Body: b, + }) + } +} From b86da3ce5fd7a71ede89399f78f9e511c1503595 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Tue, 6 May 2025 18:07:53 +0200 Subject: [PATCH 41/57] - --- libs/testserver/proxy_server.go | 199 -------------------------------- 1 file changed, 199 deletions(-) delete mode 100644 libs/testserver/proxy_server.go diff --git a/libs/testserver/proxy_server.go b/libs/testserver/proxy_server.go deleted file mode 100644 index 953a7c5b6b..0000000000 --- a/libs/testserver/proxy_server.go +++ /dev/null @@ -1,199 +0,0 @@ -package testserver - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "net/http" - "net/http/httptest" - "sync" - - "github.com/databricks/cli/internal/testutil" - "github.com/databricks/databricks-sdk-go/apierr" - "github.com/databricks/databricks-sdk-go/client" - "github.com/databricks/databricks-sdk-go/config" - "github.com/gorilla/mux" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -type ProxyServer struct { - *httptest.Server - - t testutil.TestingT - mu *sync.Mutex - router *mux.Router - - apiClient *client.DatabricksClient - requestCallback func(request *Request) - responseCallback func(request *Request, response *EncodedResponse) -} - -// This creates a reverse proxy server that sits in front of a real Databricks -// workspace. This is useful for recording API requests and responses in -// integration tests. -// -// Note: We cannot simply proxy the request from a localhost URL to a real -// workspace. This is because auth resolution in the Databricks SDK relies -// what the URL actually looks like to determine the auth method to use. -// For example, in OAuth flows, the SDK can make requests to different Microsoft -// OAuth endpoints based on the nature of the URL. -// For reference, see: -// https://github.com/databricks/databricks-sdk-go/blob/79e4b3a6e9b0b7dcb1af9ad4025deb447b01d933/common/environment/environments.go#L57 -func NewProxyServer(t testutil.TestingT) *ProxyServer { - router := mux.NewRouter() - server := httptest.NewServer(router) - t.Cleanup(server.Close) - - s := &ProxyServer{ - Server: server, - t: t, - router: router, - mu: &sync.Mutex{}, - } - - // Create an API client using the current authentication context. - // In CI test environments this would read the appropriate environment - // variables. - var err error - s.apiClient, err = client.New(&config.Config{}) - require.NoError(t, err) - - // Set up the proxy handler as the default handler for all requests. - router.NotFoundHandler = http.HandlerFunc(s.proxyToCloud) - return s -} - -func (s *ProxyServer) reqBody(r Request) any { - // The SDK expects the query parameters to be specified in the "request body" - // argument for GET, DELETE, and HEAD requests in the .Do method. - if r.Method == "GET" || r.Method == "DELETE" || r.Method == "HEAD" { - queryParams := make(map[string]any) - for k, v := range r.URL.Query() { - queryParams[k] = v[0] - } - return queryParams - } - - // The SDK does not support directly passing a JSON serialized request - // body. So we convert it to a map if the body is a JSON object. - if json.Valid(r.Body) { - m := make(map[string]any) - err := json.Unmarshal(r.Body, &m) - assert.NoError(s.t, err) - return m - } - - // Otherwise, return the raw body. - return r.Body -} - -func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { - // Process requests sequentially. It's slower but is easier to reason about. - // For example, requestCallback and responseCallback functions do not need - // to be thread-safe. - s.mu.Lock() - defer s.mu.Unlock() - - request := NewRequest(s.t, r, nil) - if s.requestCallback != nil { - s.requestCallback(&request) - } - - headers := make(map[string]string) - for k, v := range r.Header { - // Authorization headers will be set by the SDK. Do not pass them along here. - if k == "Authorization" { - continue - } - // The default HTTP client in Go sets the Accept-Encoding header to - // "gzip". Since it is originally meant to be read by the server and - // will be set again when the SDK makes the request to the workspace, do - // not pass it along here. - if k == "Accept-Encoding" { - continue - } - headers[k] = v[0] - } - - queryParams := make(map[string]any) - for k, v := range r.URL.Query() { - queryParams[k] = v[0] - } - - reqBody := s.reqBody(request) - respBody := bytes.Buffer{} - err := s.apiClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, reqBody, &respBody) - - // API errors from the SDK are expected to be of the type [apierr.APIError]. If we - // get an API error then parse the error and forward it back to the client - // in an appropriate format. - apiErr := &apierr.APIError{} - if errors.As(err, &apiErr) { - body := map[string]string{ - "error_code": apiErr.ErrorCode, - "message": apiErr.Message, - } - - b, err := json.Marshal(body) - assert.NoError(s.t, err) - - w.WriteHeader(apiErr.StatusCode) - _, err = w.Write(b) - assert.NoError(s.t, err) - - if s.responseCallback != nil { - s.responseCallback(&request, &EncodedResponse{ - StatusCode: apiErr.StatusCode, - Body: []byte(apiErr.Message), - }) - } - - return - } - - // Something else went wrong. - if err != nil { - w.WriteHeader(http.StatusInternalServerError) - _, err := w.Write([]byte(err.Error())) - assert.NoError(s.t, err) - - if s.responseCallback != nil { - s.responseCallback(&request, &EncodedResponse{ - StatusCode: 500, - Body: []byte(err.Error()), - }) - } - - return - } - - // Successful response - w.WriteHeader(200) - b := respBody.Bytes() - - _, err = w.Write(b) - assert.NoError(s.t, err) - - if s.responseCallback != nil { - s.responseCallback(&request, &EncodedResponse{ - StatusCode: 200, - Body: b, - }) - } -} - -// Eventually we can implement this function to allow for per-test overrides -// even in integration tests. -func (s *ProxyServer) Handle(method, path string, handler HandlerFunc) { - s.t.Fatalf("Not implemented") -} - -func (s *ProxyServer) SetRequestCallback(callback func(request *Request)) { - s.requestCallback = callback -} - -func (s *ProxyServer) SetResponseCallback(callback func(request *Request, response *EncodedResponse)) { - s.responseCallback = callback -} From a2639f00079cdf8f0b41ede25ebec2502a7f7fc5 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Tue, 6 May 2025 18:20:27 +0200 Subject: [PATCH 42/57] - --- acceptance/internal/prepare_server.go | 3 --- libs/testdiff/replacement.go | 5 ----- 2 files changed, 8 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 668ca1f6ee..3ac2aec609 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -13,9 +13,6 @@ import ( "time" "unicode/utf8" - sdkconfig "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/service/iam" - "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testproxy" "github.com/databricks/cli/libs/testserver" diff --git a/libs/testdiff/replacement.go b/libs/testdiff/replacement.go index 7262cb8936..61889b83ba 100644 --- a/libs/testdiff/replacement.go +++ b/libs/testdiff/replacement.go @@ -137,11 +137,6 @@ func (r *ReplacementsContext) SetPathWithParents(old, new string) { r.SetPath(filepath.Dir(filepath.Dir(old)), new+"_GPARENT") } -func PrepareReplacementsWorkspaceConfig(t testutil.TestingT, r *ReplacementsContext, cfg *config.Config) { - t.Helper() - PrepareReplacementsWorkspaceConfig(t, r, w.Config) -} - func PrepareReplacementsWorkspaceConfig(t testutil.TestingT, r *ReplacementsContext, cfg *config.Config) { t.Helper() // in some clouds (gcp) w.Config.Host includes "https://" prefix in others it's really just a host (azure) From ca98d884d4ff522c1483d9ebe50d5f6e96ba450f Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 11:26:39 +0200 Subject: [PATCH 43/57] - --- acceptance/internal/prepare_server.go | 11 +++++++++++ libs/testproxy/server.go | 6 ------ 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 3ac2aec609..f12df0d35e 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -9,6 +9,7 @@ import ( "path/filepath" "slices" "strings" + "sync" "testing" "time" "unicode/utf8" @@ -97,7 +98,12 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o } func recordRequestsCallback(t *testing.T, includeHeaders []string, outputDir string) func(request *testserver.Request) { + mu := sync.Mutex{} + return func(request *testserver.Request) { + mu.Lock() + defer mu.Unlock() + req := getLoggedRequest(request, includeHeaders) reqJson, err := json.MarshalIndent(req, "", " ") assert.NoErrorf(t, err, "Failed to json-encode: %#v", req) @@ -113,7 +119,12 @@ func recordRequestsCallback(t *testing.T, includeHeaders []string, outputDir str } func logResponseCallback(t *testing.T) func(request *testserver.Request, response *testserver.EncodedResponse) { + mu := sync.Mutex{} + return func(request *testserver.Request, response *testserver.EncodedResponse) { + mu.Lock() + defer mu.Unlock() + t.Logf("%d %s %s\n%s\n%s", response.StatusCode, request.Method, request.URL, formatHeadersAndBody("> ", request.Headers, request.Body), diff --git a/libs/testproxy/server.go b/libs/testproxy/server.go index 5553013eb5..0ecf899d2f 100644 --- a/libs/testproxy/server.go +++ b/libs/testproxy/server.go @@ -91,12 +91,6 @@ func (s *ProxyServer) reqBody(r testserver.Request) any { } func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { - // Process requests sequentially. It's slower but is easier to reason about. - // For example, requestCallback and responseCallback functions do not need - // to be thread-safe. - s.mu.Lock() - defer s.mu.Unlock() - request := testserver.NewRequest(s.t, r, nil) if s.RequestCallback != nil { s.RequestCallback(&request) From 0f5a61c2eedaad61683737115212709ff63cd670 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 11:57:32 +0200 Subject: [PATCH 44/57] - --- acceptance/internal/prepare_server.go | 5 ----- 1 file changed, 5 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index f12df0d35e..644b689c49 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -119,12 +119,7 @@ func recordRequestsCallback(t *testing.T, includeHeaders []string, outputDir str } func logResponseCallback(t *testing.T) func(request *testserver.Request, response *testserver.EncodedResponse) { - mu := sync.Mutex{} - return func(request *testserver.Request, response *testserver.EncodedResponse) { - mu.Lock() - defer mu.Unlock() - t.Logf("%d %s %s\n%s\n%s", response.StatusCode, request.Method, request.URL, formatHeadersAndBody("> ", request.Headers, request.Body), From f326379a20e3dc46dde4bf3e6cce2522d2f07459 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 12:10:37 +0200 Subject: [PATCH 45/57] - --- .../selftest/record_cloud/pipeline-crud/out.requests.txt | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt b/acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt index c5c203e309..5d69eb6fd1 100644 --- a/acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt +++ b/acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt @@ -29,8 +29,7 @@ } } ], - "name": "test-pipeline-2", - "pipeline_id": "[UUID]" + "name": "test-pipeline-2" } } { From a7480389444369fe0101d4a21ead20417fce3c5a Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 12:12:55 +0200 Subject: [PATCH 46/57] - --- acceptance/acceptance_test.go | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index f2537b8b49..4b815720c1 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -415,19 +415,6 @@ func runTest(t *testing.T, cmd.Env = append(cmd.Env, "UNIQUE_NAME="+uniqueName) cmd.Env = append(cmd.Env, "TEST_TMP_DIR="+tmpDir) - // In inprocess mode, we need to modify the environment of the test runner, so that the script sees the - // appropriate environment variables. This is necessary because any $CLI invocations in the "script" - // will be executed in the same process as the test runner. - processEnv := auth.ProcessEnv(cfg) - if inprocessMode { - testutil.NullEnvironment(t) - for _, kv := range processEnv { - parts := strings.SplitN(kv, "=", 2) - require.Len(t, parts, 2) - t.Setenv(parts[0], parts[1]) - } - } - // Must be added PrepareReplacementsUser, otherwise conflicts with [USERNAME] testdiff.PrepareReplacementsUUID(t, &repls) From 6932a8171b02ff2f34f9f8bb557ab4b95835ad5a Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 12:13:29 +0200 Subject: [PATCH 47/57] - --- acceptance/acceptance_test.go | 1 - 1 file changed, 1 deletion(-) diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 4b815720c1..55b189d669 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -403,7 +403,6 @@ func runTest(t *testing.T, } ctx, cancelFunc := context.WithTimeout(context.Background(), timeout) defer cancelFunc() - args := []string{"bash", "-euo", "pipefail", EntryPointScript} cmd := exec.CommandContext(ctx, args[0], args[1:]...) From a52d28eaa6c554cb024280b2d0420ff8498b57e1 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 12:14:25 +0200 Subject: [PATCH 48/57] - --- acceptance/internal/prepare_server.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 644b689c49..58cd449a8d 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -40,8 +40,6 @@ func isTruePtr(value *bool) bool { return value != nil && *value } -// This function returns the workspace auth configuration that should be used to authenticate -// the "script" that's executed as part of the acceptance test. func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User) { cloudEnv := os.Getenv("CLOUD_ENV") recordRequests := isTruePtr(config.RecordRequests) From 74cc3a7b44a8cef64af5c9cc6b9e9d42b8ee5c5f Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 12:15:59 +0200 Subject: [PATCH 49/57] - --- libs/testproxy/server.go | 3 --- 1 file changed, 3 deletions(-) diff --git a/libs/testproxy/server.go b/libs/testproxy/server.go index 0ecf899d2f..e303a995d2 100644 --- a/libs/testproxy/server.go +++ b/libs/testproxy/server.go @@ -7,7 +7,6 @@ import ( "errors" "net/http" "net/http/httptest" - "sync" "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/testserver" @@ -23,7 +22,6 @@ type ProxyServer struct { *httptest.Server t testutil.TestingT - mu *sync.Mutex router *mux.Router apiClient *client.DatabricksClient @@ -51,7 +49,6 @@ func New(t testutil.TestingT) *ProxyServer { Server: server, t: t, router: router, - mu: &sync.Mutex{}, } // Create an API client using the current authentication context. From 51a7e265e3f4064292f0e44e3692a16474f9a738 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 12:59:54 +0200 Subject: [PATCH 50/57] - --- libs/testproxy/server.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/testproxy/server.go b/libs/testproxy/server.go index e303a995d2..9d00eb8ccd 100644 --- a/libs/testproxy/server.go +++ b/libs/testproxy/server.go @@ -33,8 +33,8 @@ type ProxyServer struct { // workspace. This is useful for recording API requests and responses in // integration tests. // -// Note: We cannot simply proxy the request from a localhost URL to a real -// workspace. This is because auth resolution in the Databricks SDK relies +// Note: We cannot directly proxy the request from a localhost URL to a real +// workspace as is. This is because auth resolution in the Databricks SDK relies // what the URL actually looks like to determine the auth method to use. // For example, in OAuth flows, the SDK can make requests to different Microsoft // OAuth endpoints based on the nature of the URL. From 5a38af7034b7fe351f8fc650c5d388f746313ad4 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 13:22:09 +0200 Subject: [PATCH 51/57] - --- libs/testserver/server.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index f8d686f2b6..6bba7a5865 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -19,7 +19,7 @@ import ( type Server struct { *httptest.Server - router *mux.Router + Router *mux.Router t testutil.TestingT @@ -183,13 +183,13 @@ func New(t testutil.TestingT) *Server { s := &Server{ Server: server, - router: router, + Router: router, t: t, fakeWorkspaces: map[string]*FakeWorkspace{}, } // Set up the not found handler as fallback. - s.router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + s.Router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { pattern := r.Method + " " + r.URL.Path bodyBytes, err := io.ReadAll(r.Body) var body string @@ -248,7 +248,7 @@ func (s *Server) getWorkspaceForToken(token string) *FakeWorkspace { type HandlerFunc func(req Request) any func (s *Server) Handle(method, path string, handler HandlerFunc) { - s.router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + s.Router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { // Each test uses unique DATABRICKS_TOKEN, we simulate each token having // it's own fake fakeWorkspace to avoid interference between tests. fakeWorkspace := s.getWorkspaceForToken(getToken(r)) From 9111f6fb1b5da3aace8fbdc7e31fd81b49fc0de8 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 13:22:57 +0200 Subject: [PATCH 52/57] - --- libs/testserver/server.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index 6bba7a5865..0a0a3b0909 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -44,6 +44,7 @@ type Response struct { Headers http.Header Body any } + type EncodedResponse struct { StatusCode int Headers http.Header @@ -189,7 +190,7 @@ func New(t testutil.TestingT) *Server { } // Set up the not found handler as fallback. - s.Router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { pattern := r.Method + " " + r.URL.Path bodyBytes, err := io.ReadAll(r.Body) var body string From 027d51220ddcc25bd98e29d8abb7ede9c29ef09d Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 13:23:19 +0200 Subject: [PATCH 53/57] - --- libs/testserver/server.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index 0a0a3b0909..f118e59911 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -189,7 +189,7 @@ func New(t testutil.TestingT) *Server { fakeWorkspaces: map[string]*FakeWorkspace{}, } - // Set up the not found handler as fallback. + // Set up the not found handler as fallback router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { pattern := r.Method + " " + r.URL.Path bodyBytes, err := io.ReadAll(r.Body) From f9ecc8326f8737239b9959fc15982e060c9193e6 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 13:33:39 +0200 Subject: [PATCH 54/57] use http server --- libs/testproxy/server.go | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/libs/testproxy/server.go b/libs/testproxy/server.go index 9d00eb8ccd..38db1abbc7 100644 --- a/libs/testproxy/server.go +++ b/libs/testproxy/server.go @@ -13,7 +13,6 @@ import ( "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/config" - "github.com/gorilla/mux" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -21,8 +20,7 @@ import ( type ProxyServer struct { *httptest.Server - t testutil.TestingT - router *mux.Router + t testutil.TestingT apiClient *client.DatabricksClient RequestCallback func(request *testserver.Request) @@ -41,14 +39,8 @@ type ProxyServer struct { // For reference, see: // https://github.com/databricks/databricks-sdk-go/blob/79e4b3a6e9b0b7dcb1af9ad4025deb447b01d933/common/environment/environments.go#L57 func New(t testutil.TestingT) *ProxyServer { - router := mux.NewRouter() - server := httptest.NewServer(router) - t.Cleanup(server.Close) - s := &ProxyServer{ - Server: server, - t: t, - router: router, + t: t, } // Create an API client using the current authentication context. @@ -59,7 +51,10 @@ func New(t testutil.TestingT) *ProxyServer { require.NoError(t, err) // Set up the proxy handler as the default handler for all requests. - router.NotFoundHandler = http.HandlerFunc(s.proxyToCloud) + server := httptest.NewServer(http.HandlerFunc(s.proxyToCloud)) + t.Cleanup(server.Close) + + s.Server = server return s } From bcfa2d24302188b092b0abab2dc622ee1742f7cc Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 13:45:22 +0200 Subject: [PATCH 55/57] - --- libs/testproxy/server.go | 9 --------- 1 file changed, 9 deletions(-) diff --git a/libs/testproxy/server.go b/libs/testproxy/server.go index 38db1abbc7..f189c26518 100644 --- a/libs/testproxy/server.go +++ b/libs/testproxy/server.go @@ -69,15 +69,6 @@ func (s *ProxyServer) reqBody(r testserver.Request) any { return queryParams } - // The SDK does not support directly passing a JSON serialized request - // body. So we convert it to a map if the body is a JSON object. - if json.Valid(r.Body) { - m := make(map[string]any) - err := json.Unmarshal(r.Body, &m) - assert.NoError(s.t, err) - return m - } - // Otherwise, return the raw body. return r.Body } From 9a5e10e5096fb233212755250766ce145358d8c0 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 16:11:54 +0200 Subject: [PATCH 56/57] address comments --- libs/testproxy/server.go | 50 ++++++++++++++++------------------------ 1 file changed, 20 insertions(+), 30 deletions(-) diff --git a/libs/testproxy/server.go b/libs/testproxy/server.go index f189c26518..b5dfbf2f10 100644 --- a/libs/testproxy/server.go +++ b/libs/testproxy/server.go @@ -104,6 +104,8 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { respBody := bytes.Buffer{} err := s.apiClient.Do(context.Background(), r.Method, r.URL.Path, headers, queryParams, reqBody, &respBody) + var encodedResponse *testserver.EncodedResponse + // API errors from the SDK are expected to be of the type [apierr.APIError]. If we // get an API error then parse the error and forward it back to the client // in an appropriate format. @@ -117,47 +119,35 @@ func (s *ProxyServer) proxyToCloud(w http.ResponseWriter, r *http.Request) { b, err := json.Marshal(body) assert.NoError(s.t, err) - w.WriteHeader(apiErr.StatusCode) - _, err = w.Write(b) - assert.NoError(s.t, err) - - if s.ResponseCallback != nil { - s.ResponseCallback(&request, &testserver.EncodedResponse{ - StatusCode: apiErr.StatusCode, - Body: []byte(apiErr.Message), - }) + encodedResponse = &testserver.EncodedResponse{ + StatusCode: apiErr.StatusCode, + Body: b, } - - return } // Something else went wrong. - if err != nil { - w.WriteHeader(http.StatusInternalServerError) - _, err := w.Write([]byte(err.Error())) - assert.NoError(s.t, err) - - if s.ResponseCallback != nil { - s.ResponseCallback(&request, &testserver.EncodedResponse{ - StatusCode: 500, - Body: []byte(err.Error()), - }) + if encodedResponse == nil && err != nil { + encodedResponse = &testserver.EncodedResponse{ + StatusCode: 500, + Body: []byte(err.Error()), } - - return } // Successful response - w.WriteHeader(200) - b := respBody.Bytes() + if encodedResponse == nil { + encodedResponse = &testserver.EncodedResponse{ + StatusCode: 200, + Body: respBody.Bytes(), + } + } - _, err = w.Write(b) + // Send response to client. + w.WriteHeader(encodedResponse.StatusCode) + + _, err = w.Write(encodedResponse.Body) assert.NoError(s.t, err) if s.ResponseCallback != nil { - s.ResponseCallback(&request, &testserver.EncodedResponse{ - StatusCode: 200, - Body: b, - }) + s.ResponseCallback(&request, encodedResponse) } } From 084a5c3c2ec2adf293447d42e1736ee04a421cc2 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 7 May 2025 16:23:19 +0200 Subject: [PATCH 57/57] print requests inline --- .../pipeline-crud/out.requests.txt | 46 -------- .../record_cloud/pipeline-crud/output.txt | 58 ++++++++++ .../record_cloud/pipeline-crud/script | 17 +++ .../record_cloud/volume-io/out.requests.txt | 79 ------------- .../record_cloud/volume-io/output.txt | 105 ++++++++++++++++++ .../selftest/record_cloud/volume-io/script | 31 ++++++ .../workspace-file-io/out.requests.txt | 58 ---------- .../record_cloud/workspace-file-io/output.txt | 76 +++++++++++++ .../record_cloud/workspace-file-io/script | 23 ++++ 9 files changed, 310 insertions(+), 183 deletions(-) delete mode 100644 acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt delete mode 100644 acceptance/selftest/record_cloud/volume-io/out.requests.txt delete mode 100644 acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt diff --git a/acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt b/acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt deleted file mode 100644 index 5d69eb6fd1..0000000000 --- a/acceptance/selftest/record_cloud/pipeline-crud/out.requests.txt +++ /dev/null @@ -1,46 +0,0 @@ -{ - "method": "POST", - "path": "/api/2.0/pipelines", - "body": { - "allow_duplicate_names": true, - "libraries": [ - { - "file": { - "path": "/whatever.py" - } - } - ], - "name": "test-pipeline-1" - } -} -{ - "method": "GET", - "path": "/api/2.0/pipelines/[UUID]" -} -{ - "method": "PUT", - "path": "/api/2.0/pipelines/[UUID]", - "body": { - "allow_duplicate_names": true, - "libraries": [ - { - "file": { - "path": "/whatever.py" - } - } - ], - "name": "test-pipeline-2" - } -} -{ - "method": "GET", - "path": "/api/2.0/pipelines/[UUID]" -} -{ - "method": "DELETE", - "path": "/api/2.0/pipelines/[UUID]" -} -{ - "method": "GET", - "path": "/api/2.0/pipelines/[UUID]" -} diff --git a/acceptance/selftest/record_cloud/pipeline-crud/output.txt b/acceptance/selftest/record_cloud/pipeline-crud/output.txt index cceca41f90..2aef802cd9 100644 --- a/acceptance/selftest/record_cloud/pipeline-crud/output.txt +++ b/acceptance/selftest/record_cloud/pipeline-crud/output.txt @@ -1,22 +1,80 @@ === Create a pipeline +>>> print_requests +{ + "method": "POST", + "path": "/api/2.0/pipelines", + "body": { + "allow_duplicate_names": true, + "libraries": [ + { + "file": { + "path": "/whatever.py" + } + } + ], + "name": "test-pipeline-1" + } +} + === Get the pipeline >>> [CLI] pipelines get [UUID] "test-pipeline-1" +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/pipelines/[UUID]" +} + === Update the pipeline >>> [CLI] pipelines update [UUID] --json @pipeline2.json +>>> print_requests +{ + "method": "PUT", + "path": "/api/2.0/pipelines/[UUID]", + "body": { + "allow_duplicate_names": true, + "libraries": [ + { + "file": { + "path": "/whatever.py" + } + } + ], + "name": "test-pipeline-2" + } +} + === Verify the update >>> [CLI] pipelines get [UUID] "test-pipeline-2" +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/pipelines/[UUID]" +} + === Delete the pipeline >>> [CLI] pipelines delete [UUID] +>>> print_requests +{ + "method": "DELETE", + "path": "/api/2.0/pipelines/[UUID]" +} + === Verify the deletion >>> [CLI] pipelines get [UUID] Error: The specified pipeline [UUID] was not found. Exit code: 1 + +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/pipelines/[UUID]" +} diff --git a/acceptance/selftest/record_cloud/pipeline-crud/script b/acceptance/selftest/record_cloud/pipeline-crud/script index 965ffe0e4d..85de4aa322 100644 --- a/acceptance/selftest/record_cloud/pipeline-crud/script +++ b/acceptance/selftest/record_cloud/pipeline-crud/script @@ -1,19 +1,36 @@ +print_requests() { + cat out.requests.txt + rm out.requests.txt +} + # Verify that the entire crud lifecycle of a pipeline works. title "Create a pipeline" export pipeline_id=$($CLI pipelines create --json @pipeline1.json | jq -r .pipeline_id) echo "" +trace print_requests + title "Get the pipeline" trace $CLI pipelines get $pipeline_id | jq .name +trace print_requests + title "Update the pipeline" trace $CLI pipelines update $pipeline_id --json @pipeline2.json +trace print_requests + title "Verify the update" trace $CLI pipelines get $pipeline_id | jq .name +trace print_requests + title "Delete the pipeline" trace $CLI pipelines delete $pipeline_id +trace print_requests + title "Verify the deletion" errcode trace $CLI pipelines get $pipeline_id + +trace print_requests diff --git a/acceptance/selftest/record_cloud/volume-io/out.requests.txt b/acceptance/selftest/record_cloud/volume-io/out.requests.txt deleted file mode 100644 index 36425422d5..0000000000 --- a/acceptance/selftest/record_cloud/volume-io/out.requests.txt +++ /dev/null @@ -1,79 +0,0 @@ -{ - "method": "POST", - "path": "/api/2.1/unity-catalog/schemas", - "body": { - "catalog_name": "main", - "name": "schema-[UNIQUE_NAME]" - } -} -{ - "method": "GET", - "path": "/api/2.1/unity-catalog/schemas/main.schema-[UNIQUE_NAME]" -} -{ - "method": "POST", - "path": "/api/2.1/unity-catalog/volumes", - "body": { - "catalog_name": "main", - "name": "volume-[UNIQUE_NAME]", - "schema_name": "schema-[UNIQUE_NAME]", - "volume_type": "MANAGED" - } -} -{ - "method": "GET", - "path": "/api/2.1/unity-catalog/volumes/main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]" -} -{ - "method": "HEAD", - "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" -} -{ - "method": "HEAD", - "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" -} -{ - "method": "PUT", - "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt", - "raw_body": "hello, world" -} -{ - "method": "GET", - "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" -} -{ - "method": "GET", - "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" -} -{ - "method": "HEAD", - "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" -} -{ - "method": "HEAD", - "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" -} -{ - "method": "DELETE", - "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" -} -{ - "method": "GET", - "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" -} -{ - "method": "DELETE", - "path": "/api/2.1/unity-catalog/volumes/main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]" -} -{ - "method": "GET", - "path": "/api/2.1/unity-catalog/volumes/main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]" -} -{ - "method": "DELETE", - "path": "/api/2.1/unity-catalog/schemas/main.schema-[UNIQUE_NAME]" -} -{ - "method": "GET", - "path": "/api/2.1/unity-catalog/schemas/main.schema-[UNIQUE_NAME]" -} diff --git a/acceptance/selftest/record_cloud/volume-io/output.txt b/acceptance/selftest/record_cloud/volume-io/output.txt index 3ab411a8a4..8407516403 100644 --- a/acceptance/selftest/record_cloud/volume-io/output.txt +++ b/acceptance/selftest/record_cloud/volume-io/output.txt @@ -5,46 +5,151 @@ "owner": "[USERNAME]" } +>>> print_requests +{ + "method": "POST", + "path": "/api/2.1/unity-catalog/schemas", + "body": { + "catalog_name": "main", + "name": "schema-[UNIQUE_NAME]" + } +} + >>> [CLI] schemas get main.schema-[UNIQUE_NAME] { "full_name": "main.schema-[UNIQUE_NAME]", "owner": "[USERNAME]" } +>>> print_requests +{ + "method": "GET", + "path": "/api/2.1/unity-catalog/schemas/main.schema-[UNIQUE_NAME]" +} + >>> [CLI] volumes create main schema-[UNIQUE_NAME] volume-[UNIQUE_NAME] MANAGED { "full_name": "main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]", "owner": "[USERNAME]" } +>>> print_requests +{ + "method": "POST", + "path": "/api/2.1/unity-catalog/volumes", + "body": { + "catalog_name": "main", + "name": "volume-[UNIQUE_NAME]", + "schema_name": "schema-[UNIQUE_NAME]", + "volume_type": "MANAGED" + } +} + >>> [CLI] volumes read main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME] { "full_name": "main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]", "owner": "[USERNAME]" } +>>> print_requests +{ + "method": "GET", + "path": "/api/2.1/unity-catalog/volumes/main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]" +} + >>> [CLI] fs cp ./hello.txt dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME] ./hello.txt -> dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt +>>> print_requests +{ + "method": "HEAD", + "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" +} +{ + "method": "HEAD", + "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" +} +{ + "method": "PUT", + "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt", + "raw_body": "hello, world" +} + >>> [CLI] fs ls dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME] hello.txt +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" +} + >>> [CLI] fs cat dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt hello, world +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" +} + >>> [CLI] fs rm dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt +>>> print_requests +{ + "method": "HEAD", + "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" +} +{ + "method": "HEAD", + "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" +} +{ + "method": "DELETE", + "path": "/api/2.0/fs/files/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]/hello.txt" +} + >>> [CLI] fs ls dbfs:/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME] +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/fs/directories/Volumes/main/schema-[UNIQUE_NAME]/volume-[UNIQUE_NAME]" +} + >>> [CLI] volumes delete main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME] +>>> print_requests +{ + "method": "DELETE", + "path": "/api/2.1/unity-catalog/volumes/main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]" +} + >>> [CLI] volumes read main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME] Error: Volume 'main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]' does not exist. Exit code: 1 +>>> print_requests +{ + "method": "GET", + "path": "/api/2.1/unity-catalog/volumes/main.schema-[UNIQUE_NAME].volume-[UNIQUE_NAME]" +} + >>> [CLI] schemas delete main.schema-[UNIQUE_NAME] +>>> print_requests +{ + "method": "DELETE", + "path": "/api/2.1/unity-catalog/schemas/main.schema-[UNIQUE_NAME]" +} + >>> [CLI] schemas get main.schema-[UNIQUE_NAME] Error: Schema 'main.schema-[UNIQUE_NAME]' does not exist. Exit code: 1 + +>>> print_requests +{ + "method": "GET", + "path": "/api/2.1/unity-catalog/schemas/main.schema-[UNIQUE_NAME]" +} diff --git a/acceptance/selftest/record_cloud/volume-io/script b/acceptance/selftest/record_cloud/volume-io/script index 04e297d6b9..741eec1957 100644 --- a/acceptance/selftest/record_cloud/volume-io/script +++ b/acceptance/selftest/record_cloud/volume-io/script @@ -1,25 +1,56 @@ +print_requests() { + cat out.requests.txt + rm out.requests.txt +} + trace $CLI schemas create schema-$UNIQUE_NAME main | jq '{full_name, owner}' +trace print_requests + trace $CLI schemas get main.schema-$UNIQUE_NAME | jq '{full_name, owner}' +trace print_requests + trace $CLI volumes create main schema-$UNIQUE_NAME volume-$UNIQUE_NAME MANAGED | jq '{full_name, owner}' +trace print_requests + trace $CLI volumes read main.schema-$UNIQUE_NAME.volume-$UNIQUE_NAME | jq '{full_name, owner}' +trace print_requests + trace $CLI fs cp ./hello.txt dbfs:/Volumes/main/schema-$UNIQUE_NAME/volume-$UNIQUE_NAME +trace print_requests + trace $CLI fs ls dbfs:/Volumes/main/schema-$UNIQUE_NAME/volume-$UNIQUE_NAME +trace print_requests + trace $CLI fs cat dbfs:/Volumes/main/schema-$UNIQUE_NAME/volume-$UNIQUE_NAME/hello.txt +trace print_requests + trace $CLI fs rm dbfs:/Volumes/main/schema-$UNIQUE_NAME/volume-$UNIQUE_NAME/hello.txt +trace print_requests + trace $CLI fs ls dbfs:/Volumes/main/schema-$UNIQUE_NAME/volume-$UNIQUE_NAME +trace print_requests + trace $CLI volumes delete main.schema-$UNIQUE_NAME.volume-$UNIQUE_NAME +trace print_requests + errcode trace $CLI volumes read main.schema-$UNIQUE_NAME.volume-$UNIQUE_NAME +trace print_requests + trace $CLI schemas delete main.schema-$UNIQUE_NAME +trace print_requests + errcode trace $CLI schemas get main.schema-$UNIQUE_NAME + +trace print_requests diff --git a/acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt b/acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt deleted file mode 100644 index 2679e24262..0000000000 --- a/acceptance/selftest/record_cloud/workspace-file-io/out.requests.txt +++ /dev/null @@ -1,58 +0,0 @@ -{ - "method": "GET", - "path": "/api/2.0/preview/scim/v2/Me" -} -{ - "method": "POST", - "path": "/api/2.0/workspace/mkdirs", - "body": { - "path": "/Users/[USERNAME]/[UNIQUE_NAME]" - } -} -{ - "method": "POST", - "path": "/api/2.0/workspace/import", - "body": { - "content": "[HELLO-WORLD]", - "format": "AUTO", - "path": "/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt" - } -} -{ - "method": "GET", - "path": "/api/2.0/workspace/list" -} -{ - "method": "GET", - "path": "/api/2.0/workspace/get-status" -} -{ - "method": "GET", - "path": "/api/2.0/workspace/export" -} -{ - "method": "POST", - "path": "/api/2.0/workspace/delete", - "body": { - "path": "/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt" - } -} -{ - "method": "GET", - "path": "/api/2.0/workspace/get-status" -} -{ - "method": "GET", - "path": "/api/2.0/workspace/list" -} -{ - "method": "POST", - "path": "/api/2.0/workspace/delete", - "body": { - "path": "/Users/[USERNAME]/[UNIQUE_NAME]" - } -} -{ - "method": "GET", - "path": "/api/2.0/workspace/list" -} diff --git a/acceptance/selftest/record_cloud/workspace-file-io/output.txt b/acceptance/selftest/record_cloud/workspace-file-io/output.txt index 1fb61c716b..39b5a2cf8d 100644 --- a/acceptance/selftest/record_cloud/workspace-file-io/output.txt +++ b/acceptance/selftest/record_cloud/workspace-file-io/output.txt @@ -2,9 +2,33 @@ === create a folder >>> [CLI] workspace mkdirs /Users/[USERNAME]/[UNIQUE_NAME] +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/preview/scim/v2/Me" +} +{ + "method": "POST", + "path": "/api/2.0/workspace/mkdirs", + "body": { + "path": "/Users/[USERNAME]/[UNIQUE_NAME]" + } +} + === upload a file >>> [CLI] workspace import /Users/[USERNAME]/[UNIQUE_NAME]/hello.txt --format AUTO --file ./hello.txt +>>> print_requests +{ + "method": "POST", + "path": "/api/2.0/workspace/import", + "body": { + "content": "[HELLO-WORLD]", + "format": "AUTO", + "path": "/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt" + } +} + === list the folder >>> [CLI] workspace list /Users/[USERNAME]/[UNIQUE_NAME] --output json { @@ -12,6 +36,12 @@ "object_type": "FILE" } +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/workspace/list" +} + === stat the file >>> [CLI] workspace get-status /Users/[USERNAME]/[UNIQUE_NAME]/hello.txt { @@ -19,21 +49,48 @@ "object_type": "FILE" } +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/workspace/get-status" +} + === download the file >>> [CLI] workspace export /Users/[USERNAME]/[UNIQUE_NAME]/hello.txt --format AUTO --file ./hello2.txt >>> cat hello2.txt hello, world +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/workspace/export" +} + === delete the file >>> [CLI] workspace delete /Users/[USERNAME]/[UNIQUE_NAME]/hello.txt +>>> print_requests +{ + "method": "POST", + "path": "/api/2.0/workspace/delete", + "body": { + "path": "/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt" + } +} + === stat the file again >>> [CLI] workspace get-status /Users/[USERNAME]/[UNIQUE_NAME]/hello.txt Error: Path (/Users/[USERNAME]/[UNIQUE_NAME]/hello.txt) doesn't exist. Exit code: 1 +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/workspace/get-status" +} + === list the folder >>> [CLI] workspace list /Users/[USERNAME]/[UNIQUE_NAME] ID Type Language Path @@ -41,8 +98,27 @@ ID Type Language Path === delete the folder >>> [CLI] workspace delete /Users/[USERNAME]/[UNIQUE_NAME] +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/workspace/list" +} +{ + "method": "POST", + "path": "/api/2.0/workspace/delete", + "body": { + "path": "/Users/[USERNAME]/[UNIQUE_NAME]" + } +} + === list the folder again >>> [CLI] workspace list /Users/[USERNAME]/[UNIQUE_NAME] Error: Path (/Users/[USERNAME]/[UNIQUE_NAME]) doesn't exist. Exit code: 1 + +>>> print_requests +{ + "method": "GET", + "path": "/api/2.0/workspace/list" +} diff --git a/acceptance/selftest/record_cloud/workspace-file-io/script b/acceptance/selftest/record_cloud/workspace-file-io/script index 911b6a6df2..990b785691 100644 --- a/acceptance/selftest/record_cloud/workspace-file-io/script +++ b/acceptance/selftest/record_cloud/workspace-file-io/script @@ -5,35 +5,58 @@ export username=$($CLI current-user me | jq -r .userName) # Setting this environment variable prevents that conversion on windows. export MSYS_NO_PATHCONV=1 +print_requests() { + cat out.requests.txt + rm out.requests.txt +} + title "create a folder" trace $CLI workspace mkdirs /Users/$username/$UNIQUE_NAME +trace print_requests + title "upload a file" trace $CLI workspace import /Users/$username/$UNIQUE_NAME/hello.txt --format AUTO --file ./hello.txt +trace print_requests + title "list the folder" trace $CLI workspace list /Users/$username/$UNIQUE_NAME --output json | jq '.[] | {path, object_type}' +trace print_requests + title "stat the file" trace $CLI workspace get-status /Users/$username/$UNIQUE_NAME/hello.txt | jq '{path, object_type}' +trace print_requests + title "download the file" trace $CLI workspace export /Users/$username/$UNIQUE_NAME/hello.txt --format AUTO --file ./hello2.txt trace cat hello2.txt rm hello2.txt +trace print_requests + title "delete the file" trace $CLI workspace delete /Users/$username/$UNIQUE_NAME/hello.txt +trace print_requests + title "stat the file again" errcode trace $CLI workspace get-status /Users/$username/$UNIQUE_NAME/hello.txt +trace print_requests + title "list the folder" trace $CLI workspace list /Users/$username/$UNIQUE_NAME title "delete the folder" trace $CLI workspace delete /Users/$username/$UNIQUE_NAME +trace print_requests + title "list the folder again" errcode trace $CLI workspace list /Users/$username/$UNIQUE_NAME + +trace print_requests