From 4c4285d7561fdb69ab0a1a66a98e2be0ceb29845 Mon Sep 17 00:00:00 2001 From: Gleb Kanterov Date: Tue, 28 Jan 2025 22:15:31 +0100 Subject: [PATCH 1/3] Add acceptance tests for bundle deploy --- acceptance/acceptance_test.go | 14 +- .../my_project/databricks.yml | 9 + .../my_project/my_notebook.py | 2 + .../my_project/resources.py | 24 ++ .../deploy/experimental-python/output.txt | 34 +++ .../bundle/deploy/experimental-python/script | 7 + .../python-notebook/my_project/databricks.yml | 13 + .../python-notebook/my_project/my_notebook.py | 2 + .../bundle/deploy/python-notebook/output.txt | 34 +++ .../bundle/deploy/python-notebook/script | 7 + acceptance/cmd_server_test.go | 2 +- acceptance/server_test.go | 161 ++++++------ libs/testserver/fake_workspace.go | 237 ++++++++++++++++++ libs/testserver/server.go | 49 +++- 14 files changed, 513 insertions(+), 82 deletions(-) create mode 100644 acceptance/bundle/deploy/experimental-python/my_project/databricks.yml create mode 100644 acceptance/bundle/deploy/experimental-python/my_project/my_notebook.py create mode 100644 acceptance/bundle/deploy/experimental-python/my_project/resources.py create mode 100644 acceptance/bundle/deploy/experimental-python/output.txt create mode 100644 acceptance/bundle/deploy/experimental-python/script create mode 100644 acceptance/bundle/deploy/python-notebook/my_project/databricks.yml create mode 100644 acceptance/bundle/deploy/python-notebook/my_project/my_notebook.py create mode 100644 acceptance/bundle/deploy/python-notebook/output.txt create mode 100644 acceptance/bundle/deploy/python-notebook/script create mode 100644 libs/testserver/fake_workspace.go diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index f205217ffc..c7affc8170 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -19,6 +19,8 @@ import ( "time" "unicode/utf8" + "github.com/google/uuid" + "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testdiff" @@ -123,7 +125,6 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int { AddHandlers(defaultServer) // Redirect API access to local server: t.Setenv("DATABRICKS_HOST", defaultServer.URL) - t.Setenv("DATABRICKS_TOKEN", "dapi1234") homeDir := t.TempDir() // Do not read user's ~/.databrickscfg @@ -146,10 +147,12 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int { // do it last so that full paths match first: repls.SetPath(buildDir, "[BUILD_DIR]") - workspaceClient, err := databricks.NewWorkspaceClient() + config := databricks.Config{Token: "dbapi1234"} + workspaceClient, err := databricks.NewWorkspaceClient(&config) require.NoError(t, err) user, err := workspaceClient.CurrentUser.Me(ctx) + require.NoError(t, err) require.NotNil(t, user) testdiff.PrepareReplacementsUser(t, &repls, *user) @@ -264,7 +267,7 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont for _, stub := range config.Server { require.NotEmpty(t, stub.Pattern) - server.Handle(stub.Pattern, func(req *http.Request) (any, int) { + server.Handle(stub.Pattern, func(fakeWorkspace *testserver.FakeWorkspace, req *http.Request) (any, int) { statusCode := http.StatusOK if stub.Response.StatusCode != 0 { statusCode = stub.Response.StatusCode @@ -285,6 +288,11 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont cmd.Env = append(cmd.Env, "GOCOVERDIR="+coverDir) } + // Each test should use a new token that will result into a new fake workspace, + // so that test don't interfere with each other. + token := strings.ReplaceAll(uuid.NewString(), "-", "") + cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN=dbapi"+token) + // Write combined output to a file out, err := os.Create(filepath.Join(tmpDir, "output.txt")) require.NoError(t, err) diff --git a/acceptance/bundle/deploy/experimental-python/my_project/databricks.yml b/acceptance/bundle/deploy/experimental-python/my_project/databricks.yml new file mode 100644 index 0000000000..e1baf535e1 --- /dev/null +++ b/acceptance/bundle/deploy/experimental-python/my_project/databricks.yml @@ -0,0 +1,9 @@ +bundle: + name: my_project + +sync: {} # dont need to copy files + +experimental: + python: + resources: + - "resources:load_resources" diff --git a/acceptance/bundle/deploy/experimental-python/my_project/my_notebook.py b/acceptance/bundle/deploy/experimental-python/my_project/my_notebook.py new file mode 100644 index 0000000000..259b8197ad --- /dev/null +++ b/acceptance/bundle/deploy/experimental-python/my_project/my_notebook.py @@ -0,0 +1,2 @@ +# Databricks notebook source +1 + 1 diff --git a/acceptance/bundle/deploy/experimental-python/my_project/resources.py b/acceptance/bundle/deploy/experimental-python/my_project/resources.py new file mode 100644 index 0000000000..0e380398c0 --- /dev/null +++ b/acceptance/bundle/deploy/experimental-python/my_project/resources.py @@ -0,0 +1,24 @@ +from databricks.bundles.core import Bundle, Resources +from databricks.bundles.jobs import Job + + +def load_resources(bundle: Bundle) -> Resources: + resources = Resources() + + my_job = Job.from_dict( + { + "name": "My Job", + "tasks": [ + { + "task_key": "my_notebook", + "notebook_task": { + "notebook_path": "my_notebook.py", + }, + }, + ], + } + ) + + resources.add_job("job1", my_job) + + return resources diff --git a/acceptance/bundle/deploy/experimental-python/output.txt b/acceptance/bundle/deploy/experimental-python/output.txt new file mode 100644 index 0000000000..2eb70066eb --- /dev/null +++ b/acceptance/bundle/deploy/experimental-python/output.txt @@ -0,0 +1,34 @@ + +>>> uv run --quiet --python 3.12 --with databricks-bundles==0.7.0 -- [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/my_project/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> [CLI] jobs list --output json +[ + { + "job_id": 1, + "settings": { + "deployment": { + "kind": "BUNDLE", + "metadata_file_path": "/Workspace/Users/[USERNAME]/.bundle/my_project/default/state/metadata.json" + }, + "edit_mode": "UI_LOCKED", + "format": "MULTI_TASK", + "max_concurrent_runs": 1, + "name": "My Job", + "queue": { + "enabled": true + }, + "tasks": [ + { + "notebook_task": { + "notebook_path": "/Workspace/Users/[USERNAME]/.bundle/my_project/default/files/my_notebook" + }, + "task_key": "my_notebook" + } + ] + } + } +] diff --git a/acceptance/bundle/deploy/experimental-python/script b/acceptance/bundle/deploy/experimental-python/script new file mode 100644 index 0000000000..d660236967 --- /dev/null +++ b/acceptance/bundle/deploy/experimental-python/script @@ -0,0 +1,7 @@ +cd my_project + +trace uv run --quiet --python 3.12 --with databricks-bundles==0.7.0 -- $CLI bundle deploy + +trace $CLI jobs list --output json + +rm -rf .databricks __pycache__ .gitignore diff --git a/acceptance/bundle/deploy/python-notebook/my_project/databricks.yml b/acceptance/bundle/deploy/python-notebook/my_project/databricks.yml new file mode 100644 index 0000000000..cb6071f6df --- /dev/null +++ b/acceptance/bundle/deploy/python-notebook/my_project/databricks.yml @@ -0,0 +1,13 @@ +bundle: + name: my_project + +sync: {} # dont need to copy files + +resources: + jobs: + my_job: + name: My Job + tasks: + - task_key: my_notebook + notebook_task: + notebook_path: my_notebook.py diff --git a/acceptance/bundle/deploy/python-notebook/my_project/my_notebook.py b/acceptance/bundle/deploy/python-notebook/my_project/my_notebook.py new file mode 100644 index 0000000000..259b8197ad --- /dev/null +++ b/acceptance/bundle/deploy/python-notebook/my_project/my_notebook.py @@ -0,0 +1,2 @@ +# Databricks notebook source +1 + 1 diff --git a/acceptance/bundle/deploy/python-notebook/output.txt b/acceptance/bundle/deploy/python-notebook/output.txt new file mode 100644 index 0000000000..232ae556ff --- /dev/null +++ b/acceptance/bundle/deploy/python-notebook/output.txt @@ -0,0 +1,34 @@ + +>>> $CLI bundle deploy +Uploading bundle files to /Workspace/Users/$USERNAME/.bundle/my_project/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> $CLI jobs list --output json +[ + { + "job_id": 1, + "settings": { + "deployment": { + "kind": "BUNDLE", + "metadata_file_path": "/Workspace/Users/$USERNAME/.bundle/my_project/default/state/metadata.json" + }, + "edit_mode": "UI_LOCKED", + "format": "MULTI_TASK", + "max_concurrent_runs": 1, + "name": "My Job", + "queue": { + "enabled": true + }, + "tasks": [ + { + "notebook_task": { + "notebook_path": "/Workspace/Users/$USERNAME/.bundle/my_project/default/files/my_notebook" + }, + "task_key": "my_notebook" + } + ] + } + } +] diff --git a/acceptance/bundle/deploy/python-notebook/script b/acceptance/bundle/deploy/python-notebook/script new file mode 100644 index 0000000000..dd97de73d0 --- /dev/null +++ b/acceptance/bundle/deploy/python-notebook/script @@ -0,0 +1,7 @@ +cd my_project + +trace $CLI bundle deploy + +trace $CLI jobs list --output json + +rm -rf .gitignore .databricks diff --git a/acceptance/cmd_server_test.go b/acceptance/cmd_server_test.go index 04d56c7d42..0166dfe32a 100644 --- a/acceptance/cmd_server_test.go +++ b/acceptance/cmd_server_test.go @@ -15,7 +15,7 @@ import ( func StartCmdServer(t *testing.T) *testserver.Server { server := testserver.New(t) - server.Handle("/", func(r *http.Request) (any, int) { + server.Handle("/", func(w *testserver.FakeWorkspace, r *http.Request) (any, int) { q := r.URL.Query() args := strings.Split(q.Get("args"), " ") diff --git a/acceptance/server_test.go b/acceptance/server_test.go index a7695b21e1..a99818fb1d 100644 --- a/acceptance/server_test.go +++ b/acceptance/server_test.go @@ -1,97 +1,116 @@ package acceptance_test import ( + "bytes" + "encoding/json" + "fmt" "net/http" + "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go/service/catalog" - "github.com/databricks/databricks-sdk-go/service/compute" - "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/workspace" ) func AddHandlers(server *testserver.Server) { - server.Handle("GET /api/2.0/policies/clusters/list", func(r *http.Request) (any, int) { - return compute.ListPoliciesResponse{ - Policies: []compute.Policy{ - { - PolicyId: "5678", - Name: "wrong-cluster-policy", - }, - { - PolicyId: "9876", - Name: "some-test-cluster-policy", - }, - }, - }, http.StatusOK + server.Handle("GET /api/2.0/policies/clusters/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + return fakeWorkspace.PoliciesList() }) - server.Handle("GET /api/2.0/instance-pools/list", func(r *http.Request) (any, int) { - return compute.ListInstancePools{ - InstancePools: []compute.InstancePoolAndStats{ - { - InstancePoolName: "some-test-instance-pool", - InstancePoolId: "1234", - }, - }, - }, http.StatusOK + server.Handle("GET /api/2.0/instance-pools/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + return fakeWorkspace.InstancePoolsList() }) - server.Handle("GET /api/2.1/clusters/list", func(r *http.Request) (any, int) { - return compute.ListClustersResponse{ - Clusters: []compute.ClusterDetails{ - { - ClusterName: "some-test-cluster", - ClusterId: "4321", - }, - { - ClusterName: "some-other-cluster", - ClusterId: "9876", - }, - }, - }, http.StatusOK + server.Handle("GET /api/2.1/clusters/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + return fakeWorkspace.ClustersList() }) - server.Handle("GET /api/2.0/preview/scim/v2/Me", func(r *http.Request) (any, int) { - return iam.User{ - Id: "1000012345", - UserName: "tester@databricks.com", - }, http.StatusOK + server.Handle("GET /api/2.1/unity-catalog/current-metastore-assignment", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + return fakeWorkspace.CurrentMetastoreAssignment() }) - server.Handle("GET /api/2.0/workspace/get-status", func(r *http.Request) (any, int) { - return workspace.ObjectInfo{ - ObjectId: 1001, - ObjectType: "DIRECTORY", - Path: "", - ResourceId: "1001", - }, http.StatusOK + server.Handle("GET /api/2.0/permissions/directories/{objectId}", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + objectId := r.PathValue("objectId") + + return fakeWorkspace.DirectoryPermissions(objectId) + }) + + server.Handle("GET /api/2.0/preview/scim/v2/Me", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + return fakeWorkspace.Me() + }) + + server.Handle("GET /api/2.0/workspace/get-status", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + path := r.URL.Query().Get("path") + + return fakeWorkspace.GetStatus(path) + }) + + server.Handle("POST /api/2.0/workspace/mkdirs", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + request := workspace.Mkdirs{} + decoder := json.NewDecoder(r.Body) + + err := decoder.Decode(&request) + if err != nil { + return internalError(err) + } + + return fakeWorkspace.WorkspaceMkdirs(request) + }) + + server.Handle("GET /api/2.0/workspace/export", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + path := r.URL.Query().Get("path") + + return fakeWorkspace.WorkspaceExport(path) }) - server.Handle("GET /api/2.1/unity-catalog/current-metastore-assignment", func(r *http.Request) (any, int) { - return catalog.MetastoreAssignment{ - DefaultCatalogName: "main", - }, http.StatusOK + server.Handle("POST /api/2.0/workspace/delete", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + path := r.URL.Query().Get("path") + recursiveStr := r.URL.Query().Get("recursive") + var recursive bool + + if recursiveStr == "true" { + recursive = true + } else { + recursive = false + } + + return fakeWorkspace.WorkspaceDelete(path, recursive) }) - server.Handle("GET /api/2.0/permissions/directories/1001", func(r *http.Request) (any, int) { - return workspace.WorkspaceObjectPermissions{ - ObjectId: "1001", - ObjectType: "DIRECTORY", - AccessControlList: []workspace.WorkspaceObjectAccessControlResponse{ - { - UserName: "tester@databricks.com", - AllPermissions: []workspace.WorkspaceObjectPermission{ - { - PermissionLevel: "CAN_MANAGE", - }, - }, - }, - }, - }, http.StatusOK + server.Handle("POST /api/2.0/workspace-files/import-file/{path}", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + path := r.PathValue("path") + + body := new(bytes.Buffer) + _, err := body.ReadFrom(r.Body) + if err != nil { + return internalError(err) + } + + return fakeWorkspace.WorkspaceImportFiles(path, body.Bytes()) }) - server.Handle("POST /api/2.0/workspace/mkdirs", func(r *http.Request) (any, int) { - return "{}", http.StatusOK + server.Handle("POST /api/2.1/jobs/create", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + request := jobs.CreateJob{} + decoder := json.NewDecoder(r.Body) + + err := decoder.Decode(&request) + if err != nil { + return internalError(err) + } + + return fakeWorkspace.JobsCreate(request) }) + + server.Handle("GET /api/2.1/jobs/get", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + jobId := r.URL.Query().Get("job_id") + + return fakeWorkspace.JobsGet(jobId) + }) + + server.Handle("GET /api/2.1/jobs/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) { + return fakeWorkspace.JobsList() + }) +} + +func internalError(err error) (any, int) { + return fmt.Errorf("internal error: %w", err), http.StatusInternalServerError } diff --git a/libs/testserver/fake_workspace.go b/libs/testserver/fake_workspace.go new file mode 100644 index 0000000000..5b31771a35 --- /dev/null +++ b/libs/testserver/fake_workspace.go @@ -0,0 +1,237 @@ +package testserver + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "strconv" + "strings" + + "github.com/databricks/databricks-sdk-go/service/catalog" + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/databricks/databricks-sdk-go/service/workspace" +) + +// FakeWorkspace holds a state of a workspace for acceptance tests. +type FakeWorkspace struct { + directories map[string]bool + files map[string][]byte + // normally, ids are not sequential, but we make them sequential for deterministic diff + nextJobId int64 + jobs map[int64]jobs.Job +} + +func NewFakeWorkspace() *FakeWorkspace { + return &FakeWorkspace{ + directories: map[string]bool{ + "/Workspace": true, + }, + files: map[string][]byte{}, + jobs: map[int64]jobs.Job{}, + nextJobId: 1, + } +} + +func (s *FakeWorkspace) Me() (iam.User, int) { + return iam.User{ + Id: "1000012345", + UserName: "tester@databricks.com", + }, http.StatusOK +} + +func (s *FakeWorkspace) CurrentMetastoreAssignment() (catalog.MetastoreAssignment, int) { + return catalog.MetastoreAssignment{ + DefaultCatalogName: "main", + }, http.StatusOK +} + +func (s *FakeWorkspace) DirectoryPermissions(objectId string) (workspace.WorkspaceObjectPermissions, int) { + return workspace.WorkspaceObjectPermissions{ + ObjectId: objectId, + ObjectType: "DIRECTORY", + AccessControlList: []workspace.WorkspaceObjectAccessControlResponse{ + { + UserName: "tester@databricks.com", + AllPermissions: []workspace.WorkspaceObjectPermission{ + { + PermissionLevel: "CAN_MANAGE", + }, + }, + }, + }, + }, http.StatusOK +} + +func (s *FakeWorkspace) GetStatus(path string) (workspace.ObjectInfo, int) { + if s.directories[path] { + return workspace.ObjectInfo{ + ObjectType: "DIRECTORY", + Path: path, + }, http.StatusOK + } else if _, ok := s.files[path]; ok { + return workspace.ObjectInfo{ + ObjectType: "FILE", + Path: path, + Language: "SCALA", + }, http.StatusOK + } else { + return workspace.ObjectInfo{}, http.StatusNotFound + } +} + +func (s *FakeWorkspace) WorkspaceMkdirs(request workspace.Mkdirs) (string, int) { + s.directories[request.Path] = true + + return "{}", http.StatusOK +} + +func (s *FakeWorkspace) WorkspaceExport(path string) ([]byte, int) { + file := s.files[path] + + if file == nil { + return nil, http.StatusNotFound + } + + return file, http.StatusOK +} + +func (s *FakeWorkspace) WorkspaceDelete(path string, recursive bool) (string, int) { + if !recursive { + s.files[path] = nil + } else { + for key := range s.files { + if strings.HasPrefix(key, path) { + s.files[key] = nil + } + } + } + + return "{}", http.StatusOK +} + +func (s *FakeWorkspace) WorkspaceImportFiles(path string, body []byte) (any, int) { + if !strings.HasPrefix(path, "/") { + path = "/" + path + } + + s.files[path] = body + + return "{}", http.StatusOK +} + +func (s *FakeWorkspace) JobsCreate(request jobs.CreateJob) (any, int) { + jobId := s.nextJobId + s.nextJobId++ + + jobSettings := jobs.JobSettings{} + err := jsonConvert(request, &jobSettings) + if err != nil { + return internalError(err) + } + + s.jobs[jobId] = jobs.Job{ + JobId: jobId, + Settings: &jobSettings, + } + + return jobs.CreateResponse{JobId: jobId}, http.StatusOK +} + +func (s *FakeWorkspace) JobsGet(jobId string) (any, int) { + id := jobId + + jobIdInt, err := strconv.ParseInt(id, 10, 64) + if err != nil { + return internalError(fmt.Errorf("failed to parse job id: %s", err)) + } + + job, ok := s.jobs[jobIdInt] + if !ok { + return jobs.Job{}, http.StatusNotFound + } + + return job, http.StatusOK +} + +func (s *FakeWorkspace) JobsList() (any, int) { + list := make([]jobs.BaseJob, 0, len(s.jobs)) + for _, job := range s.jobs { + baseJob := jobs.BaseJob{} + err := jsonConvert(job, &baseJob) + if err != nil { + return internalError(fmt.Errorf("failed to convert job to base job: %w", err)) + } + + list = append(list, baseJob) + } + + return jobs.ListJobsResponse{ + Jobs: list, + }, http.StatusOK +} + +func (s *FakeWorkspace) InstancePoolsList() (compute.ListInstancePools, int) { + return compute.ListInstancePools{ + InstancePools: []compute.InstancePoolAndStats{ + { + InstancePoolName: "some-test-instance-pool", + InstancePoolId: "1234", + }, + }, + }, http.StatusOK +} + +func (s *FakeWorkspace) ClustersList() (compute.ListClustersResponse, int) { + return compute.ListClustersResponse{ + Clusters: []compute.ClusterDetails{ + { + ClusterName: "some-test-cluster", + ClusterId: "4321", + }, + { + ClusterName: "some-other-cluster", + ClusterId: "9876", + }, + }, + }, http.StatusOK +} + +func (s *FakeWorkspace) PoliciesList() (any, int) { + return compute.ListPoliciesResponse{ + Policies: []compute.Policy{ + { + PolicyId: "5678", + Name: "wrong-cluster-policy", + }, + { + PolicyId: "9876", + Name: "some-test-cluster-policy", + }, + }, + }, http.StatusOK +} + +// jsonConvert saves input to a value pointed by output +func jsonConvert(input, output any) error { + writer := new(bytes.Buffer) + encoder := json.NewEncoder(writer) + err := encoder.Encode(input) + if err != nil { + return fmt.Errorf("failed to encode: %w", err) + } + + decoder := json.NewDecoder(writer) + err = decoder.Decode(output) + if err != nil { + return fmt.Errorf("failed to decode: %w", err) + } + + return nil +} + +func internalError(err error) (string, int) { + return fmt.Sprintf("internal error: %s", err), http.StatusInternalServerError +} diff --git a/libs/testserver/server.go b/libs/testserver/server.go index 5e3efe1c54..b0f01a24ed 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -6,6 +6,8 @@ import ( "net/http" "net/http/httptest" "slices" + "strings" + "sync" "github.com/stretchr/testify/assert" @@ -18,6 +20,9 @@ type Server struct { t testutil.TestingT + fakeWorkspaces map[string]*FakeWorkspace + mu *sync.Mutex + RecordRequests bool IncludeRequestHeaders []string @@ -37,17 +42,35 @@ func New(t testutil.TestingT) *Server { t.Cleanup(server.Close) return &Server{ - Server: server, - Mux: mux, - t: t, + Server: server, + Mux: mux, + t: t, + mu: &sync.Mutex{}, + fakeWorkspaces: map[string]*FakeWorkspace{}, } } -type HandlerFunc func(req *http.Request) (resp any, statusCode int) +type HandlerFunc func(fakeWorkspace *FakeWorkspace, req *http.Request) (resp any, statusCode int) func (s *Server) Handle(pattern string, handler HandlerFunc) { s.Mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) { - resp, statusCode := handler(r) + // 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 + if token := getToken(r); token != "" { + if _, ok := s.fakeWorkspaces[token]; !ok { + s.fakeWorkspaces[token] = NewFakeWorkspace() + } + + fakeWorkspace = s.fakeWorkspaces[token] + } + + resp, statusCode := handler(fakeWorkspace, r) if s.RecordRequests { body, err := io.ReadAll(r.Body) @@ -75,9 +98,10 @@ func (s *Server) Handle(pattern string, handler HandlerFunc) { var respBytes []byte var err error - respString, ok := resp.(string) - if ok { + if respString, ok := resp.(string); ok { respBytes = []byte(respString) + } else if respBytes0, ok := resp.([]byte); ok { + respBytes = respBytes0 } else { respBytes, err = json.MarshalIndent(resp, "", " ") if err != nil { @@ -92,3 +116,14 @@ func (s *Server) Handle(pattern string, handler HandlerFunc) { } }) } + +func getToken(r *http.Request) string { + header := r.Header.Get("Authorization") + prefix := "Bearer " + + if !strings.HasPrefix(header, prefix) { + return "" + } + + return header[len(prefix):] +} From 6219f55154b6196a16dd88e7b57aeac96e344b2f Mon Sep 17 00:00:00 2001 From: Gleb Kanterov Date: Thu, 6 Feb 2025 09:44:24 +0100 Subject: [PATCH 2/3] Fix tests --- acceptance/acceptance_test.go | 9 ++++++--- .../bundle/deploy/python-notebook/output.txt | 10 +++++----- acceptance/bundle/scripts/output.txt | 14 ++++++-------- acceptance/workspace/jobs/create/out.requests.txt | 2 +- libs/testserver/server.go | 4 +++- 5 files changed, 21 insertions(+), 18 deletions(-) diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index c7affc8170..a3ab4ad22b 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -290,8 +290,10 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont // Each test should use a new token that will result into a new fake workspace, // so that test don't interfere with each other. - token := strings.ReplaceAll(uuid.NewString(), "-", "") - cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN=dbapi"+token) + tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "") + token := "dbapi" + tokenSuffix + cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+token) + repls.Set(token, "[DATABRICKS_TOKEN]") // Write combined output to a file out, err := os.Create(filepath.Join(tmpDir, "output.txt")) @@ -311,7 +313,8 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont reqJson, err := json.Marshal(req) require.NoError(t, err) - line := fmt.Sprintf("%s\n", reqJson) + reqJsonWithRepls := repls.Replace(string(reqJson)) + line := fmt.Sprintf("%s\n", reqJsonWithRepls) _, err = f.WriteString(line) require.NoError(t, err) } diff --git a/acceptance/bundle/deploy/python-notebook/output.txt b/acceptance/bundle/deploy/python-notebook/output.txt index 232ae556ff..2372664d63 100644 --- a/acceptance/bundle/deploy/python-notebook/output.txt +++ b/acceptance/bundle/deploy/python-notebook/output.txt @@ -1,18 +1,18 @@ ->>> $CLI bundle deploy -Uploading bundle files to /Workspace/Users/$USERNAME/.bundle/my_project/default/files... +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/my_project/default/files... Deploying resources... Updating deployment state... Deployment complete! ->>> $CLI jobs list --output json +>>> [CLI] jobs list --output json [ { "job_id": 1, "settings": { "deployment": { "kind": "BUNDLE", - "metadata_file_path": "/Workspace/Users/$USERNAME/.bundle/my_project/default/state/metadata.json" + "metadata_file_path": "/Workspace/Users/[USERNAME]/.bundle/my_project/default/state/metadata.json" }, "edit_mode": "UI_LOCKED", "format": "MULTI_TASK", @@ -24,7 +24,7 @@ Deployment complete! "tasks": [ { "notebook_task": { - "notebook_path": "/Workspace/Users/$USERNAME/.bundle/my_project/default/files/my_notebook" + "notebook_path": "/Workspace/Users/[USERNAME]/.bundle/my_project/default/files/my_notebook" }, "task_key": "my_notebook" } diff --git a/acceptance/bundle/scripts/output.txt b/acceptance/bundle/scripts/output.txt index 2deedb0e75..68afb2fecc 100644 --- a/acceptance/bundle/scripts/output.txt +++ b/acceptance/bundle/scripts/output.txt @@ -42,11 +42,9 @@ from myscript.py 0 postbuild: hello stderr! Executing 'predeploy' script from myscript.py 0 predeploy: hello stdout! from myscript.py 0 predeploy: hello stderr! -Error: unable to deploy to /Workspace/Users/[USERNAME]/.bundle/scripts/default/state as [USERNAME]. -Please make sure the current user or one of their groups is listed under the permissions of this bundle. -For assistance, contact the owners of this project. -They may need to redeploy the bundle to apply the new permissions. -Please refer to https://docs.databricks.com/dev-tools/bundles/permissions.html for more on managing permissions. - - -Exit code: 1 +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/scripts/default/files... +Deploying resources... +Deployment complete! +Executing 'postdeploy' script +from myscript.py 0 postdeploy: hello stdout! +from myscript.py 0 postdeploy: hello stderr! diff --git a/acceptance/workspace/jobs/create/out.requests.txt b/acceptance/workspace/jobs/create/out.requests.txt index 4a85c4c43a..60977e3e31 100644 --- a/acceptance/workspace/jobs/create/out.requests.txt +++ b/acceptance/workspace/jobs/create/out.requests.txt @@ -1 +1 @@ -{"headers":{"Authorization":"Bearer dapi1234","User-Agent":"cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/jobs_create cmd-exec-id/[UUID] auth/pat"},"method":"POST","path":"/api/2.1/jobs/create","body":{"name":"abc"}} +{"headers":{"Authorization":"Bearer [DATABRICKS_TOKEN]","User-Agent":"cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/jobs_create cmd-exec-id/[UUID] auth/pat"},"method":"POST","path":"/api/2.1/jobs/create","body":{"name":"abc"}} diff --git a/libs/testserver/server.go b/libs/testserver/server.go index b0f01a24ed..ad1fc889aa 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -62,7 +62,8 @@ func (s *Server) Handle(pattern string, handler HandlerFunc) { // 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 - if token := getToken(r); token != "" { + token := getToken(r) + if token != "" { if _, ok := s.fakeWorkspaces[token]; !ok { s.fakeWorkspaces[token] = NewFakeWorkspace() } @@ -81,6 +82,7 @@ func (s *Server) Handle(pattern string, handler HandlerFunc) { if len(v) == 0 || !slices.Contains(s.IncludeRequestHeaders, k) { continue } + headers[k] = v[0] } From 062a252f243371b419b802d42b64c5b946b206d3 Mon Sep 17 00:00:00 2001 From: Gleb Kanterov Date: Thu, 6 Feb 2025 09:45:26 +0100 Subject: [PATCH 3/3] Fix tests --- libs/testserver/server.go | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index ad1fc889aa..ffb83a49c2 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -82,7 +82,6 @@ func (s *Server) Handle(pattern string, handler HandlerFunc) { if len(v) == 0 || !slices.Contains(s.IncludeRequestHeaders, k) { continue } - headers[k] = v[0] }