fix: write file cache entries atomically to avoid read races

sqlmesh/utils/cache.py

@@ -2,8 +2,10 @@
 
 import gzip
 import logging
+import os
 import pickle
 import shutil
+import tempfile
 import typing as t
 from pathlib import Path
 
@@ -125,8 +127,25 @@ def put(self, name: str, entry_id: str = "", *, value: T) -> None:
         if not self._path.is_dir():
             raise SQLMeshError(f"Cache path '{self._path}' is not a directory.")
 
-        with gzip.open(self._cache_entry_path(name, entry_id), "wb", compresslevel=1) as fd:
-            pickle.dump(value, fd)
+        # Write to a temporary file and then atomically move it into place. Writing the
+        # entry in place ("wb" truncates the target file first) means that a concurrent
+        # reader (e.g. another pytest-xdist worker) could observe a partially written
+        # entry and fail to unpickle it.
+        tmp_fd, tmp_name = tempfile.mkstemp(dir=self._path, prefix=f"{self._cache_version}__tmp")
+        try:
+            with os.fdopen(tmp_fd, "wb") as raw_fd:
+                with gzip.open(raw_fd, "wb", compresslevel=1) as fd:
+                    pickle.dump(value, fd)
+            try:
+                os.replace(tmp_name, self._cache_entry_path(name, entry_id))
+            except OSError as ex:
+                # Windows os.replace fails if a concurrent reader still has the target file open.
+                logger.warning("Failed to store a cache entry '%s': %s", name, ex)
+        finally:
+            try:
+                os.unlink(tmp_name)
+            except OSError:
+                pass
 
     def exists(self, name: str, entry_id: str = "") -> bool:
         """Returns true if the cache entry with the given name and ID exists, false otherwise.

tests/utils/test_cache.py

@@ -42,6 +42,22 @@ def test_file_cache(tmp_path: Path, mocker: MockerFixture):
     assert "客户数据" in cache._cache_entry_path("客户数据").name
 
 
+def test_file_cache_put_is_atomic(tmp_path: Path, mocker: MockerFixture) -> None:
+    cache: FileCache[_TestEntry] = FileCache(tmp_path)
+
+    old_entry = _TestEntry(value="old")
+    cache.put("test_name", value=old_entry)
+
+    # Simulate os.replace failing, e.g. on Windows when a concurrent reader still has the
+    # target file open. The existing entry must never be truncated / partially overwritten.
+    mocker.patch("sqlmesh.utils.cache.os.replace", side_effect=PermissionError("file in use"))
+    cache.put("test_name", value=_TestEntry(value="new"))
+
+    assert cache.get("test_name") == old_entry
+    # The temporary file should have been cleaned up.
+    assert len(list(tmp_path.glob("*"))) == 1
+
+
 def test_optimized_query_cache(tmp_path: Path, mocker: MockerFixture):
     model = SqlModel(
         name="test_model",