Merge branch 'main' into feature/DRM/resolve-test-collisions-in-parallel-run

Author: cmgoffena13

.github/workflows/pr.yaml

@@ -457,7 +457,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        dbt-version: ['1.3', '1.4', '1.5', '1.6', '1.7', '1.8', '1.9', '1.10', '1.11']
+        dbt-version: ['1.3', '1.4', '1.5', '1.6', '1.7', '1.8', '1.9', '1.10', '1.11', '1.12']
     steps:
       - uses: actions/checkout@v7
       - name: Set up Python

Makefile

@@ -33,7 +33,7 @@ install-dev-dbt-%:
 	echo "Installing dbt version: $$version"; \
 	cp pyproject.toml pyproject.toml.backup; \
 	$(SED_INPLACE) 's/"pydantic>=2.0.0"/"pydantic"/g' pyproject.toml; \
-	if [ "$$version" = "1.10.0" ] || [ "$$version" = "1.11.0" ]; then \
+	if [ "$$version" = "1.10.0" ] || [ "$$version" = "1.11.0" ] || [ "$$version" = "1.12.0" ]; then \
 		echo "Applying special handling for dbt $$version"; \
 		$(SED_INPLACE) -E 's/"(dbt-core)[^"]*"/"\1~='"$$version"'"/g' pyproject.toml; \
 		$(SED_INPLACE) -E 's/"(dbt-(bigquery|duckdb|snowflake|athena-community|clickhouse|redshift|trino))[^"]*"/"\1"/g' pyproject.toml; \

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",