Add libs/atomicfile and use it for in-place file saves

.nextchanges/cli/atomic-file-writes.md

@@ -0,0 +1 @@
+* Write local state, cache, and config files atomically so an interrupted or concurrent write cannot corrupt them. ([#6708](https://github.com/databricks/cli/pull/6708))

bundle/configsync/diff.go

@@ -7,7 +7,6 @@ import (
 	"io"
 	"io/fs"
 	"os"
-	"path/filepath"
 	"strings"
 
 	"github.com/databricks/cli/bundle"
@@ -17,6 +16,7 @@ import (
 	"github.com/databricks/cli/bundle/deployplan"
 	"github.com/databricks/cli/bundle/direct"
 	"github.com/databricks/cli/bundle/direct/dstate"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/cli/libs/dyn/convert"
 	"github.com/databricks/cli/libs/log"
@@ -260,13 +260,7 @@ func ensureSnapshotAvailable(ctx context.Context, b *bundle.Bundle, engine engin
 		return fmt.Errorf("reading snapshot content: %w", err)
 	}
 
-	localStateDir := filepath.Dir(localPathSnapshot)
-	err = os.MkdirAll(localStateDir, 0o700)
-	if err != nil {
-		return fmt.Errorf("creating snapshot directory: %w", err)
-	}
-
-	err = os.WriteFile(localPathSnapshot, content, 0o600)
+	err = atomicfile.Write(localPathSnapshot, content, 0o600, atomicfile.MkDir(0o700))
 	if err != nil {
 		return fmt.Errorf("writing snapshot file: %w", err)
 	}

bundle/configsync/output.go

@@ -6,10 +6,9 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"os"
-	"path/filepath"
 
 	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/telemetry"
 	"github.com/databricks/cli/libs/telemetry/protos"
 )
@@ -91,12 +90,7 @@ func WriteResult(out io.Writer, jsonOutput bool, stats *Stats, files []FileChang
 // SaveFiles writes all file changes to disk.
 func SaveFiles(ctx context.Context, b *bundle.Bundle, files []FileChange) error {
 	for _, file := range files {
-		err := os.MkdirAll(filepath.Dir(file.Path), 0o755)
-		if err != nil {
-			return err
-		}
-
-		err = os.WriteFile(file.Path, []byte(file.ModifiedContent), 0o644)
+		err := atomicfile.Write(file.Path, []byte(file.ModifiedContent), 0o644, atomicfile.MkDir(0o755))
 		if err != nil {
 			return err
 		}

bundle/deploy/state_pull.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/deploy/files"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/log"
@@ -44,37 +45,22 @@ func (s *statePull) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostic
 		return diag.FromErr(err)
 	}
 
-	local, err := os.OpenFile(statePath, os.O_CREATE|os.O_RDWR, 0o600)
-	if err != nil {
-		return diag.FromErr(err)
-	}
-	defer local.Close()
-
 	data := remote.Bytes()
 	err = validateRemoteStateCompatibility(bytes.NewReader(data))
 	if err != nil {
 		return diag.FromErr(err)
 	}
 
-	if !isLocalStateStale(local, bytes.NewReader(data)) {
+	// A missing or unreadable local file counts as stale, so we write the remote copy.
+	localData, _ := os.ReadFile(statePath)
+	if !isLocalStateStale(bytes.NewReader(localData), bytes.NewReader(data)) {
 		log.Infof(ctx, "Local deployment state is the same or newer, ignoring remote state")
 		return nil
 	}
 
-	// Truncating the file before writing
-	err = local.Truncate(0)
-	if err != nil {
-		return diag.FromErr(err)
-	}
-	_, err = local.Seek(0, 0)
-	if err != nil {
-		return diag.FromErr(err)
-	}
-
 	// Write file to disk.
 	log.Infof(ctx, "Writing remote deployment state file to local cache directory")
-	_, err = io.Copy(local, bytes.NewReader(data))
-	if err != nil {
+	if err := atomicfile.Write(statePath, data, 0o600); err != nil {
 		return diag.FromErr(err)
 	}
 

bundle/deploy/state_update.go

@@ -1,17 +1,16 @@
 package deploy
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
-	"io"
 	"io/fs"
 	"os"
 	"time"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/internal/build"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/log"
 	"github.com/google/uuid"
@@ -56,21 +55,14 @@ func (s *stateUpdate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnost
 	if err != nil {
 		return diag.FromErr(err)
 	}
-	// Write the state back to the file.
-	f, err := os.OpenFile(statePath, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0o600)
-	if err != nil {
-		log.Infof(ctx, "Unable to open deployment state file: %s", err)
-		return diag.FromErr(err)
-	}
-	defer f.Close()
-
 	data, err := json.Marshal(state)
 	if err != nil {
 		return diag.FromErr(err)
 	}
 
-	_, err = io.Copy(f, bytes.NewReader(data))
-	if err != nil {
+	// Write the state back to the file.
+	if err := atomicfile.Write(statePath, data, 0o600); err != nil {
+		log.Infof(ctx, "Unable to write deployment state file: %s", err)
 		return diag.FromErr(err)
 	}
 

bundle/direct/dstate/state.go

@@ -19,6 +19,7 @@ import (
 	"github.com/databricks/cli/bundle/deployplan"
 	"github.com/databricks/cli/bundle/statemgmt/resourcestate"
 	"github.com/databricks/cli/internal/build"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/dms"
 	"github.com/databricks/cli/libs/log"
@@ -998,31 +999,7 @@ func (db *DeploymentState) unlockedSave() error {
 		return err
 	}
 
-	dir := filepath.Dir(db.Path)
-	if err := os.MkdirAll(dir, 0o755); err != nil {
-		return fmt.Errorf("failed to create directory %#v: %w", dir, err)
-	}
-
-	// CreateTemp creates the file with mode 0o600, matching the state file.
-	tmp, err := os.CreateTemp(dir, "."+filepath.Base(db.Path)+".tmp-*")
-	if err != nil {
-		return fmt.Errorf("failed to create temp file for %#v: %w", db.Path, err)
-	}
-	tmpPath := tmp.Name()
-	// Cleans up the temp file on failure; a no-op once the rename succeeded.
-	defer os.Remove(tmpPath)
-
-	if _, err := tmp.Write(data); err != nil {
-		tmp.Close()
-		return fmt.Errorf("failed to write %#v: %w", tmpPath, err)
-	}
-
-	// Close before the rename: on Windows the file must not be open for writing.
-	if err := tmp.Close(); err != nil {
-		return fmt.Errorf("failed to close %#v: %w", tmpPath, err)
-	}
-
-	if err := os.Rename(tmpPath, db.Path); err != nil {
+	if err := atomicfile.Write(db.Path, data, 0o600, atomicfile.MkDir(0o755)); err != nil {
 		return fmt.Errorf("failed to save resources state to %#v: %w", db.Path, err)
 	}
 

bundle/statemgmt/state_pull.go

@@ -16,6 +16,7 @@ import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config/engine"
 	"github.com/databricks/cli/bundle/deploy"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/log"
@@ -194,16 +195,7 @@ func PullResourcesState(ctx context.Context, b *bundle.Bundle, alwaysPull Always
 			localStatePath = localPathDirect
 		}
 
-		localStateDir := filepath.Dir(localStatePath)
-
-		err := os.MkdirAll(localStateDir, 0o700)
-		if err != nil {
-			logdiag.LogError(ctx, err)
-			return ctx, winner
-		}
-
-		// TODO: write + rename
-		err = os.WriteFile(localStatePath, winner.Content, 0o600)
+		err := atomicfile.Write(localStatePath, winner.Content, 0o600, atomicfile.MkDir(0o700))
 		if err != nil {
 			logdiag.LogError(ctx, err)
 			return ctx, winner

cmd/genie/conversations.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"time"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/env"
 )
 
@@ -82,25 +83,7 @@ func saveStore(path string, store conversationStore) {
 	if err != nil {
 		return
 	}
-	if err := os.MkdirAll(filepath.Dir(path), conversationDirPerm); err != nil {
-		return
-	}
-	tmp, err := os.CreateTemp(filepath.Dir(path), ".genie-conversations-*.tmp")
-	if err != nil {
-		return
-	}
-	defer os.Remove(tmp.Name())
-	defer tmp.Close()
-	if _, err := tmp.Write(raw); err != nil {
-		return
-	}
-	if err := tmp.Chmod(conversationFilePerm); err != nil {
-		return
-	}
-	if err := tmp.Close(); err != nil {
-		return
-	}
-	_ = os.Rename(tmp.Name(), path)
+	_ = atomicfile.Write(path, raw, conversationFilePerm, atomicfile.MkDir(conversationDirPerm))
 }
 
 // lookupConversationID returns the server conversation id mapped to sessionID on

cmd/labs/localcache/jsonfile.go

@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"time"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/log"
 )
 
@@ -73,19 +74,7 @@ func (r *LocalCache[T]) writeCache(ctx context.Context, data T) (T, error) {
 		return r.zero, fmt.Errorf("json marshal: %w", err)
 	}
 	cacheFile := r.FileName()
-	err = os.WriteFile(cacheFile, raw, userRW)
-	if errors.Is(err, fs.ErrNotExist) {
-		cacheDir := filepath.Dir(cacheFile)
-		err := os.MkdirAll(cacheDir, ownerRWXworldRX)
-		if err != nil {
-			return r.zero, fmt.Errorf("create %s: %w", cacheDir, err)
-		}
-		err = os.WriteFile(cacheFile, raw, userRW)
-		if err != nil {
-			return r.zero, fmt.Errorf("retry save cache: %w", err)
-		}
-		return data, nil
-	} else if err != nil {
+	if err := atomicfile.Write(cacheFile, raw, userRW, atomicfile.MkDir(ownerRWXworldRX)); err != nil {
 		return r.zero, fmt.Errorf("save cache: %w", err)
 	}
 	return data, nil

cmd/labs/project/login.go

@@ -4,9 +4,9 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"os"
 
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg/cfgpickers"
 	"github.com/databricks/cli/libs/log"
@@ -120,5 +120,5 @@ func (lc *loginConfig) save(ctx context.Context) error {
 		return err
 	}
 	log.Debugf(ctx, "Writing auth configuration to: %s", authFile)
-	return os.WriteFile(authFile, raw, ownerRW)
+	return atomicfile.Write(authFile, raw, ownerRW)
 }

cmd/labs/project/project.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/databricks/cli/cmd/labs/github"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
@@ -291,7 +292,7 @@ func (p *Project) writeVersionFile(ctx context.Context, ver string) error {
 		return err
 	}
 	log.Debugf(ctx, "Writing installed version info to: %s", versionFile)
-	return os.WriteFile(versionFile, raw, ownerRW)
+	return atomicfile.Write(versionFile, raw, ownerRW)
 }
 
 // checkUpdates is called before every command of an installed project,

cmd/sandbox/sshconfig.go

@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 )
@@ -119,15 +120,7 @@ func writeManagedConfig(path, content string) error {
 	if existing, err := os.ReadFile(path); err == nil && bytes.Equal(existing, []byte(content)) {
 		return nil
 	}
-	tmp := path + ".tmp"
-	if err := os.WriteFile(tmp, []byte(content), 0o600); err != nil {
-		return fmt.Errorf("writing %s: %w", tmp, err)
-	}
-	if err := os.Rename(tmp, path); err != nil {
-		_ = os.Remove(tmp)
-		return fmt.Errorf("renaming %s to %s: %w", tmp, path, err)
-	}
-	return nil
+	return atomicfile.Write(path, []byte(content), 0o600)
 }
 
 // ensureMainIncludesManaged makes sure ~/.ssh/config begins with an
@@ -162,15 +155,7 @@ func ensureMainIncludesManaged(mainPath, managedPath string) error {
 		buf.Write(existing)
 	}
 
-	tmp := mainPath + ".tmp"
-	if err := os.WriteFile(tmp, buf.Bytes(), 0o600); err != nil {
-		return fmt.Errorf("writing %s: %w", tmp, err)
-	}
-	if err := os.Rename(tmp, mainPath); err != nil {
-		_ = os.Remove(tmp)
-		return fmt.Errorf("renaming %s to %s: %w", tmp, mainPath, err)
-	}
-	return nil
+	return atomicfile.Write(mainPath, buf.Bytes(), 0o600)
 }
 
 // hasOurMarkedBlock reports whether the given config text already has

cmd/sandbox/state.go

@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"slices"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/env"
 )
 
@@ -87,24 +88,12 @@ func saveState(ctx context.Context, state *stateFile) error {
 		return err
 	}
 
-	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
-		return err
-	}
-
 	data, err := json.MarshalIndent(state, "", "  ")
 	if err != nil {
 		return err
 	}
 
-	tmp := path + ".tmp"
-	if err := os.WriteFile(tmp, data, 0o600); err != nil {
-		return fmt.Errorf("writing %s: %w", tmp, err)
-	}
-	if err := os.Rename(tmp, path); err != nil {
-		_ = os.Remove(tmp)
-		return fmt.Errorf("renaming %s to %s: %w", tmp, path, err)
-	}
-	return nil
+	return atomicfile.Write(path, data, 0o600, atomicfile.MkDir(0o700))
 }
 
 func getDefault(ctx context.Context, profile string) string {