From 9aeffa93f732e1ef6e74d110cb0d7e1f48ac2373 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Tue, 21 Jul 2026 18:50:34 +0000 Subject: [PATCH] air run: git-pin code_source snapshots Add git pinning to the CLI code_source snapshot path: when code_source.snapshot.git sets a commit or branch, resolve it locally and `git archive` that commit into the tarball (instead of plain-tarring the working tree), then upload it through the same DABs artifact plumbing as the tar path. buildSnapshotTarball now resolves a snapshotPlan (resolveSnapshotPlan) and uses createGitArchiveSnapshot for the git_archive mode. Reintroduces the git helpers (snapshot_git.go, snapshot_resolve.go, snapshot_cachekey.go) removed in the tar PR. Git resolution stays CLI-side (it is a train.yaml construct with no DABs config surface); only the tarball upload reuses DABs plumbing. Co-authored-by: Isaac --- experimental/air/cmd/runsubmit_test.go | 41 +++ experimental/air/cmd/snapshot_cachekey.go | 34 ++ .../air/cmd/snapshot_cachekey_test.go | 62 ++++ experimental/air/cmd/snapshot_dabs.go | 29 +- experimental/air/cmd/snapshot_git.go | 309 ++++++++++++++++++ experimental/air/cmd/snapshot_git_test.go | 296 +++++++++++++++++ experimental/air/cmd/snapshot_package.go | 26 +- experimental/air/cmd/snapshot_package_test.go | 45 ++- experimental/air/cmd/snapshot_resolve.go | 120 +++++++ experimental/air/cmd/snapshot_resolve_test.go | 114 +++++++ 10 files changed, 1058 insertions(+), 18 deletions(-) create mode 100644 experimental/air/cmd/snapshot_cachekey.go create mode 100644 experimental/air/cmd/snapshot_cachekey_test.go create mode 100644 experimental/air/cmd/snapshot_git.go create mode 100644 experimental/air/cmd/snapshot_git_test.go create mode 100644 experimental/air/cmd/snapshot_resolve.go create mode 100644 experimental/air/cmd/snapshot_resolve_test.go diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index e5e4dace41..9a3d45cca7 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -199,6 +199,47 @@ code_source: assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) } +// A git-pinned code_source is git-archived at the commit, uploaded via DABs' artifact +// plumbing, and its remote code_source_path attached to the submitted task. +func TestSubmitWorkloadWithGitPinnedCodeSource(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + var got jobs.SubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return jobs.SubmitRunResponse{RunId: 555} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + // A git repo committed at HEAD, referenced by commit so packaging is git_archive. + repo := newTestRepo(t) + writeRepoFile(t, repo, "train.py", "print()") + sha := commitAll(t, repo, "init") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` + git: + commit: ` + sha + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + ctx := cmdio.MockDiscard(t.Context()) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem") + require.NoError(t, err) + + at := got.Tasks[0].AiRuntimeTask + assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/.internal/") + assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) +} + func TestSubmitWorkloadGuards(t *testing.T) { w := newFakeWorkspaceClient(t) cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) diff --git a/experimental/air/cmd/snapshot_cachekey.go b/experimental/air/cmd/snapshot_cachekey.go new file mode 100644 index 0000000000..44c58ee903 --- /dev/null +++ b/experimental/air/cmd/snapshot_cachekey.go @@ -0,0 +1,34 @@ +package aircmd + +// This file packages a local code directory into a tarball, uploads it to the +// workspace (or a Volume), and records git provenance sidecars for cache +// invalidation — the Go port of the Python CLI's code_source snapshot path. + +import ( + "crypto/sha256" + "encoding/hex" + "slices" + "strings" +) + +// snapshotPackagingVersion is bumped when packaging logic changes in a way that invalidates existing caches +const snapshotPackagingVersion = "v1" + +// computeSnapshotCacheKey returns a stable cache key for a snapshot tarball: the +// SHA-256 digest of (commitSHA, normalized includePaths, snapshotPackagingVersion). +// Changing any input yields a different entry. +func computeSnapshotCacheKey(commitSHA string, includePaths []string) string { + var normalizedPaths string + if len(includePaths) > 0 { + trimmed := make([]string, len(includePaths)) + for i, p := range includePaths { + trimmed[i] = strings.TrimSpace(p) + } + slices.Sort(trimmed) + normalizedPaths = strings.Join(trimmed, "\n") + } + + keyMaterial := commitSHA + "\n" + normalizedPaths + "\n" + snapshotPackagingVersion + sum := sha256.Sum256([]byte(keyMaterial)) + return hex.EncodeToString(sum[:]) +} diff --git a/experimental/air/cmd/snapshot_cachekey_test.go b/experimental/air/cmd/snapshot_cachekey_test.go new file mode 100644 index 0000000000..5743217c00 --- /dev/null +++ b/experimental/air/cmd/snapshot_cachekey_test.go @@ -0,0 +1,62 @@ +package aircmd + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type goldenCase struct { + Name string `json:"name"` + CommitSHA string `json:"commit_sha"` + IncludePaths []string `json:"include_paths"` + CacheKey string `json:"cache_key"` +} + +// TestComputeSnapshotCacheKeyGolden asserts byte-for-byte parity with golden +// fixtures across the local-only matrix (commit + include_paths permutations). +func TestComputeSnapshotCacheKeyGolden(t *testing.T) { + data, err := os.ReadFile(filepath.Join("testdata", "cache_keys.json")) + require.NoError(t, err) + + var cases []goldenCase + require.NoError(t, json.Unmarshal(data, &cases)) + require.NotEmpty(t, cases) + + for _, tc := range cases { + t.Run(tc.Name, func(t *testing.T) { + assert.Equal(t, tc.CacheKey, computeSnapshotCacheKey(tc.CommitSHA, tc.IncludePaths)) + }) + } +} + +// TestComputeSnapshotCacheKeyProperties pins the normalization behavior the golden cases +// encode, so a regression is legible without decoding hashes. +func TestComputeSnapshotCacheKeyProperties(t *testing.T) { + sha := "a3492b801c0ffee00000000000000000000dead" + + // Order-independent: sorting means unsorted input yields the sorted key. + assert.Equal(t, + computeSnapshotCacheKey(sha, []string{"a", "b", "c"}), + computeSnapshotCacheKey(sha, []string{"c", "a", "b"}), + ) + + // nil and empty include_paths are equivalent (both contribute an empty line). + assert.Equal(t, computeSnapshotCacheKey(sha, nil), computeSnapshotCacheKey(sha, []string{})) + + // Paths are trimmed before hashing. + assert.Equal(t, + computeSnapshotCacheKey(sha, []string{"research", "data"}), + computeSnapshotCacheKey(sha, []string{" research ", " data "}), + ) + + // Duplicates are NOT collapsed — they are sorted and kept, matching Python. + assert.NotEqual(t, computeSnapshotCacheKey(sha, []string{"x", "y"}), computeSnapshotCacheKey(sha, []string{"x", "x", "y"})) + + // The version constant participates: a different version is a different key. + assert.NotEqual(t, snapshotPackagingVersion, "") +} diff --git a/experimental/air/cmd/snapshot_dabs.go b/experimental/air/cmd/snapshot_dabs.go index afe63fd6ba..6e4b2b5989 100644 --- a/experimental/air/cmd/snapshot_dabs.go +++ b/experimental/air/cmd/snapshot_dabs.go @@ -31,16 +31,13 @@ func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, s if snap.RemoteVolume != nil { return snapshotResult{}, errors.New("code_source.snapshot.remote_volume is not yet supported") } - if snap.Git != nil { - return snapshotResult{}, errors.New("code_source.snapshot.git is not yet supported on this path") - } repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) if err != nil { return snapshotResult{}, err } - tarball, cleanup, err := buildSnapshotTarball(ctx, repoPath, snap.IncludePaths) + tarball, cleanup, err := buildSnapshotTarball(ctx, repoPath, snap.Git, snap.IncludePaths) if err != nil { return snapshotResult{}, err } @@ -49,10 +46,11 @@ func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, s return uploadSnapshotViaDABs(ctx, w, tarball) } -// buildSnapshotTarball writes a plain-tar snapshot of repoPath to a temp file named +// buildSnapshotTarball writes a snapshot of repoPath to a temp file named // .tar.gz (the basename becomes the uploaded filename), returning the path -// and a cleanup func. -func buildSnapshotTarball(ctx context.Context, repoPath string, includePaths []string) (string, func(), error) { +// and a cleanup func. With a git ref it archives the resolved commit (git archive); +// otherwise it plain-tars the working tree. +func buildSnapshotTarball(ctx context.Context, repoPath string, ref *gitRef, includePaths []string) (string, func(), error) { noop := func() {} tmp, err := os.MkdirTemp("", "air-snapshot-*") if err != nil { @@ -61,7 +59,22 @@ func buildSnapshotTarball(ctx context.Context, repoPath string, includePaths []s cleanup := func() { _ = os.RemoveAll(tmp) } tarball := filepath.Join(tmp, filepath.Base(repoPath)+".tar.gz") - if err := createPlainTarball(ctx, repoPath, tarball, includePaths); err != nil { + dirName := filepath.Base(repoPath) + + git := newGitRepo(repoPath) + plan, err := resolveSnapshotPlan(ctx, git, ref, includePaths) + if err != nil { + cleanup() + return "", noop, err + } + + switch plan.mode { + case modeGitArchive: + err = createGitArchiveSnapshot(ctx, git, plan.commitSHA, tarball, dirName, includePaths) + default: + err = createPlainTarball(ctx, repoPath, tarball, includePaths) + } + if err != nil { cleanup() return "", noop, err } diff --git a/experimental/air/cmd/snapshot_git.go b/experimental/air/cmd/snapshot_git.go new file mode 100644 index 0000000000..616b3049f7 --- /dev/null +++ b/experimental/air/cmd/snapshot_git.go @@ -0,0 +1,309 @@ +package aircmd + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "os/exec" + "strings" + "time" +) + +// Local, no-network git introspection and the git-state provenance sidecar, ported +// from the Python CLI's cli/utils/git_state.py. The remote-fetch helpers +// (fetch_branch_sha, remote detection, partial clone) are deliberately not ported: +// the snapshot archives the local copy only, so a ref must resolve to a local commit. + +// gitRepo runs git subcommands scoped to one repository via `git -C`. Arguments are +// passed as a slice, never a shell string, so branch/commit values can't inject. +type gitRepo struct { + path string +} + +func newGitRepo(path string) gitRepo { + return gitRepo{path: path} +} + +// run executes `git ` and returns stdout; a non-zero exit wraps stderr. +func (g gitRepo) run(ctx context.Context, args ...string) (string, error) { + out, err := g.runBytes(ctx, args...) + return string(out), err +} + +// runBytes is run returning raw stdout bytes, for the dirty-diff capture which needs +// exact bytes and a size measurement. +func (g gitRepo) runBytes(ctx context.Context, args ...string) ([]byte, error) { + full := append([]string{"-C", g.path}, args...) + cmd := exec.CommandContext(ctx, "git", full...) + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + if err := cmd.Run(); err != nil { + msg := strings.TrimSpace(stderr.String()) + if msg == "" { + return nil, fmt.Errorf("git %s: %w", strings.Join(args, " "), err) + } + return nil, fmt.Errorf("git %s: %w: %s", strings.Join(args, " "), err, msg) + } + return stdout.Bytes(), nil +} + +// isRepository reports whether the path is inside a git work tree. Using +// `rev-parse --is-inside-work-tree` (not a .git lookup) means a subdirectory of a +// repo counts — the common case when root_path is a subfolder of a monorepo. +func (g gitRepo) isRepository(ctx context.Context) bool { + out, err := g.run(ctx, "rev-parse", "--is-inside-work-tree") + if err != nil { + return false + } + return strings.TrimSpace(out) == "true" +} + +// headSHA returns the current HEAD commit SHA. +func (g gitRepo) headSHA(ctx context.Context) (string, error) { + out, err := g.run(ctx, "rev-parse", "HEAD") + if err != nil { + return "", err + } + return strings.TrimSpace(out), nil +} + +// hasUncommittedChanges reports whether there are staged or unstaged changes under +// the repo subtree. The `-- .` pathspec scopes the check so a subfolder snapshot +// considers only changes that could land in it. +func (g gitRepo) hasUncommittedChanges(ctx context.Context) (bool, error) { + out, err := g.run(ctx, "status", "--porcelain", "--", ".") + if err != nil { + return false, err + } + return strings.TrimSpace(out) != "", nil +} + +// hasUncommittedChangesInPaths reports whether there are uncommitted changes within +// the include paths (empty includePaths yields false). +// +// The pathspecs limit `git status` to those subtrees (it is O(working tree), slow on +// a large monorepo) and already scope the output to what could land in the snapshot. +// Unlike the Python source we don't re-parse the entries to filter by name: git +// reports a rename as `R \x00`, so a name-based re-filter keys off the old +// path and could miss a rename into an include path. The only caller needs the bool. +func (g gitRepo) hasUncommittedChangesInPaths(ctx context.Context, includePaths []string) (bool, error) { + var pathspecs []string + for _, p := range includePaths { + if s := strings.TrimRight(p, "/"); s != "" { + pathspecs = append(pathspecs, s) + } + } + if len(pathspecs) == 0 { + return false, nil + } + + args := append([]string{"status", "--porcelain", "--"}, pathspecs...) + out, err := g.run(ctx, args...) + if err != nil { + return false, err + } + return strings.TrimSpace(out) != "", nil +} + +// resolveLocalBranchSHA resolves a branch to its local-HEAD commit. No remote is +// contacted; the branch must exist locally. +func (g gitRepo) resolveLocalBranchSHA(ctx context.Context, branch string) (string, error) { + out, err := g.run(ctx, "rev-parse", "refs/heads/"+branch) + if err != nil { + return "", fmt.Errorf("failed to resolve local branch %q; ensure the branch exists locally and root_path is correct: %w", branch, err) + } + return strings.TrimSpace(out), nil +} + +// commitExistsLocally reports whether commitSHA is in the local object store, without +// triggering a promisor/lazy fetch. +func (g gitRepo) commitExistsLocally(ctx context.Context, commitSHA string) bool { + _, err := g.run(ctx, "cat-file", "-e", commitSHA) + return err == nil +} + +// currentBranch returns the branch name, or "" for a detached HEAD or on error. +func (g gitRepo) currentBranch(ctx context.Context) string { + out, err := g.run(ctx, "rev-parse", "--abbrev-ref", "HEAD") + if err != nil { + return "" + } + branch := strings.TrimSpace(out) + if branch == "HEAD" { + return "" + } + return branch +} + +// remoteURL returns the URL of the named remote, or "" if it has none. +func (g gitRepo) remoteURL(ctx context.Context, remoteName string) string { + out, err := g.run(ctx, "remote", "get-url", remoteName) + if err != nil { + return "" + } + return strings.TrimSpace(out) +} + +// mergeBaseWithUpstream resolves the merge-base of HEAD and a likely upstream ref, +// trying /HEAD, /main, then /master. It reads only local remote-tracking +// refs (no fetch), returning "" if none resolve. +func (g gitRepo) mergeBaseWithUpstream(ctx context.Context, remoteName string) string { + for _, ref := range []string{remoteName + "/HEAD", remoteName + "/main", remoteName + "/master"} { + out, err := g.run(ctx, "merge-base", "HEAD", ref) + if err != nil { + continue + } + if base := strings.TrimSpace(out); base != "" { + return base + } + } + return "" +} + +// validateIncludePathsExist checks that every include path exists at commitSHA. +// `git ls-tree` (without -d, so both blobs and trees count) reports an entry when the +// path exists; empty output means missing. +func (g gitRepo) validateIncludePathsExist(ctx context.Context, commitSHA string, includePaths []string) error { + var missing []string + for _, p := range includePaths { + out, err := g.run(ctx, "ls-tree", commitSHA, p) + if err != nil { + return err + } + if strings.TrimSpace(out) == "" { + missing = append(missing, p) + } + } + if len(missing) > 0 { + return fmt.Errorf("include_paths do not exist at commit %s: %s", shortSHA(commitSHA), strings.Join(missing, ", ")) + } + return nil +} + +// shortSHA abbreviates a commit SHA to 8 chars for log/error messages, tolerating +// user-supplied abbreviations shorter than that. +func shortSHA(sha string) string { + return sha[:min(len(sha), 8)] +} + +// --- git-state provenance sidecar (git_state.json + git_diff.patch) --- +// +// The backend reads git_state.json next to the tarball to tag the MLflow run with +// base/tip/dirty provenance, and logs git_diff.patch when the tree was dirty. +// Producing the sidecar is best-effort: callers warn and continue, never fail submit. + +// snapshotStateSchemaVersion is the git_state.json schema version. Bump only in +// coordination with the backend reader. +const snapshotStateSchemaVersion = 1 + +// defaultRemoteName is the remote consulted for merge-base and repo URL (local refs +// only — the remote-fetch path is gone). +const defaultRemoteName = "origin" + +// dirtyDiffSizeCapBytes caps the git_diff.patch sidecar; a larger diff records +// size_exceeded and is skipped to keep the upload small. +const dirtyDiffSizeCapBytes = 1024 * 1024 + +// dirtyDiffTimeout bounds `git diff HEAD` so provenance never delays submission. +const dirtyDiffTimeout = 5 * time.Second + +// packaging_mode values: how the uploaded tarball was produced. +const ( + packagingModeGitArchive = "git_archive" + packagingModePlainTar = "plain_tar" +) + +// diff_status values recorded in the sidecar. +const ( + diffStatusClean = "clean" + diffStatusCaptured = "captured" + diffStatusSizeExceeded = "size_exceeded" + diffStatusTimeout = "timeout" +) + +// gitStateSidecar is the git_state.json record. Field names and the null-for-absent +// encoding match the Python source, so nullable fields are *string (absent → null). +type gitStateSidecar struct { + SchemaVersion int `json:"schema_version"` + PackagingMode string `json:"packaging_mode"` + BaseCommit *string `json:"base_commit"` + TipCommit *string `json:"tip_commit"` + Branch *string `json:"branch"` + RepoURL *string `json:"repo_url"` + Dirty bool `json:"dirty"` + DiffStatus string `json:"diff_status"` + DiffPath *string `json:"diff_path"` + GeneratedAtUTC string `json:"generated_at_utc"` +} + +// nilIfEmpty maps "" to nil so an absent value serializes as JSON null. +func nilIfEmpty(s string) *string { + if s == "" { + return nil + } + return &s +} + +// buildGitStateSidecar gathers git provenance. pinnedTip overrides the HEAD-derived +// tip for git_archive (the tarball reflects that commit, not HEAD); pass "" for +// plain_tar. Metadata is best-effort — unavailable fields become null. +func buildGitStateSidecar(ctx context.Context, git gitRepo, packagingMode, pinnedTip string, now time.Time) (gitStateSidecar, error) { + tip := pinnedTip + if tip == "" { + head, err := git.headSHA(ctx) + if err != nil { + return gitStateSidecar{}, err + } + tip = head + } + + dirty, err := git.hasUncommittedChanges(ctx) + if err != nil { + return gitStateSidecar{}, err + } + + return gitStateSidecar{ + SchemaVersion: snapshotStateSchemaVersion, + PackagingMode: packagingMode, + BaseCommit: nilIfEmpty(git.mergeBaseWithUpstream(ctx, defaultRemoteName)), + TipCommit: nilIfEmpty(tip), + Branch: nilIfEmpty(git.currentBranch(ctx)), + RepoURL: nilIfEmpty(git.remoteURL(ctx, defaultRemoteName)), + Dirty: dirty, + DiffStatus: diffStatusClean, + DiffPath: nil, + GeneratedAtUTC: now.UTC().Format("2006-01-02T15:04:05.000000") + "Z", + }, nil +} + +// marshal renders the sidecar as indented JSON (matching Python's json.dump indent=2). +func (s gitStateSidecar) marshal() ([]byte, error) { + return json.MarshalIndent(s, "", " ") +} + +// captureDirtyDiff runs `git diff HEAD` over the repo subtree, returning a diff_status +// and the diff bytes (non-nil only when captured): clean (no changes or diff failed), +// captured (under the cap), size_exceeded, or timeout. +func captureDirtyDiff(ctx context.Context, git gitRepo, sizeCapBytes int, timeout time.Duration) (string, []byte) { + diffCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + out, err := git.runBytes(diffCtx, "diff", "HEAD", "--", ".") + if err != nil { + if errors.Is(diffCtx.Err(), context.DeadlineExceeded) { + return diffStatusTimeout, nil + } + return diffStatusClean, nil + } + if len(out) == 0 { + return diffStatusClean, nil + } + if len(out) > sizeCapBytes { + return diffStatusSizeExceeded, nil + } + return diffStatusCaptured, out +} diff --git a/experimental/air/cmd/snapshot_git_test.go b/experimental/air/cmd/snapshot_git_test.go new file mode 100644 index 0000000000..cf1a821b1d --- /dev/null +++ b/experimental/air/cmd/snapshot_git_test.go @@ -0,0 +1,296 @@ +package aircmd + +import ( + "encoding/json" + "os" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newTestRepo initializes a git repo in a temp dir with a deterministic identity +// and returns its path. Tests build up real commits/branches/dirty states on top, +// mirroring the Python git_state tests (which drive real repos, not a fake). +func newTestRepo(t *testing.T) string { + t.Helper() + dir := t.TempDir() + runGit(t, dir, "init", "-q", "-b", "main") + // Deterministic identity so commits succeed in a bare CI environment. + runGit(t, dir, "config", "user.email", "test@example.test") + runGit(t, dir, "config", "user.name", "Test") + return dir +} + +// runGit runs a git command in dir and fails the test on error. +func runGit(t *testing.T, dir string, args ...string) { + t.Helper() + cmd := exec.Command("git", append([]string{"-C", dir}, args...)...) + out, err := cmd.CombinedOutput() + require.NoError(t, err, "git %v: %s", args, out) +} + +// writeRepoFile writes a file at a repo-relative path, creating parent dirs. +func writeRepoFile(t *testing.T, repo, rel, content string) { + t.Helper() + full := filepath.Join(repo, rel) + require.NoError(t, os.MkdirAll(filepath.Dir(full), 0o755)) + require.NoError(t, os.WriteFile(full, []byte(content), 0o600)) +} + +// commitAll stages everything and commits, returning the new HEAD SHA. +func commitAll(t *testing.T, repo, msg string) string { + t.Helper() + runGit(t, repo, "add", "-A") + runGit(t, repo, "commit", "-q", "-m", msg) + sha, err := newGitRepo(repo).headSHA(t.Context()) + require.NoError(t, err) + return sha +} + +func TestGitRepo_IsRepository(t *testing.T) { + ctx := t.Context() + + repo := newTestRepo(t) + assert.True(t, newGitRepo(repo).isRepository(ctx)) + + // A subdirectory of a repo is still inside the work tree. + writeRepoFile(t, repo, "sub/x.txt", "hi") + assert.True(t, newGitRepo(filepath.Join(repo, "sub")).isRepository(ctx)) + + // A plain temp dir with no repo is not. + assert.False(t, newGitRepo(t.TempDir()).isRepository(ctx)) +} + +func TestGitRepo_HeadSHA(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + sha := commitAll(t, repo, "init") + + got, err := newGitRepo(repo).headSHA(ctx) + require.NoError(t, err) + assert.Equal(t, sha, got) + assert.Len(t, got, 40) +} + +func TestGitRepo_HasUncommittedChanges(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + + // Clean tree. + dirty, err := newGitRepo(repo).hasUncommittedChanges(ctx) + require.NoError(t, err) + assert.False(t, dirty) + + // Unstaged modification. + writeRepoFile(t, repo, "a.txt", "2") + dirty, err = newGitRepo(repo).hasUncommittedChanges(ctx) + require.NoError(t, err) + assert.True(t, dirty) +} + +func TestGitRepo_HasUncommittedChangesInPaths(t *testing.T) { + ctx := t.Context() + + repo := newTestRepo(t) + writeRepoFile(t, repo, "src/model.py", "1") + writeRepoFile(t, repo, "other/x.py", "1") + commitAll(t, repo, "init") + g := newGitRepo(repo) + + // No paths: no changes, and git is never consulted. + dirty, err := g.hasUncommittedChangesInPaths(ctx, nil) + require.NoError(t, err) + assert.False(t, dirty) + + // A change outside the included paths is ignored. + writeRepoFile(t, repo, "other/x.py", "2") + dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"src", "configs"}) + require.NoError(t, err) + assert.False(t, dirty) + + // A change inside an included path is reported. + writeRepoFile(t, repo, "src/model.py", "2") + dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"src", "configs"}) + require.NoError(t, err) + assert.True(t, dirty) + + // Trailing slashes on include paths are trimmed for the pathspec. + dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"other/"}) + require.NoError(t, err) + assert.True(t, dirty) +} + +func TestGitRepo_HasUncommittedChangesInPaths_Rename(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "src/old.py", "content") + commitAll(t, repo, "init") + + // A rename within an included path counts as a change, however git classifies + // it (rename vs delete+add); we only assert the boolean. + runGit(t, repo, "mv", "src/old.py", "src/new.py") + dirty, err := newGitRepo(repo).hasUncommittedChangesInPaths(ctx, []string{"src"}) + require.NoError(t, err) + assert.True(t, dirty) +} + +func TestGitRepo_ResolveLocalBranchSHA(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + mainSHA := commitAll(t, repo, "init") + + // A second branch at its own commit. + runGit(t, repo, "checkout", "-q", "-b", "feature") + writeRepoFile(t, repo, "b.txt", "2") + featSHA := commitAll(t, repo, "feature work") + g := newGitRepo(repo) + + got, err := g.resolveLocalBranchSHA(ctx, "main") + require.NoError(t, err) + assert.Equal(t, mainSHA, got) + + got, err = g.resolveLocalBranchSHA(ctx, "feature") + require.NoError(t, err) + assert.Equal(t, featSHA, got) + + // A branch that does not exist locally errors (no remote is contacted). + _, err = g.resolveLocalBranchSHA(ctx, "nope") + require.Error(t, err) + assert.Contains(t, err.Error(), "resolve local branch") +} + +func TestGitRepo_CommitExistsLocally(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + sha := commitAll(t, repo, "init") + g := newGitRepo(repo) + + assert.True(t, g.commitExistsLocally(ctx, sha)) + assert.False(t, g.commitExistsLocally(ctx, "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef")) +} + +func TestGitRepo_ValidateIncludePathsExist(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "src/model.py", "1") + writeRepoFile(t, repo, "configs/train.yaml", "x") + writeRepoFile(t, repo, "train.py", "print()") + sha := commitAll(t, repo, "init") + g := newGitRepo(repo) + + // Both directory and file include_paths are accepted (ls-tree without -d). + require.NoError(t, g.validateIncludePathsExist(ctx, sha, []string{"src", "configs", "train.py"})) + + err := g.validateIncludePathsExist(ctx, sha, []string{"src", "missing"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "missing") + assert.Contains(t, err.Error(), sha[:8]) +} + +func TestBuildGitStateSidecar_PlainTarClean(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + head := commitAll(t, repo, "init") + + sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) + require.NoError(t, err) + assert.Equal(t, snapshotStateSchemaVersion, sc.SchemaVersion) + assert.Equal(t, packagingModePlainTar, sc.PackagingMode) + require.NotNil(t, sc.TipCommit) + assert.Equal(t, head, *sc.TipCommit) + assert.False(t, sc.Dirty) + assert.Equal(t, diffStatusClean, sc.DiffStatus) + assert.Nil(t, sc.DiffPath) + // No remote in a bare test repo → base_commit and repo_url are null. + assert.Nil(t, sc.BaseCommit) + assert.Nil(t, sc.RepoURL) + require.NotNil(t, sc.Branch) + assert.Equal(t, "main", *sc.Branch) + assert.Equal(t, "2026-07-10T12:00:00.000000Z", sc.GeneratedAtUTC) +} + +func TestBuildGitStateSidecar_GitArchivePinsTip(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + first := commitAll(t, repo, "init") + // Advance HEAD; the pinned tip must win over HEAD. + writeRepoFile(t, repo, "b.txt", "2") + commitAll(t, repo, "second") + + sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModeGitArchive, first, fixedNow) + require.NoError(t, err) + require.NotNil(t, sc.TipCommit) + assert.Equal(t, first, *sc.TipCommit) + assert.Equal(t, packagingModeGitArchive, sc.PackagingMode) +} + +func TestBuildGitStateSidecar_Dirty(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + writeRepoFile(t, repo, "a.txt", "2") // uncommitted + + sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) + require.NoError(t, err) + assert.True(t, sc.Dirty) +} + +func TestGitStateSidecar_MarshalNullsAbsentFields(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + + sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) + require.NoError(t, err) + data, err := sc.marshal() + require.NoError(t, err) + + // Absent fields serialize as JSON null (not "" or omitted), matching Python. + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + require.Contains(t, raw, "base_commit") + assert.Nil(t, raw["base_commit"]) + require.Contains(t, raw, "repo_url") + assert.Nil(t, raw["repo_url"]) + require.Contains(t, raw, "diff_path") + assert.Nil(t, raw["diff_path"]) + assert.EqualValues(t, 1, raw["schema_version"]) +} + +func TestCaptureDirtyDiff(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "one\n") + commitAll(t, repo, "init") + + // Clean tree → no diff. + status, diff := captureDirtyDiff(ctx, newGitRepo(repo), dirtyDiffSizeCapBytes, dirtyDiffTimeout) + assert.Equal(t, diffStatusClean, status) + assert.Nil(t, diff) + + // Dirty tree → captured, and the diff mentions the changed file. + writeRepoFile(t, repo, "a.txt", "two\n") + status, diff = captureDirtyDiff(ctx, newGitRepo(repo), dirtyDiffSizeCapBytes, dirtyDiffTimeout) + assert.Equal(t, diffStatusCaptured, status) + assert.Contains(t, string(diff), "a.txt") + + // A tiny size cap forces size_exceeded and drops the bytes. + status, diff = captureDirtyDiff(ctx, newGitRepo(repo), 1, dirtyDiffTimeout) + assert.Equal(t, diffStatusSizeExceeded, status) + assert.Nil(t, diff) +} + +var fixedNow = time.Date(2026, 7, 10, 12, 0, 0, 0, time.UTC) diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index 38fdb76c0b..dd043fdfa5 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -10,10 +10,28 @@ import ( "strings" ) -// Tar builder ported from cli/utils/snapshot.py. It shells out to `tar` for parity -// and to reuse tar's symlink, gitignore, and AppleDouble handling. The tarball's -// top-level dir name is load-bearing — the remote entry_script extracts to -// /databricks/code_source/ — so the `-C parent dir` form preserves it. +// Tar builders ported from cli/utils/snapshot.py. Both shell out (git archive / tar) +// for parity and to reuse git's/tar's symlink, gitignore, and AppleDouble handling. +// The tarball's top-level dir name is load-bearing — the remote entry_script extracts +// to /databricks/code_source/ — so the --prefix / `-C parent dir` forms preserve it. + +// createGitArchiveSnapshot writes a gzipped tar of commitSHA to outputTarball via +// `git archive`, with every entry prefixed by directoryName/. When includePaths is +// set, only those paths are archived. +func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outputTarball, directoryName string, includePaths []string) error { + args := []string{ + "archive", + "--format=tar.gz", + "--prefix=" + directoryName + "/", + "-o", outputTarball, + commitSHA, + } + args = append(args, includePaths...) + if _, err := git.run(ctx, args...); err != nil { + return fmt.Errorf("failed to create git archive: %w", err) + } + return nil +} // createPlainTarball writes a gzipped tar of repoPath's working tree to // outputTarball via `tar`. The archive preserves repoPath's directory name as the diff --git a/experimental/air/cmd/snapshot_package_test.go b/experimental/air/cmd/snapshot_package_test.go index a3cdaf72b0..22561f4e2c 100644 --- a/experimental/air/cmd/snapshot_package_test.go +++ b/experimental/air/cmd/snapshot_package_test.go @@ -6,6 +6,7 @@ import ( "os" "path/filepath" "slices" + "strings" "testing" "github.com/stretchr/testify/assert" @@ -36,12 +37,44 @@ func tarballEntries(t *testing.T, path string) []string { return names } -// writeRepoFile writes a file at a dir-relative path, creating parent dirs. -func writeRepoFile(t *testing.T, dir, rel, content string) { - t.Helper() - full := filepath.Join(dir, rel) - require.NoError(t, os.MkdirAll(filepath.Dir(full), 0o755)) - require.NoError(t, os.WriteFile(full, []byte(content), 0o600)) +func TestCreateGitArchiveSnapshot(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + writeRepoFile(t, repo, "src/model.py", "print()") + sha := commitAll(t, repo, "init") + + out := filepath.Join(t.TempDir(), "snap.tar.gz") + dirName := filepath.Base(repo) + require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, nil)) + + entries := tarballEntries(t, out) + // Every real entry is prefixed with the directory name. git archive also emits a + // `pax_global_header` pseudo-entry (no prefix) that tar ignores on extraction. + assert.Contains(t, entries, dirName+"/a.txt") + assert.Contains(t, entries, dirName+"/src/model.py") + for _, e := range entries { + if e == "pax_global_header" { + continue + } + assert.True(t, strings.HasPrefix(e, dirName+"/"), "entry %q lacks prefix", e) + } +} + +func TestCreateGitArchiveSnapshot_IncludePaths(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + writeRepoFile(t, repo, "src/model.py", "print()") + sha := commitAll(t, repo, "init") + + out := filepath.Join(t.TempDir(), "snap.tar.gz") + dirName := filepath.Base(repo) + require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, []string{"src"})) + + entries := tarballEntries(t, out) + assert.Contains(t, entries, dirName+"/src/model.py") + assert.NotContains(t, entries, dirName+"/a.txt") } func TestCreatePlainTarball(t *testing.T) { diff --git a/experimental/air/cmd/snapshot_resolve.go b/experimental/air/cmd/snapshot_resolve.go new file mode 100644 index 0000000000..1146b641b9 --- /dev/null +++ b/experimental/air/cmd/snapshot_resolve.go @@ -0,0 +1,120 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" +) + +// This file ports the mode/ref resolution from the Python CLI's cli_entrypoint +// snapshot block (the if/elif at lines ~1541–1722), local-only. The remote-fetch +// branches are dropped: a git ref must resolve to a commit already present +// locally (git.remote is rejected at validation — see gitRef.validate). + +// snapshotMode is how the snapshot tarball is produced. +type snapshotMode int + +const ( + // modeGitArchive packages a pinned commit via `git archive`. The commit is + // deterministic, so the tarball is cacheable by (commit, include_paths). + modeGitArchive snapshotMode = iota + // modePlainTar packages the working tree (including uncommitted changes) via + // `tar`. Not cacheable — working-tree content isn't pinned to a SHA. + modePlainTar +) + +// snapshotPlan is the outcome of resolving how to package a snapshot: the mode, +// the commit SHA to archive (git_archive only; empty for plain_tar), and whether +// the working tree under the snapshot root has uncommitted changes. +type snapshotPlan struct { + mode snapshotMode + commitSHA string + hasUncommit bool + isGitRepo bool + includePaths []string +} + +// resolveSnapshotPlan decides how to package the snapshot (local-only): +// - git.commit → pin the SHA (must exist locally) → git_archive. +// - git.branch → the branch's local HEAD SHA → git_archive. +// - no ref / non-git dir → the working tree → plain_tar (no caching). +// +// The dirty check runs at most once (git status is O(working tree)) and is threaded +// into the plan. Dirty + git.branch is an error: the committed HEAD wouldn't include +// the uncommitted changes. +func resolveSnapshotPlan(ctx context.Context, git gitRepo, ref *gitRef, includePaths []string) (snapshotPlan, error) { + plan := snapshotPlan{includePaths: includePaths} + plan.isGitRepo = git.isRepository(ctx) + + // Detect uncommitted changes once. When include_paths is set, only changes + // under those paths can land in the snapshot, so scope the check to them — + // both more correct and cheaper than scanning the whole repo. + if plan.isGitRepo { + var err error + if len(includePaths) > 0 { + plan.hasUncommit, err = git.hasUncommittedChangesInPaths(ctx, includePaths) + } else { + plan.hasUncommit, err = git.hasUncommittedChanges(ctx) + } + if err != nil { + return snapshotPlan{}, err + } + } + + // Non-git directory: plain tar, no ref allowed. gitRef.validate already rejects + // git.* on a non-git dir at load time, but guard here too since this function + // is the single decision point. + if !plan.isGitRepo { + if ref != nil { + return snapshotPlan{}, fmt.Errorf("git.* is set but %s is not a git repository", git.path) + } + plan.mode = modePlainTar + return plan, nil + } + + // git repo, no ref: package the working tree as plain tar (uncommitted changes + // included). Provenance is captured separately via the git_state sidecar. + if ref == nil { + plan.mode = modePlainTar + return plan, nil + } + + switch { + case ref.Commit != nil: + // git.commit pins a committed SHA; local uncommitted changes are irrelevant + // and won't be included. The commit must exist locally — no remote fetch. + commit := *ref.Commit + if !git.commitExistsLocally(ctx, commit) { + return snapshotPlan{}, fmt.Errorf("commit %q does not exist locally; fetch it (e.g. `git fetch`) before submitting — the snapshot archives your local copy and does not fetch from a remote", commit) + } + plan.mode = modeGitArchive + plan.commitSHA = commit + + case ref.Branch != nil: + // git.branch deploys the branch's local HEAD. A dirty tree here is an error: + // the committed HEAD wouldn't include the uncommitted changes. + if plan.hasUncommit { + return snapshotPlan{}, fmt.Errorf("uncommitted changes under %s would not be included: git.branch deploys the committed HEAD of %q. Commit your changes, or use git.commit to pin a specific revision", git.path, *ref.Branch) + } + sha, err := git.resolveLocalBranchSHA(ctx, *ref.Branch) + if err != nil { + return snapshotPlan{}, err + } + plan.mode = modeGitArchive + plan.commitSHA = sha + + default: + // gitRef.validate guarantees exactly one of branch/commit is set. + return snapshotPlan{}, errors.New("git: must specify either 'branch' or 'commit'") + } + + // For git_archive with include_paths, verify each path exists at the resolved + // commit so a typo fails fast rather than producing an empty subtree. + if len(includePaths) > 0 { + if err := git.validateIncludePathsExist(ctx, plan.commitSHA, includePaths); err != nil { + return snapshotPlan{}, err + } + } + + return plan, nil +} diff --git a/experimental/air/cmd/snapshot_resolve_test.go b/experimental/air/cmd/snapshot_resolve_test.go new file mode 100644 index 0000000000..c8c946f839 --- /dev/null +++ b/experimental/air/cmd/snapshot_resolve_test.go @@ -0,0 +1,114 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestResolveSnapshotPlan_Commit(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + sha := commitAll(t, repo, "init") + + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, nil) + require.NoError(t, err) + assert.Equal(t, modeGitArchive, plan.mode) + assert.Equal(t, sha, plan.commitSHA) + assert.True(t, plan.isGitRepo) + + // A commit pin is valid even with a dirty tree: local changes are irrelevant. + writeRepoFile(t, repo, "a.txt", "2") + plan, err = resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, nil) + require.NoError(t, err) + assert.Equal(t, modeGitArchive, plan.mode) + assert.True(t, plan.hasUncommit) +} + +func TestResolveSnapshotPlan_CommitNotLocal(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + + absent := "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef" + _, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(absent)}, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not exist locally") +} + +func TestResolveSnapshotPlan_BranchLocalHead(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + mainSHA := commitAll(t, repo, "init") + + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Branch: new("main")}, nil) + require.NoError(t, err) + assert.Equal(t, modeGitArchive, plan.mode) + assert.Equal(t, mainSHA, plan.commitSHA) +} + +func TestResolveSnapshotPlan_BranchDirtyIsError(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + writeRepoFile(t, repo, "a.txt", "2") // uncommitted + + _, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Branch: new("main")}, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "uncommitted changes") + assert.Contains(t, err.Error(), "git.commit") +} + +func TestResolveSnapshotPlan_NoRefPlainTar(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + writeRepoFile(t, repo, "a.txt", "2") // dirty tree is fine for plain tar + + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), nil, nil) + require.NoError(t, err) + assert.Equal(t, modePlainTar, plan.mode) + assert.Empty(t, plan.commitSHA) + assert.True(t, plan.isGitRepo) + assert.True(t, plan.hasUncommit) +} + +func TestResolveSnapshotPlan_NonGitDir(t *testing.T) { + ctx := t.Context() + dir := t.TempDir() + + plan, err := resolveSnapshotPlan(ctx, newGitRepo(dir), nil, nil) + require.NoError(t, err) + assert.Equal(t, modePlainTar, plan.mode) + assert.False(t, plan.isGitRepo) + + // A git ref on a non-git directory is an error. + _, err = resolveSnapshotPlan(ctx, newGitRepo(dir), &gitRef{Branch: new("main")}, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "not a git repository") +} + +func TestResolveSnapshotPlan_IncludePaths(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "src/model.py", "1") + writeRepoFile(t, repo, "configs/train.yaml", "x") + sha := commitAll(t, repo, "init") + + // All include paths exist at the commit. + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, []string{"src", "configs"}) + require.NoError(t, err) + assert.Equal(t, modeGitArchive, plan.mode) + assert.Equal(t, []string{"src", "configs"}, plan.includePaths) + + // A missing include path fails fast. + _, err = resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, []string{"src", "missing"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "missing") +}