diff --git a/sqlmesh/utils/cache.py b/sqlmesh/utils/cache.py index e1ff59a4a7..d7df2ef0d2 100644 --- a/sqlmesh/utils/cache.py +++ b/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. diff --git a/tests/utils/test_cache.py b/tests/utils/test_cache.py index e6e041e30a..aa11b6c846 100644 --- a/tests/utils/test_cache.py +++ b/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",