diff --git a/libs/localenv/constraints.go b/libs/localenv/constraints.go new file mode 100644 index 00000000000..22a2c49f66e --- /dev/null +++ b/libs/localenv/constraints.go @@ -0,0 +1,303 @@ +package localenv + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "regexp" + "strings" + "time" + + "github.com/BurntSushi/toml" + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/log" +) + +// EnvConstraintRepo names the environment variable that supplies the GitHub repo +// ("owner/name") hosting the environment constraint artifacts. +const EnvConstraintRepo = "DATABRICKS_LOCALENV_CONSTRAINT_REPO" + +// defaultConstraintRepo is the GitHub repo that hosts the constraint artifacts. +// It is intentionally empty for now: the artifacts must move to a +// Databricks-owned repo (databricks/environments) before this ships, but that +// repo can't publish them yet (its GitHub Actions are disabled). Until then the +// repo is supplied via the EnvConstraintRepo environment variable rather than +// hardcoding a personal repo, so no untrusted default controls what the CLI +// installs. Once databricks/environments is ready this becomes that constant. +const defaultConstraintRepo = "" + +// RepoConstraintBaseURL resolves the base URL for constraint artifacts from the +// hosting GitHub repo: EnvConstraintRepo overrides the (currently empty) built-in +// default, and the repo is turned into a raw.githubusercontent.com main-branch +// URL. It returns "" when no repo is configured; the caller must not fall back to +// any other source, and FetchConstraints reports the missing configuration as a +// fetch-phase error so it surfaces through the normal phase/JSON reporting rather +// than aborting the command before the phase table is rendered. +func RepoConstraintBaseURL(ctx context.Context) string { + repo := defaultConstraintRepo + if v, ok := env.Lookup(ctx, EnvConstraintRepo); ok && strings.TrimSpace(v) != "" { + repo = strings.TrimSpace(v) + } + if repo == "" { + return "" + } + return "https://raw.githubusercontent.com/" + repo + "/main" +} + +// errEnvKeyNotFound is returned by fetchURL when the constraint artifact does +// not exist for the requested env key (HTTP 404). It is distinct from a +// transport failure so FetchConstraints can classify it as E_ENV_UNSUPPORTED +// (a resolvable target with no published environment) rather than E_FETCH. +var errEnvKeyNotFound = errors.New("environment key not found") + +// maxConstraintBytes caps the constraint artifact read. The body is untrusted +// remote content and a pyproject.toml is small; 1 MiB is far above any real +// artifact while preventing a misbehaving or hostile host from exhausting memory. +const maxConstraintBytes int64 = 1 << 20 + +// constraintHTTPClient fetches constraint artifacts with an explicit timeout, so +// the request is bounded even when the caller's context has no deadline. +var constraintHTTPClient = &http.Client{Timeout: 30 * time.Second} + +// Constraints holds the parsed contents of a per-environment pyproject.toml. +type Constraints struct { + // EnvKey is the environment key used to look up the constraints. + EnvKey string + // SourceURL is the URL from which the constraints were fetched. + SourceURL string + // FromCache is true when the data came from the on-disk cache rather than a live fetch. + FromCache bool + // RequiresPython is the PEP 440 python version specifier from [project].requires-python. + RequiresPython string + // DatabricksConnect is the full dependency string for databricks-connect from [dependency-groups].dev. + DatabricksConnect string + // ConstraintDeps is the list of entries from [tool.uv].constraint-dependencies. + ConstraintDeps []string +} + +// cacheFileName maps an env key to a single, collision-free cache filename. +// It keeps a readable slug (path separators flattened to double-underscores so +// the file stays inside cacheDir on every OS) and appends a short hash of the +// raw env key. The hash guarantees injectivity: distinct env keys that would +// otherwise flatten to the same slug (e.g. "a/b" and "a__b") get distinct +// filenames, so a cache hit can never serve another environment's constraints. +func cacheFileName(envKey string) string { + slug := strings.ReplaceAll(envKey, "/", "__") + slug = strings.ReplaceAll(slug, "\\", "__") + sum := sha256.Sum256([]byte(envKey)) + return fmt.Sprintf("%s-%s.toml", slug, hex.EncodeToString(sum[:8])) +} + +// writeCacheAtomic writes data to path via a temp file and rename, creating the +// parent directory first. The rename is atomic on the same filesystem, so a +// concurrent reader never observes a truncated or partial cache file (os.WriteFile +// truncates in place, which a fallback reader could catch mid-write). +func writeCacheAtomic(path string, data []byte) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o755); err != nil { + return err + } + tmp, err := os.CreateTemp(dir, ".constraints-*.tmp") + if err != nil { + return err + } + tmpName := tmp.Name() + if _, err := tmp.Write(data); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + if err := os.Chmod(tmpName, 0o600); err != nil { + os.Remove(tmpName) + return err + } + if err := os.Rename(tmpName, path); err != nil { + os.Remove(tmpName) + return err + } + return nil +} + +// FetchConstraints fetches the pyproject.toml for envKey from baseURL and caches it in +// cacheDir. On a transport or non-404 HTTP failure it falls back to the cached copy if one +// exists (E_FETCH otherwise). A 404 means the env key is not published (E_ENV_UNSUPPORTED) +// and does not fall back to cache — a resolvable target with no environment is a distinct, +// non-transient condition. +// +// baseURL points at the repo hosting the constraint artifacts (see +// RepoConstraintBaseURL); it is empty when no source is configured, which is +// reported below as a fetch-phase error. +func FetchConstraints(ctx context.Context, baseURL, envKey, cacheDir string) (*Constraints, error) { + if baseURL == "" { + // No constraint host is configured. This is reported at the fetch phase (as + // E_FETCH) rather than aborting earlier, so the failure flows through the + // same phase/JSON reporting as any other fetch error. + return nil, NewError(ErrFetch, nil, + "no constraint source configured: set %s to the GitHub repo (owner/name) that hosts the environment constraints", EnvConstraintRepo) + } + url := baseURL + "/" + envKey + "/pyproject.toml" + cachePath := filepath.Join(cacheDir, cacheFileName(envKey)) + + data, fetchErr := fetchURL(ctx, url) + if fetchErr == nil { + // Parse before caching: a malformed 2xx body must not overwrite a valid + // cached copy, or a later transport-failure run would serve the poisoned + // cache and fail to parse instead of falling back to the last-good file. + rp, dbc, deps, err := parseConstraints(data) + if err != nil { + return nil, fmt.Errorf("parse constraints for %s: %w", envKey, err) + } + // Write the cache copy (creating cacheDir if needed, atomically); non-fatal + // so a read-only cacheDir doesn't break the command. + if err := writeCacheAtomic(cachePath, data); err != nil { + log.Debugf(ctx, "failed to write constraint cache %s: %v", filepath.ToSlash(cachePath), err) + } + return &Constraints{ + EnvKey: envKey, + SourceURL: url, + FromCache: false, + RequiresPython: rp, + DatabricksConnect: dbc, + ConstraintDeps: deps, + }, nil + } + + // A missing env key (404) is not a transport failure and has no useful cache + // fallback: the target resolved to an environment that isn't published. + if errors.Is(fetchErr, errEnvKeyNotFound) { + return nil, NewError(ErrEnvUnsupported, fetchErr, + "no published environment for %q. If this is a new runtime, try the latest LTS target (e.g. --serverless v4 or a supported --cluster DBR)", envKey) + } + + // Network or HTTP failure: attempt to serve from cache. + cached, readErr := os.ReadFile(cachePath) + if readErr != nil { + return nil, NewError(ErrFetch, fetchErr, "fetch constraints for %s", envKey) + } + + log.Warnf(ctx, "constraint fetch failed, using cached copy: %v", fetchErr) + rp, dbc, deps, err := parseConstraints(cached) + if err != nil { + return nil, fmt.Errorf("parse cached constraints for %s: %w", envKey, err) + } + return &Constraints{ + EnvKey: envKey, + SourceURL: url, + FromCache: true, + RequiresPython: rp, + DatabricksConnect: dbc, + ConstraintDeps: deps, + }, nil +} + +// fetchURL performs an HTTP GET and returns the body bytes, or an error on non-2xx or transport failure. +func fetchURL(ctx context.Context, url string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, fmt.Errorf("build request for %s: %w", url, err) + } + // Use a client with an explicit timeout rather than http.DefaultClient, which + // has none: the fetch of remote content must be bounded even if the caller's + // context carries no deadline. + resp, err := constraintHTTPClient.Do(req) + if err != nil { + return nil, fmt.Errorf("GET %s: %w", url, err) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusNotFound { + return nil, fmt.Errorf("GET %s: %w", url, errEnvKeyNotFound) + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("GET %s: unexpected status %s", url, resp.Status) + } + // Cap the read: the body is untrusted remote content and a pyproject.toml is + // small, so an oversized (or hostile) response must not be read into memory + // unbounded. Read one byte past the cap to detect an over-limit body. + data, err := io.ReadAll(io.LimitReader(resp.Body, maxConstraintBytes+1)) + if err != nil { + return nil, fmt.Errorf("read body from %s: %w", url, err) + } + if int64(len(data)) > maxConstraintBytes { + return nil, fmt.Errorf("constraint artifact from %s exceeds %d bytes", url, maxConstraintBytes) + } + return data, nil +} + +// pyprojectTOML mirrors the pyproject.toml fields we care about. +type pyprojectTOML struct { + Project struct { + RequiresPython string `toml:"requires-python"` + } `toml:"project"` + DependencyGroups struct { + Dev []string `toml:"dev"` + } `toml:"dependency-groups"` + Tool struct { + UV struct { + ConstraintDependencies []string `toml:"constraint-dependencies"` + } `toml:"uv"` + } `toml:"tool"` +} + +// parseConstraints parses a pyproject.toml byte slice and extracts requires-python, +// the databricks-connect entry from dependency-groups.dev, and constraint-dependencies. +// A body that is valid TOML but carries no requires-python is rejected: it is not a +// usable constraint artifact, and silently accepting it would cache an empty result +// and only surface a confusing failure later in the pipeline. +func parseConstraints(data []byte) (requiresPython, dbconnect string, deps []string, err error) { + var p pyprojectTOML + if err = toml.Unmarshal(data, &p); err != nil { + return "", "", nil, fmt.Errorf("unmarshal pyproject.toml: %w", err) + } + + requiresPython = p.Project.RequiresPython + if strings.TrimSpace(requiresPython) == "" { + return "", "", nil, errors.New("constraint artifact has no [project].requires-python") + } + + for _, entry := range p.DependencyGroups.Dev { + if isDatabricksConnectDep(entry) { + dbconnect = entry + break + } + } + + deps = p.Tool.UV.ConstraintDependencies + return requiresPython, dbconnect, deps, nil +} + +// depNameSepRe matches the first PEP 508 delimiter that ends a requirement's +// package name: a version specifier, extra, marker, url, or list separator. +var depNameSepRe = regexp.MustCompile(`[<>=!~;,@\[( \t]`) + +// isDatabricksConnectDep reports whether a dependency-group entry is the +// databricks-connect requirement. It extracts the leading package name (up to +// the first PEP 508 delimiter) and compares it under PEP 503 normalization, so +// case, and runs of "-", "_", or "." are all treated as equivalent: +// "Databricks-Connect", "databricks_connect", and "databricks.connect" all match, +// while a distinct package like "databricks-connectors" does not. +func isDatabricksConnectDep(entry string) bool { + name := strings.TrimSpace(entry) + if i := depNameSepRe.FindStringIndex(name); i != nil { + name = name[:i[0]] + } + return normalizePackageName(name) == "databricks-connect" +} + +// pep503SepRe matches runs of "-", "_", or "." for PEP 503 name normalization. +var pep503SepRe = regexp.MustCompile(`[-_.]+`) + +// normalizePackageName applies PEP 503 normalization: lowercase and collapse any +// run of "-", "_", or "." to a single "-". +func normalizePackageName(name string) string { + return pep503SepRe.ReplaceAllString(strings.ToLower(strings.TrimSpace(name)), "-") +} diff --git a/libs/localenv/constraints_test.go b/libs/localenv/constraints_test.go new file mode 100644 index 00000000000..8df363d943e --- /dev/null +++ b/libs/localenv/constraints_test.go @@ -0,0 +1,205 @@ +package localenv + +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/databricks/cli/libs/env" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRepoConstraintBaseURL(t *testing.T) { + // With no repo configured (empty built-in default), it returns "" so the caller + // can report the missing source at the fetch phase rather than aborting early. + assert.Empty(t, RepoConstraintBaseURL(t.Context())) + + // The env var supplies the repo and is turned into a raw main-branch URL. + ctx := env.Set(t.Context(), EnvConstraintRepo, "databricks/environments") + assert.Equal(t, "https://raw.githubusercontent.com/databricks/environments/main", RepoConstraintBaseURL(ctx)) + + // Whitespace-only is treated as unset. + ctx = env.Set(t.Context(), EnvConstraintRepo, " ") + assert.Empty(t, RepoConstraintBaseURL(ctx)) +} + +func TestFetchConstraintsNoSourceConfigured(t *testing.T) { + // An empty base URL means no constraint host is configured; it must classify as + // E_FETCH (surfaced at the fetch phase) and name the env var to set. + _, err := FetchConstraints(t.Context(), "", "serverless/serverless-v4", t.TempDir()) + var pe *PipelineError + require.ErrorAs(t, err, &pe) + assert.Equal(t, ErrFetch, pe.Code) + assert.Contains(t, pe.Error(), EnvConstraintRepo) +} + +const sampleToml = `[project] +requires-python = "==3.12.*" + +[dependency-groups] +dev = [ + "databricks-connect~=17.2.0", + "pytest~=8.0", +] + +[tool.uv] +constraint-dependencies = [ + "pydantic~=2.10.6", + "anyio~=4.6.2", +] +` + +func TestParseConstraints(t *testing.T) { + rp, dbc, deps, err := parseConstraints([]byte(sampleToml)) + require.NoError(t, err) + assert.Equal(t, "==3.12.*", rp) + assert.Equal(t, "databricks-connect~=17.2.0", dbc) + assert.Equal(t, []string{"pydantic~=2.10.6", "anyio~=4.6.2"}, deps) +} + +func TestParseConstraintsRejectsMissingRequiresPython(t *testing.T) { + // Valid TOML but no requires-python is not a usable artifact; it must error + // rather than return an empty result that would be cached and fail later. + _, _, _, err := parseConstraints([]byte("[project]\nname = \"x\"\n")) + require.Error(t, err) +} + +func TestParseConstraintsDatabricksConnectNameBoundary(t *testing.T) { + // A sibling package whose name merely starts with "databricks-connect" must + // not be mistaken for the databricks-connect requirement. + toml := `[project] +requires-python = ">=3.10" + +[dependency-groups] +dev = [ + "databricks-connectors==1.0", + "databricks-connect~=17.2.0", +] +` + _, dbc, _, err := parseConstraints([]byte(toml)) + require.NoError(t, err) + assert.Equal(t, "databricks-connect~=17.2.0", dbc) +} + +func TestParseConstraintsDatabricksConnectPEP503(t *testing.T) { + // PEP 503: package names are case-insensitive and runs of -, _, . are + // equivalent. Every spelling of databricks-connect must be detected, with the + // original entry preserved verbatim in the result. + for _, entry := range []string{ + "Databricks-Connect==16.4.0", + "databricks_connect==16.4.0", + "databricks.connect==16.4.0", + "databricks-connect ~= 17.2", + } { + toml := "[project]\nrequires-python = \">=3.10\"\n\n[dependency-groups]\ndev = [\"" + entry + "\"]\n" + _, dbc, _, err := parseConstraints([]byte(toml)) + require.NoError(t, err, entry) + assert.Equal(t, entry, dbc, "entry %q", entry) + } + // A distinct sibling package must NOT match. + toml := "[project]\nrequires-python = \">=3.10\"\n\n[dependency-groups]\ndev = [\"databricks-connectors==1.0\"]\n" + _, dbc, _, err := parseConstraints([]byte(toml)) + require.NoError(t, err) + assert.Empty(t, dbc) +} + +func TestFetchConstraintsCreatesCacheDir(t *testing.T) { + // The cache directory may not exist yet on a fresh machine; the fetch must + // create it so the cache actually populates (and offline fallback works). + cacheDir := filepath.Join(t.TempDir(), "does", "not", "exist", "yet") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(sampleToml)) + })) + defer srv.Close() + + _, err := FetchConstraints(t.Context(), srv.URL, "serverless/serverless-v4", cacheDir) + require.NoError(t, err) + // The cache file was written into the freshly created directory. + written, err := os.ReadFile(filepath.Join(cacheDir, cacheFileName("serverless/serverless-v4"))) + require.NoError(t, err) + assert.Equal(t, sampleToml, string(written)) +} + +func TestCacheFileNameInjective(t *testing.T) { + // Distinct env keys that flatten to the same slug must not collide, so a + // cache hit can never serve another environment's constraints. + assert.NotEqual(t, cacheFileName("a/b"), cacheFileName("a__b")) + // The filename stays inside cacheDir (no separators leak through). + assert.NotContains(t, cacheFileName("a/b"), "/") + assert.NotContains(t, cacheFileName("a\\b"), "\\") +} + +func TestFetchConstraintsHTTP(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/serverless/serverless-v4/pyproject.toml", r.URL.Path) + _, _ = w.Write([]byte(sampleToml)) + })) + defer srv.Close() + + c, err := FetchConstraints(t.Context(), srv.URL, "serverless/serverless-v4", t.TempDir()) + require.NoError(t, err) + assert.False(t, c.FromCache) + assert.Equal(t, "databricks-connect~=17.2.0", c.DatabricksConnect) + assert.Len(t, c.ConstraintDeps, 2) +} + +func TestFetchConstraintsEnvKeyNotFound(t *testing.T) { + // A 404 for a resolved env key means the environment is not published; this + // must classify as E_ENV_UNSUPPORTED, not E_FETCH, and not fall back to cache. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "not found", http.StatusNotFound) + })) + defer srv.Close() + + _, err := FetchConstraints(t.Context(), srv.URL, "dbr/99.9.x-scala2.12", t.TempDir()) + var pe *PipelineError + require.ErrorAs(t, err, &pe) + assert.Equal(t, ErrEnvUnsupported, pe.Code) +} + +func TestFetchConstraintsTransportFailureNoCache(t *testing.T) { + // A transport failure with no cache classifies as E_FETCH. + down := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + url := down.URL + down.Close() + + _, err := FetchConstraints(t.Context(), url, "serverless/serverless-v4", t.TempDir()) + var pe *PipelineError + require.ErrorAs(t, err, &pe) + assert.Equal(t, ErrFetch, pe.Code) +} + +func TestFetchConstraintsRejectsOversizedBody(t *testing.T) { + // An over-cap response body must be rejected (classified E_FETCH here, with no + // cache to fall back to) rather than read unbounded into memory. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + big := strings.Repeat("x", int(maxConstraintBytes)+100) + _, _ = w.Write([]byte(big)) + })) + defer srv.Close() + + _, err := FetchConstraints(t.Context(), srv.URL, "serverless/serverless-v4", t.TempDir()) + var pe *PipelineError + require.ErrorAs(t, err, &pe) + assert.Equal(t, ErrFetch, pe.Code) +} + +func TestFetchConstraintsFallsBackToCache(t *testing.T) { + cacheDir := t.TempDir() + // First, a successful fetch populates the cache. + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(sampleToml)) + })) + _, err := FetchConstraints(t.Context(), good.URL, "serverless/serverless-v4", cacheDir) + require.NoError(t, err) + good.Close() + + // Now the server is down; fetch must serve the cache. + c, err := FetchConstraints(t.Context(), good.URL, "serverless/serverless-v4", cacheDir) + require.NoError(t, err) + assert.True(t, c.FromCache) +} diff --git a/libs/localenv/envkey.go b/libs/localenv/envkey.go new file mode 100644 index 00000000000..070668ba826 --- /dev/null +++ b/libs/localenv/envkey.go @@ -0,0 +1,85 @@ +package localenv + +import ( + "fmt" + "regexp" + "strconv" + "strings" +) + +// clauseRe splits a single requires-python clause into its operator (optional), +// MAJOR.MINOR version, and an optional patch component. A clause with no operator +// is a bare floor. The patch capture (group 4) is needed to interpret a strict +// ">" correctly: ">3.10" excludes all of 3.10.x, but ">3.10.5" is still satisfied +// by 3.10.6, so only the former bumps the minor. +var clauseRe = regexp.MustCompile(`^(>=|<=|===|==|~=|!=|<|>)?\s*(\d+)\.(\d+)(\.\d+)?`) + +// NormalizeServerless returns the canonical "vN" spelling of a serverless +// version accepting "4", "v4", or "V4". +func NormalizeServerless(version string) string { + return "v" + strings.TrimPrefix(strings.ToLower(version), "v") +} + +// EnvKeyForServerless returns the environment key for a serverless version. +func EnvKeyForServerless(version string) string { + return "serverless/serverless-" + NormalizeServerless(version) +} + +// EnvKeyForSparkVersion returns the environment key for a Spark version. +func EnvKeyForSparkVersion(sparkVersion string) string { + return "dbr/" + sparkVersion +} + +// PythonMinorFromRequires parses a PEP 440 requires-python string and returns +// the MAJOR.MINOR of the Python version to install: the effective lower bound. +// +// A requires-python is a comma-separated list of clauses in any order (e.g. +// "<3.13,>=3.10"). Each clause is classified by operator: +// - lower-bound / pinning (>=, >, ==, ~=, ===) or a bare MAJOR.MINOR with no +// operator establishes a floor; +// - upper-bound / exclusion (<, <=, !=) does not — those versions are capped +// or forbidden and must never be installed. +// +// The result is the highest floor across all floor clauses (so ">=3.8,>=3.11" +// yields 3.11, the version that satisfies every clause). A spec with no floor +// clause at all (e.g. "<3.13" or "!=3.12") is an error rather than a guess. +func PythonMinorFromRequires(requiresPython string) (string, error) { + bestMajor, bestMinor := -1, -1 + sawClause := false + for clause := range strings.SplitSeq(requiresPython, ",") { + clause = strings.TrimSpace(clause) + if clause == "" { + continue + } + m := clauseRe.FindStringSubmatch(clause) + if m == nil { + continue + } + sawClause = true + op := m[1] + // Upper-bound and exclusion operators never establish a floor. + if op == "<" || op == "<=" || op == "!=" { + continue + } + major, _ := strconv.Atoi(m[2]) + minor, _ := strconv.Atoi(m[3]) + hasPatch := m[4] != "" + // A strict ">" with no patch excludes the whole given minor series (PEP 440: + // ">3.10" matches neither 3.10 nor any 3.10.x), so the lowest installable + // minor is the next one up. But ">3.10.5" is still satisfied by 3.10.6, so a + // patch-qualified strict bound leaves the minor unchanged. + if op == ">" && !hasPatch { + minor++ + } + if major > bestMajor || (major == bestMajor && minor > bestMinor) { + bestMajor, bestMinor = major, minor + } + } + if bestMajor >= 0 { + return fmt.Sprintf("%d.%d", bestMajor, bestMinor), nil + } + if sawClause { + return "", fmt.Errorf("requires-python %q has no lower bound to install from", requiresPython) + } + return "", fmt.Errorf("cannot parse python version from %q", requiresPython) +} diff --git a/libs/localenv/envkey_test.go b/libs/localenv/envkey_test.go new file mode 100644 index 00000000000..898e93434b6 --- /dev/null +++ b/libs/localenv/envkey_test.go @@ -0,0 +1,67 @@ +package localenv + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestEnvKeyForServerless(t *testing.T) { + for _, in := range []string{"4", "v4", "V4"} { + assert.Equal(t, "serverless/serverless-v4", EnvKeyForServerless(in)) + } +} + +func TestEnvKeyForSparkVersion(t *testing.T) { + assert.Equal(t, "dbr/15.4.x-scala2.12", EnvKeyForSparkVersion("15.4.x-scala2.12")) +} + +func TestPythonMinorFromRequires(t *testing.T) { + cases := map[string]string{ + "==3.12.*": "3.12", + ">=3.12": "3.12", + "==3.12.3": "3.12", + "~=3.11": "3.11", + // Multi-clause specifiers: the lower bound is the version to install, + // regardless of clause order. Taking the first number would pick the + // excluded upper bound (e.g. 3.13 from "<3.13"). + "<3.13,>=3.10": "3.10", + ">=3.10,<3.13": "3.10", + ">=3.10, <3.13": "3.10", + "<4.0,>=3.9": "3.9", + "===3.11": "3.11", + // The effective floor is the HIGHEST lower bound, regardless of order. + ">=3.8,>=3.11": "3.11", + ">=3.11,>=3.8": "3.11", + // A bare floor alongside an exclusion is still a floor. + "!=3.11,3.12": "3.12", + "3.12,!=3.12.4": "3.12", + // Bare version with no operator. + "3.12": "3.12", + // Whitespace and patch components tolerated. + ">= 3.10 , < 3.13": "3.10", + // Strict ">" with no patch excludes the whole minor series (PEP 440), so + // the floor is the next minor up. + ">3.10": "3.11", + ">3.10,<3.13": "3.11", + ">=3.9,>3.10": "3.11", + // Strict ">" WITH a patch does not exclude the minor series: 3.10.6 + // satisfies ">3.10.5", so the floor stays 3.10. + ">3.10.5": "3.10", + ">3.10.5,<3.13": "3.10", + ">=3.10.2": "3.10", + } + for in, want := range cases { + got, err := PythonMinorFromRequires(in) + require.NoError(t, err) + assert.Equal(t, want, got, "input %q", in) + } + + // No usable floor: only upper-bound / exclusion clauses. Must error rather + // than select a forbidden/capped version. + for _, in := range []string{"<3.13", "<=3.12", "!=3.12", "<3.13,!=3.12", "garbage", ""} { + _, err := PythonMinorFromRequires(in) + assert.Error(t, err, "input %q must error", in) + } +} diff --git a/libs/localenv/result.go b/libs/localenv/result.go new file mode 100644 index 00000000000..1de8f66ba50 --- /dev/null +++ b/libs/localenv/result.go @@ -0,0 +1,195 @@ +package localenv + +import "fmt" + +// Command path components, defined once so a rename touches a single place +// (spec §0 / invariant 8 / scenario 21). The cmd layer builds the Cobra +// command tree from CommandGroup/CommandSubgroup/CommandVerb; the --json +// "command" field uses CommandName. No other string re-spells the command path. +const ( + CommandGroup = "local-env" + CommandSubgroup = "python" + CommandVerb = "sync" + CommandName = CommandGroup + " " + CommandSubgroup + " " + CommandVerb + + // SchemaVersion is the version of the --json output contract (spec §6). + // Bump it on any breaking change to the JSON shape. + SchemaVersion = 1 +) + +// Mode is the provisioning mode: a full environment (default) or the +// constraints-only variant that omits the databricks-connect dependency. +type Mode int + +const ( + ModeDefault Mode = iota + ModeConstraintsOnly +) + +// String returns the JSON/text spelling of the mode ("default" | "constraints-only"). +func (m Mode) String() string { + if m == ModeConstraintsOnly { + return "constraints-only" + } + return "default" +} + +// PhaseName is a canonical execution phase (spec §3 / §6). The set is fixed and +// ordered; the --json "phases" array reports every phase in this order. +type PhaseName string + +const ( + PhasePreflight PhaseName = "preflight" + PhaseResolve PhaseName = "resolve" + PhaseFetch PhaseName = "fetch" + PhaseMerge PhaseName = "merge" + PhaseProvision PhaseName = "provision" + PhaseValidate PhaseName = "validate" +) + +// Phase status values (spec §6.2). +const ( + StatusOK = "ok" + StatusError = "error" + StatusPending = "pending" +) + +// ErrorCode is a stable failure-class identifier surfaced in --json error.code +// (spec §7). Values are compared via the ErrorCode constants, never by +// string-matching messages, and are defined once here. +type ErrorCode string + +const ( + ErrNoTarget ErrorCode = "E_NO_TARGET" + ErrManagerUnsupported ErrorCode = "E_MANAGER_UNSUPPORTED" + ErrUvMissing ErrorCode = "E_UV_MISSING" + ErrNotWritable ErrorCode = "E_NOT_WRITABLE" + ErrResolve ErrorCode = "E_RESOLVE" + ErrEnvUnsupported ErrorCode = "E_ENV_UNSUPPORTED" + ErrFetch ErrorCode = "E_FETCH" + ErrWrite ErrorCode = "E_WRITE" + ErrMerge ErrorCode = "E_MERGE" + ErrPythonInstall ErrorCode = "E_PYTHON_INSTALL" + ErrProvision ErrorCode = "E_PROVISION" + ErrValidate ErrorCode = "E_VALIDATE" +) + +// PipelineError is a failure carrying a stable code, the phase at which it +// occurred, and whether disk was mutated before the failure. It marshals to the +// --json error object (spec §6.2). Code and FailurePhase are the stable +// contract; Err holds the wrapped cause for errors.Is/As and is not serialized. +type PipelineError struct { + Code ErrorCode `json:"code"` + FailurePhase PhaseName `json:"failurePhase"` + Msg string `json:"message"` + DiskMutated bool `json:"diskMutated"` + Err error `json:"-"` +} + +func (e *PipelineError) Error() string { + if e.Err != nil { + return e.Msg + ": " + e.Err.Error() + } + return e.Msg +} + +func (e *PipelineError) Unwrap() error { + return e.Err +} + +// NewError creates a PipelineError with a code and message. FailurePhase and +// DiskMutated are filled in by the pipeline when it records the failure. The +// message is formatted with fmt.Sprintf(format, args...); err may be nil. +func NewError(code ErrorCode, err error, format string, args ...any) *PipelineError { + return &PipelineError{ + Code: code, + Msg: fmt.Sprintf(format, args...), + Err: err, + } +} + +// TargetInfo is the resolved compute target (spec §6 "target"). Source records +// which of the four precedence sources was used. SparkVersion is the raw cluster +// runtime string the resolver read; it is folded into EnvKey (dbr/) +// and is not part of the JSON contract, kept only as intermediate resolver state. +type TargetInfo struct { + Source string `json:"source"` + ClusterID string `json:"clusterId,omitempty"` + ServerlessVersion string `json:"serverlessVersion,omitempty"` + EnvKey string `json:"envKey"` + + SparkVersion string `json:"-"` +} + +// ResolvedInfo is the resolved environment definition (spec §6 "resolved"). +// DBConnectVersion is omitted in constraints-only mode. +type ResolvedInfo struct { + PythonVersion string `json:"pythonVersion"` + DBConnectVersion string `json:"dbconnectVersion,omitempty"` + ArtifactSource string `json:"artifactSource"` +} + +// Plan describes the changes a --check run would apply (spec §6.3). +// ChangedRegions is retained for text output only and is not serialized. +type Plan struct { + WouldWrite string `json:"wouldWrite"` + WouldBackup string `json:"wouldBackup,omitempty"` + WouldInstallPython string `json:"wouldInstallPython,omitempty"` + Diff string `json:"diff"` + + ChangedRegions []string `json:"-"` +} + +// PhaseStatus is one entry in the --json "phases" array (spec §6). Detail is +// used for human-readable text output only and is not serialized. +type PhaseStatus struct { + Phase PhaseName `json:"phase"` + Status string `json:"status"` + + Detail string `json:"-"` +} + +// Warning is a non-fatal advisory surfaced in --json "warnings" (spec §6). +type Warning struct { + Code string `json:"code"` + Message string `json:"message"` +} + +// Result is the full outcome of a sync run and the root of the --json object +// (spec §6). Field order matches the spec's schema so JSON key order is stable. +// +// Phases and Warnings are non-omitempty slices, so they must always be non-nil +// before marshalling or the --json contract would emit "null" instead of "[]" — +// a distinction that trips JSON consumers and golden diffs. Construct a Result +// with NewResult (or otherwise seed both) rather than a bare Result{} literal. +type Result struct { + SchemaVersion int `json:"schemaVersion"` + Command string `json:"command"` + OK bool `json:"ok"` + Mode string `json:"mode"` + DryRun bool `json:"dryRun"` + Target *TargetInfo `json:"target,omitempty"` + Resolved *ResolvedInfo `json:"resolved,omitempty"` + Greenfield bool `json:"greenfield"` + Plan *Plan `json:"plan,omitempty"` + VenvPath string `json:"venvPath,omitempty"` + Phases []PhaseStatus `json:"phases"` + Warnings []Warning `json:"warnings"` + Error *PipelineError `json:"error"` + BackupPath string `json:"backupPath,omitempty"` + // DurationMs is part of the §6 contract but reserved for now: the pipeline + // does not measure wall time (a real clock would make acceptance goldens + // non-deterministic), so it is always emitted as 0 until timing is wired + // through a clock the tests can control. + DurationMs int64 `json:"durationMs"` +} + +// NewResult returns a Result with the non-omitempty slice fields initialized to +// empty (non-nil) slices, so the --json output always renders "phases": [] and +// "warnings": [] rather than "null". Callers fill in the remaining fields. +func NewResult() *Result { + return &Result{ + Phases: []PhaseStatus{}, + Warnings: []Warning{}, + } +} diff --git a/libs/localenv/result_test.go b/libs/localenv/result_test.go new file mode 100644 index 00000000000..75fe60ea66d --- /dev/null +++ b/libs/localenv/result_test.go @@ -0,0 +1,45 @@ +package localenv + +import ( + "encoding/json" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPipelineErrorWrapsAndExposesCode(t *testing.T) { + base := errors.New("boom") + err := NewError(ErrFetch, base, "fetch %s", "x") + assert.Equal(t, "fetch x: boom", err.Error()) + assert.Equal(t, ErrFetch, err.Code) + assert.ErrorIs(t, err, base) +} + +func TestModeString(t *testing.T) { + assert.Equal(t, "default", ModeDefault.String()) + assert.Equal(t, "constraints-only", ModeConstraintsOnly.String()) +} + +func TestCommandName(t *testing.T) { + // The --json "command" field and all help text derive from these; the + // three-part path must join to the full command a user types. + assert.Equal(t, "local-env python sync", CommandName) +} + +func TestNewResultEmitsEmptyArraysNotNull(t *testing.T) { + // The --json contract requires phases/warnings to render as [] not null; + // NewResult seeds them so consumers and golden diffs see a stable shape. + b, err := json.Marshal(NewResult()) + require.NoError(t, err) + s := string(b) + assert.Contains(t, s, `"phases":[]`) + assert.Contains(t, s, `"warnings":[]`) + assert.NotContains(t, s, `"phases":null`) + assert.NotContains(t, s, `"warnings":null`) + // A bare Result{} literal is the shape NewResult exists to avoid. + bare, err := json.Marshal(&Result{}) + require.NoError(t, err) + assert.Contains(t, string(bare), `"phases":null`, "sanity: bare literal is the null case") +} diff --git a/libs/localenv/target.go b/libs/localenv/target.go new file mode 100644 index 00000000000..b3ccb3171e1 --- /dev/null +++ b/libs/localenv/target.go @@ -0,0 +1,149 @@ +package localenv + +import ( + "context" + "fmt" + "strings" +) + +// ComputeClient is a narrow seam over the SDK so tests can stub it. +type ComputeClient interface { + // GetClusterSparkVersion returns the Spark version string for a cluster. + GetClusterSparkVersion(ctx context.Context, clusterID string) (string, error) + // GetJobSparkVersion returns either a Spark version (isServerless=false) or a + // serverless marker (isServerless=true) for a job, plus a recorded version string. + GetJobSparkVersion(ctx context.Context, jobID string) (sparkVersion string, isServerless bool, version string, err error) +} + +// TargetFlags holds the mutually-exclusive compute target flags from the CLI. +type TargetFlags struct { + Cluster string + Serverless string + Job string +} + +// BundleTarget is the three-state result of reading the bundle's configured +// target. Selected=false means nothing was configured. +type BundleTarget struct { + ClusterID string + Serverless bool + Selected bool +} + +// noTargetMessage is the actionable message shown when no target is selected, +// matching spec §2.3. +const noTargetMessage = "No compute target is selected. Select a cluster or serverless target, or pass --cluster / --serverless / --job" + +// ValidateTargetFlags returns an error if more than one of the three flags is set. +// Cobra marks them mutually exclusive too; this guards the library path. +func ValidateTargetFlags(f TargetFlags) error { + var set []string + if f.Cluster != "" { + set = append(set, "--cluster") + } + if f.Serverless != "" { + set = append(set, "--serverless") + } + if f.Job != "" { + set = append(set, "--job") + } + if len(set) > 1 { + return fmt.Errorf("flags %s are mutually exclusive; specify at most one", strings.Join(set, " and ")) + } + return nil +} + +// ResolveTarget resolves the compute target using ordered precedence: +// --cluster flag → --serverless flag → --job flag → bundle target. +// PythonVersion is left empty; it is filled later from constraint data. +// +// Incompatible flags are rejected up front: without this a library caller that +// bypasses Cobra (which also marks the flags mutually exclusive) and passes more +// than one target flag would have all but the first precedence branch silently +// ignored, resolving a different target than requested. +func ResolveTarget(ctx context.Context, f TargetFlags, c ComputeClient, bt BundleTarget) (*TargetInfo, error) { + if err := ValidateTargetFlags(f); err != nil { + return nil, NewError(ErrResolve, err, "invalid compute target flags") + } + + if f.Cluster != "" { + v, err := c.GetClusterSparkVersion(ctx, f.Cluster) + if err != nil { + return nil, NewError(ErrResolve, err, "resolving cluster %s", f.Cluster) + } + return &TargetInfo{ + Source: "cluster", + ClusterID: f.Cluster, + SparkVersion: v, + EnvKey: EnvKeyForSparkVersion(v), + }, nil + } + + if f.Serverless != "" { + return &TargetInfo{ + Source: "serverless", + ServerlessVersion: NormalizeServerless(f.Serverless), + EnvKey: EnvKeyForServerless(f.Serverless), + }, nil + } + + if f.Job != "" { + sparkVersion, isServerless, version, err := c.GetJobSparkVersion(ctx, f.Job) + if err != nil { + return nil, NewError(ErrResolve, err, "resolving job %s", f.Job) + } + if isServerless { + // Default to v4 when the job is serverless; the serverless env version + // is not recorded in the bundle/project (documented stand-in from the + // original script, spec §4.3). + v := version + if v == "" { + v = "v4" + } + return &TargetInfo{ + Source: "job", + ServerlessVersion: NormalizeServerless(v), + EnvKey: EnvKeyForServerless(v), + }, nil + } + // Classic compute: the Spark version is the first return per the + // GetJobSparkVersion contract, not the recorded-version third return. + return &TargetInfo{ + Source: "job", + SparkVersion: sparkVersion, + EnvKey: EnvKeyForSparkVersion(sparkVersion), + }, nil + } + + // Fall back to bundle target. + if !bt.Selected { + return nil, NewError(ErrNoTarget, nil, "%s", noTargetMessage) + } + + if bt.Serverless { + // Default to serverless-v4: the serverless env version is not recorded + // in the bundle/project (documented stand-in from the original script). + return &TargetInfo{ + Source: "bundle", + ServerlessVersion: "v4", + EnvKey: EnvKeyForServerless("v4"), + }, nil + } + + if bt.ClusterID != "" { + v, err := c.GetClusterSparkVersion(ctx, bt.ClusterID) + if err != nil { + return nil, NewError(ErrResolve, err, "resolving bundle cluster %s", bt.ClusterID) + } + return &TargetInfo{ + Source: "bundle", + ClusterID: bt.ClusterID, + SparkVersion: v, + EnvKey: EnvKeyForSparkVersion(v), + }, nil + } + + // Bundle target is selected but has neither serverless nor a cluster ID — + // treat this the same as nothing selected so the user gets a clear message. + return nil, NewError(ErrNoTarget, nil, "%s", noTargetMessage) +} diff --git a/libs/localenv/target_test.go b/libs/localenv/target_test.go new file mode 100644 index 00000000000..22beaee3382 --- /dev/null +++ b/libs/localenv/target_test.go @@ -0,0 +1,106 @@ +package localenv + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type stubCompute struct { + clusterVersion string + clusterErr error +} + +func (s stubCompute) GetClusterSparkVersion(_ context.Context, _ string) (string, error) { + return s.clusterVersion, s.clusterErr +} + +func (s stubCompute) GetJobSparkVersion(_ context.Context, _ string) (string, bool, string, error) { + return "", false, "", nil +} + +func TestResolveServerlessFlag(t *testing.T) { + ti, err := ResolveTarget(t.Context(), TargetFlags{Serverless: "v4"}, stubCompute{}, BundleTarget{}) + require.NoError(t, err) + assert.Equal(t, "serverless", ti.Source) + assert.Equal(t, "v4", ti.ServerlessVersion) + assert.Equal(t, "serverless/serverless-v4", ti.EnvKey) +} + +func TestResolveClusterFlag(t *testing.T) { + c := stubCompute{clusterVersion: "15.4.x-scala2.12"} + ti, err := ResolveTarget(t.Context(), TargetFlags{Cluster: "abc"}, c, BundleTarget{}) + require.NoError(t, err) + assert.Equal(t, "cluster", ti.Source) + assert.Equal(t, "15.4.x-scala2.12", ti.SparkVersion) + assert.Equal(t, "dbr/15.4.x-scala2.12", ti.EnvKey) + assert.Equal(t, "abc", ti.ClusterID) +} + +func TestResolveClusterFlagError(t *testing.T) { + c := stubCompute{clusterErr: errors.New("cluster not found")} + _, err := ResolveTarget(t.Context(), TargetFlags{Cluster: "abc"}, c, BundleTarget{}) + var pe *PipelineError + require.ErrorAs(t, err, &pe) + assert.Equal(t, ErrResolve, pe.Code) +} + +func TestResolveBundleNothingSelected(t *testing.T) { + _, err := ResolveTarget(t.Context(), TargetFlags{}, stubCompute{}, BundleTarget{Selected: false}) + var pe *PipelineError + require.ErrorAs(t, err, &pe) + assert.Equal(t, ErrNoTarget, pe.Code) +} + +func TestResolveBundleServerless(t *testing.T) { + ti, err := ResolveTarget(t.Context(), TargetFlags{}, stubCompute{}, BundleTarget{Selected: true, Serverless: true}) + require.NoError(t, err) + assert.Equal(t, "bundle", ti.Source) + assert.Equal(t, "serverless/serverless-v4", ti.EnvKey) +} + +// jobStubCompute returns distinct values for the first (sparkVersion) and third +// (recorded version) results of GetJobSparkVersion so the classic-compute branch +// can be checked against the documented contract (it must use the first). +type jobStubCompute struct { + sparkVersion string + isServerless bool + version string +} + +func (jobStubCompute) GetClusterSparkVersion(_ context.Context, _ string) (string, error) { + return "", nil +} + +func (s jobStubCompute) GetJobSparkVersion(_ context.Context, _ string) (string, bool, string, error) { + return s.sparkVersion, s.isServerless, s.version, nil +} + +func TestResolveJobClassicUsesSparkVersionReturn(t *testing.T) { + // Contract: for a classic-compute job the Spark version is the FIRST return. + // The third "recorded version" return differs here to catch use of the wrong one. + c := jobStubCompute{sparkVersion: "15.4.x-scala2.12", isServerless: false, version: "wrong-recorded"} + ti, err := ResolveTarget(t.Context(), TargetFlags{Job: "42"}, c, BundleTarget{}) + require.NoError(t, err) + assert.Equal(t, "job", ti.Source) + assert.Equal(t, "15.4.x-scala2.12", ti.SparkVersion) + assert.Equal(t, "dbr/15.4.x-scala2.12", ti.EnvKey) +} + +func TestValidateTargetFlagsMutuallyExclusive(t *testing.T) { + assert.Error(t, ValidateTargetFlags(TargetFlags{Cluster: "a", Serverless: "v4"})) + assert.NoError(t, ValidateTargetFlags(TargetFlags{Cluster: "a"})) +} + +func TestResolveTargetRejectsConflictingFlags(t *testing.T) { + // ResolveTarget must reject incompatible flags rather than silently taking + // the first precedence branch, so a library caller bypassing Cobra can't + // resolve a different target than it asked for. + _, err := ResolveTarget(t.Context(), TargetFlags{Cluster: "c", Serverless: "v4"}, stubCompute{}, BundleTarget{}) + var pe *PipelineError + require.ErrorAs(t, err, &pe) + assert.Equal(t, ErrResolve, pe.Code) +}