From fc48cff153a4c7ac943470714035d16410513c61 Mon Sep 17 00:00:00 2001
From: "Denis Bilenko (aider)" <denis.bilenko@databricks.com>
Date: Mon, 3 Mar 2025 21:10:11 +0100
Subject: [PATCH] fix: Implement atomic file write with temporary file in
 PatchWheel

---
 libs/patchwheel/patch.go | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/libs/patchwheel/patch.go b/libs/patchwheel/patch.go
index 40f10bd8f..fd22dd641 100644
--- a/libs/patchwheel/patch.go
+++ b/libs/patchwheel/patch.go
@@ -165,7 +165,16 @@ func PatchWheel(ctx context.Context, path, outputDir string) (string, error) {
 	}
 
 	// Target wheel doesn't exist, proceed with patching
-	// AI TODO: write to a temporary file in the same location. Then do rename atomically. Use .tmp<pid> suffix for temp file. Use defer to ensure .tmp<pid> is removed if this function exits early.
+	// Create a temporary file in the same directory with a unique name
+	tmpFile := outpath + fmt.Sprintf(".tmp%d", os.Getpid())
+	
+	// Ensure the temporary file is removed if we exit early
+	defer func() {
+		if _, statErr := os.Stat(tmpFile); statErr == nil {
+			os.Remove(tmpFile)
+		}
+	}()
+	
 	r, err := zip.OpenReader(path)
 	if err != nil {
 		return "", err
@@ -235,8 +244,8 @@ func PatchWheel(ctx context.Context, path, outputDir string) (string, error) {
 		return "", err
 	}
 
-	// We already calculated the output path earlier
-	outFile, err := os.Create(outpath)
+	// Write to the temporary file first
+	outFile, err := os.Create(tmpFile)
 	if err != nil {
 		return "", err
 	}
@@ -296,6 +305,14 @@ func PatchWheel(ctx context.Context, path, outputDir string) (string, error) {
 	if err := zipw.Close(); err != nil {
 		return "", err
 	}
+	
+	// Close the file before renaming
+	outFile.Close()
+	
+	// Rename the temporary file to the final output path (atomic operation)
+	if err := os.Rename(tmpFile, outpath); err != nil {
+		return "", err
+	}
 
 	return outpath, nil
 }