|
| 1 | +package acceptance_test |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "errors" |
| 6 | + "os" |
| 7 | + "os/exec" |
| 8 | + "path" |
| 9 | + "path/filepath" |
| 10 | + "slices" |
| 11 | + "strings" |
| 12 | + "sync" |
| 13 | + "testing" |
| 14 | + "time" |
| 15 | + |
| 16 | + "github.com/databricks/databricks-sdk-go" |
| 17 | + "github.com/databricks/databricks-sdk-go/apierr" |
| 18 | + "github.com/databricks/databricks-sdk-go/service/iam" |
| 19 | + "github.com/databricks/databricks-sdk-go/service/workspace" |
| 20 | +) |
| 21 | + |
| 22 | +// setupBundleCleanup arranges for every bundle deployed by this run to be |
| 23 | +// destroyed once the suite finishes. The caller invokes it only on cloud: all |
| 24 | +// cloud tests share one real workspace, whereas local tests each get a |
| 25 | +// throwaway in-memory fake workspace with nothing to clean up. |
| 26 | +// |
| 27 | +// prefix is the leg-specific "ci<runID>x<legSuffix>" that ciUniqueName stamps |
| 28 | +// into every $UNIQUE_NAME, so the cleanup sweeps exactly the deployments this |
| 29 | +// leg created and nothing else. That is what makes destroying against the shared |
| 30 | +// workspace safe even while sibling matrix legs (which share the run id and may |
| 31 | +// share the workspace) deploy concurrently. |
| 32 | +func setupBundleCleanup(t *testing.T, execPath, prefix string) { |
| 33 | + // t.Context() is canceled once the test finishes, before cleanups run, so |
| 34 | + // derive a context that survives cancellation for the cleanup's API calls. |
| 35 | + ctx := context.WithoutCancel(t.Context()) |
| 36 | + t.Cleanup(func() { |
| 37 | + cleanBundles(ctx, t, execPath, prefix) |
| 38 | + }) |
| 39 | +} |
| 40 | + |
| 41 | +// cleanBundles finds every bundle this run deployed under the current user's |
| 42 | +// ~/.bundle directory (identified by the run's prefix) and destroys each one, |
| 43 | +// logging each deployment and the total time taken. |
| 44 | +func cleanBundles(ctx context.Context, t *testing.T, execPath, prefix string) { |
| 45 | + start := time.Now() |
| 46 | + |
| 47 | + // Cleanup never fails the test (see the WARNING note below), so on any error |
| 48 | + // that prevents sweeping, log loudly and return rather than require-failing. |
| 49 | + w, err := databricks.NewWorkspaceClient() |
| 50 | + if err != nil { |
| 51 | + t.Logf("WARNING: bundle cleanup skipped, cannot create client: %s", err) |
| 52 | + return |
| 53 | + } |
| 54 | + |
| 55 | + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) |
| 56 | + if err != nil { |
| 57 | + t.Logf("WARNING: bundle cleanup skipped, cannot resolve current user: %s", err) |
| 58 | + return |
| 59 | + } |
| 60 | + |
| 61 | + // Tests deploy under the user's home .bundle by default, but some set |
| 62 | + // workspace.root_path under /Shared (e.g. resources/jobs/shared-root-path), |
| 63 | + // so sweep both. A "/Shared/..." root_path is normalized to "/Workspace/Shared/..." |
| 64 | + // by prependWorkspacePrefix, so the swept path uses that form. |
| 65 | + bundleRoots := []string{ |
| 66 | + "/Workspace/Users/" + me.UserName + "/.bundle", |
| 67 | + "/Workspace/Shared/" + me.UserName + "/.bundle", |
| 68 | + } |
| 69 | + |
| 70 | + // The run's prefix always appears in the first path segment under .bundle |
| 71 | + // (the bundle name, or the leaf of a workspace.root_path override), so match |
| 72 | + // there and only descend into this run's subtrees. This avoids walking the |
| 73 | + // thousands of directories other runs may have leaked under .bundle. |
| 74 | + var roots []string |
| 75 | + for _, bundleRoot := range bundleRoots { |
| 76 | + for _, child := range listChildDirs(ctx, t, w, bundleRoot) { |
| 77 | + if strings.Contains(path.Base(child), prefix) { |
| 78 | + roots = append(roots, findDeploymentRoots(ctx, t, w, child)...) |
| 79 | + } |
| 80 | + } |
| 81 | + } |
| 82 | + slices.Sort(roots) |
| 83 | + |
| 84 | + t.Logf("%s bundle cleanup: found %d deployment(s) with prefix %q", time.Now().Format(time.RFC3339), len(roots), prefix) |
| 85 | + |
| 86 | + // Each destroy shells out to a separate `bundle destroy` (auth + state pull + |
| 87 | + // deletes), so run them concurrently. Each is network-bound (not CPU-bound), |
| 88 | + // so cap at a fixed 20 rather than by GOMAXPROCS: it parallelizes the API |
| 89 | + // waits while staying well under the workspace rate limit. This is |
| 90 | + // best-effort, not fail-fast: a failed destroy is recorded and the rest still |
| 91 | + // run, so a plain WaitGroup with a semaphore fits better than errgroup (whose |
| 92 | + // error short-circuit we would not use). |
| 93 | + const maxConcurrentDestroys = 20 |
| 94 | + sem := make(chan struct{}, maxConcurrentDestroys) |
| 95 | + var wg sync.WaitGroup |
| 96 | + var mu sync.Mutex |
| 97 | + var failed []string |
| 98 | + for _, root := range roots { |
| 99 | + sem <- struct{}{} |
| 100 | + wg.Go(func() { |
| 101 | + defer func() { <-sem }() |
| 102 | + t.Logf("%s destroying %s", time.Now().Format(time.RFC3339), root) |
| 103 | + if out, err := destroyBundle(execPath, root); err != nil { |
| 104 | + t.Logf("%s destroy failed: %s\n%s", time.Now().Format(time.RFC3339), root, out) |
| 105 | + mu.Lock() |
| 106 | + failed = append(failed, root) |
| 107 | + mu.Unlock() |
| 108 | + } |
| 109 | + }) |
| 110 | + } |
| 111 | + wg.Wait() |
| 112 | + |
| 113 | + slices.Sort(failed) |
| 114 | + t.Logf("%s bundle cleanup: destroyed %d/%d deployment(s) in %s", time.Now().Format(time.RFC3339), len(roots)-len(failed), len(roots), time.Since(start)) |
| 115 | + |
| 116 | + // Do not fail the test on a cleanup failure: this runs in a t.Cleanup on the |
| 117 | + // root TestAccept, so failing here marks the root test failed with no failed |
| 118 | + // subtest, which makes gotestsum --rerun-fails (used by the integration task) |
| 119 | + // rerun the entire cloud suite. Cleanup is best-effort housekeeping and the |
| 120 | + // product tests already passed, so log loudly instead; leaked deployments are |
| 121 | + // reclaimed by the periodic prefix sweep (sweep_test_resources.py). |
| 122 | + if len(failed) > 0 { |
| 123 | + t.Logf("WARNING: bundle cleanup failed to destroy %d deployment(s), leaked until swept: %s", len(failed), strings.Join(failed, ", ")) |
| 124 | + } |
| 125 | +} |
| 126 | + |
| 127 | +// findDeploymentRoots walks the workspace tree under dir and returns the paths |
| 128 | +// of bundle deployment roots. A directory is a deployment root when it contains |
| 129 | +// a "state" or "files" child, which the bundle deploy writes beneath the |
| 130 | +// resolved workspace.root_path. This works regardless of whether the root is |
| 131 | +// the default ~/.bundle/<name>/<target> or a custom ~/.bundle/<...> override. |
| 132 | +func findDeploymentRoots(ctx context.Context, t *testing.T, w *databricks.WorkspaceClient, dir string) []string { |
| 133 | + var childDirs []string |
| 134 | + for _, child := range listChildDirs(ctx, t, w, dir) { |
| 135 | + if base := path.Base(child); base == "state" || base == "files" { |
| 136 | + // dir is a deployment root; don't descend into its internals. |
| 137 | + return []string{dir} |
| 138 | + } |
| 139 | + childDirs = append(childDirs, child) |
| 140 | + } |
| 141 | + |
| 142 | + var roots []string |
| 143 | + for _, child := range childDirs { |
| 144 | + roots = append(roots, findDeploymentRoots(ctx, t, w, child)...) |
| 145 | + } |
| 146 | + return roots |
| 147 | +} |
| 148 | + |
| 149 | +// listChildDirs returns the immediate subdirectory paths of dir. A missing dir |
| 150 | +// (nothing was deployed under it) yields nil silently; any other listing error |
| 151 | +// is logged loudly (it means the sweep under dir is incomplete) but does not |
| 152 | +// fail the test, since cleanup runs in the root t.Cleanup. |
| 153 | +func listChildDirs(ctx context.Context, t *testing.T, w *databricks.WorkspaceClient, dir string) []string { |
| 154 | + objects, err := w.Workspace.ListAll(ctx, workspace.ListWorkspaceRequest{Path: dir}) |
| 155 | + if err != nil { |
| 156 | + if !errors.Is(err, apierr.ErrNotFound) { |
| 157 | + t.Logf("WARNING: bundle cleanup incomplete, cannot list %s: %s", dir, err) |
| 158 | + } |
| 159 | + return nil |
| 160 | + } |
| 161 | + var dirs []string |
| 162 | + for _, o := range objects { |
| 163 | + if o.ObjectType == workspace.ObjectTypeDirectory { |
| 164 | + dirs = append(dirs, o.Path) |
| 165 | + } |
| 166 | + } |
| 167 | + return dirs |
| 168 | +} |
| 169 | + |
| 170 | +// destroyBundle destroys the bundle deployed at rootPath using the CLI binary, |
| 171 | +// returning the combined output and any error. It writes a throwaway |
| 172 | +// databricks.yml pinning workspace.root_path to the deployment root; destroy |
| 173 | +// pulls the remote state from there (auto-detecting the engine) and deletes the |
| 174 | +// resources and files. The bundle name and target are placeholders because |
| 175 | +// root_path fully determines the deployment location. --force-lock overrides a |
| 176 | +// stale deployment lock left by a test that was killed mid-deploy; these are |
| 177 | +// known-leaked bundles, so there is no concurrent deployment to conflict with. |
| 178 | +func destroyBundle(execPath, rootPath string) ([]byte, error) { |
| 179 | + dir, err := os.MkdirTemp("", "bundle-clean") //nolint:usetesting // runs in a cleanup, where t.TempDir is already removed |
| 180 | + if err != nil { |
| 181 | + return nil, err |
| 182 | + } |
| 183 | + defer os.RemoveAll(dir) |
| 184 | + |
| 185 | + databricksYML := "bundle:\n name: cleanup\nworkspace:\n root_path: " + rootPath + "\ntargets:\n default: {}\n" |
| 186 | + if err := os.WriteFile(filepath.Join(dir, "databricks.yml"), []byte(databricksYML), 0o644); err != nil { |
| 187 | + return nil, err |
| 188 | + } |
| 189 | + |
| 190 | + cmd := exec.Command(execPath, "bundle", "destroy", "--target", "default", "--auto-approve", "--force-lock") |
| 191 | + cmd.Dir = dir |
| 192 | + cmd.Env = os.Environ() |
| 193 | + |
| 194 | + return cmd.CombinedOutput() |
| 195 | +} |
0 commit comments