Fix: Concurrent dialect patching in model testing

Author: VaggelisD

sqlmesh/core/test/definition.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import datetime
+import threading
 import typing as t
 import unittest
 from collections import Counter
@@ -57,6 +58,7 @@ def __init__(
         path: Path | None = None,
         preserve_fixtures: bool = False,
         default_catalog: str | None = None,
+        lock: t.Optional[threading.Lock] = None,
     ) -> None:
         """ModelTest encapsulates a unit test for a model.
 
@@ -79,6 +81,7 @@ def __init__(
         self.preserve_fixtures = preserve_fixtures
         self.default_catalog = default_catalog
         self.dialect = dialect
+        self.lock = lock
 
         self._fixture_table_cache: t.Dict[str, exp.Table] = {}
         self._normalized_column_name_cache: t.Dict[str, str] = {}
@@ -102,6 +105,7 @@ def __init__(
         )
         self._qualified_fixture_schema = schema_(self._fixture_schema, self._fixture_catalog)
 
+        self._exec_time_transforms: t.Dict[type[exp.Expression], exp.Expression] = {}
         self._transforms = self._test_adapter_dialect.generator_class.TRANSFORMS
         self._execution_time = str(self.body.get("vars", {}).get("execution_time") or "")
 
@@ -112,20 +116,20 @@ def __init__(
         # When execution_time is set, we mock the CURRENT_* SQL expressions so they always return it
         if self._execution_time:
             exec_time = exp.Literal.string(self._execution_time)
+
+            self._exec_time_transforms = {
+                exp.CurrentDate: exp.cast(exec_time, "date", dialect=dialect),
+                exp.CurrentDatetime: exp.cast(exec_time, "datetime", dialect=dialect),
+                exp.CurrentTime: exp.cast(exec_time, "time", dialect=dialect),
+                exp.CurrentTimestamp: exp.cast(exec_time, "timestamp", dialect=dialect),
+            }
+
             self._transforms = {
                 **self._transforms,
-                exp.CurrentDate: lambda self, _: self.sql(
-                    exp.cast(exec_time, "date", dialect=dialect)
-                ),
-                exp.CurrentDatetime: lambda self, _: self.sql(
-                    exp.cast(exec_time, "datetime", dialect=dialect)
-                ),
-                exp.CurrentTime: lambda self, _: self.sql(
-                    exp.cast(exec_time, "time", dialect=dialect)
-                ),
-                exp.CurrentTimestamp: lambda self, _: self.sql(
-                    exp.cast(exec_time, "timestamp", dialect=dialect)
-                ),
+                **{
+                    key: lambda self, _: self.sql(value)
+                    for key, value in self._exec_time_transforms.items()
+                },
             }
 
         super().__init__()
@@ -310,6 +314,7 @@ def create_test(
         path: Path | None,
         preserve_fixtures: bool = False,
         default_catalog: str | None = None,
+        lock: t.Optional[threading.Lock] = None,
     ) -> t.Optional[ModelTest]:
         """Create a SqlModelTest or a PythonModelTest.
 
@@ -353,6 +358,7 @@ def create_test(
             path,
             preserve_fixtures,
             default_catalog,
+            lock=lock,
         )
 
     def __str__(self) -> str:
@@ -514,8 +520,13 @@ def _normalize_column_name(self, name: str) -> str:
 
     def _execute(self, query: exp.Query) -> pd.DataFrame:
         """Executes the given query using the testing engine adapter and returns a DataFrame."""
-        with patch.dict(self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms):
-            return self.engine_adapter.fetchdf(query)
+
+        def replace_execution_time(expression: exp.Expression) -> exp.Expression:
+            return self._exec_time_transforms.get(type(expression), expression)
+
+        return self.engine_adapter.fetchdf(
+            query.transform(replace_execution_time) if self._execution_time else query
+        )
 
     def _create_df(
         self,
@@ -626,6 +637,7 @@ def __init__(
         path: Path | None = None,
         preserve_fixtures: bool = False,
         default_catalog: str | None = None,
+        lock: t.Optional[threading.Lock] = None,
     ) -> None:
         """PythonModelTest encapsulates a unit test for a Python model.
 
@@ -651,6 +663,7 @@ def __init__(
             path,
             preserve_fixtures,
             default_catalog,
+            lock,
         )
 
         self.context = TestExecutionContext(
@@ -681,15 +694,18 @@ def _execute_model(self) -> pd.DataFrame:
         else:
             time_ctx = nullcontext()
 
-        with patch.dict(self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms):
-            with time_ctx:
-                variables = self.body.get("vars", {}).copy()
-                time_kwargs = {
-                    key: variables.pop(key) for key in TIME_KWARG_KEYS if key in variables
-                }
-                df = next(self.model.render(context=self.context, **time_kwargs, **variables))
-                assert not isinstance(df, exp.Expression)
-                return df if isinstance(df, pd.DataFrame) else df.toPandas()
+        with self.lock or nullcontext():
+            with patch.dict(
+                self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms
+            ):
+                with time_ctx:
+                    variables = self.body.get("vars", {}).copy()
+                    time_kwargs = {
+                        key: variables.pop(key) for key in TIME_KWARG_KEYS if key in variables
+                    }
+                    df = next(self.model.render(context=self.context, **time_kwargs, **variables))
+                    assert not isinstance(df, exp.Expression)
+                    return df if isinstance(df, pd.DataFrame) else df.toPandas()
 
 
 def generate_test(

sqlmesh/core/test/result.py

@@ -100,7 +100,8 @@ def log_test_report(self, test_duration: float) -> None:
         for test_case, failure in failures:
             stream.writeln(unittest.TextTestResult.separator1)
             stream.writeln(f"FAIL: {test_case}")
-            stream.writeln(f"{test_case.shortDescription()}")
+            if test_description := test_case.shortDescription():
+                stream.writeln(test_description)
             stream.writeln(unittest.TextTestResult.separator2)
             stream.writeln(failure)
 

sqlmesh/core/test/runner.py

@@ -132,6 +132,7 @@ def _run_single_test(
             path=metadata.path,
             default_catalog=default_catalog,
             preserve_fixtures=preserve_fixtures,
+            lock=lock if num_workers > 1 else None,
         )
 
         if not test:

tests/core/test_test.py

@@ -2370,3 +2370,74 @@ def test_number_of_tests_found(tmp_path: Path) -> None:
     # Case 3: The "new_test.yaml::test_example_full_model2" should amount to a single subtest
     results = context.test(tests=[f"{test_file}::test_example_full_model2"])
     assert len(results.successes) == 1
+
+
+def test_freeze_time_concurrent(tmp_path: Path) -> None:
+    tests_dir = tmp_path / "tests"
+    tests_dir.mkdir()
+
+    for model_name in ["sql_model", "py_model"]:
+        for i in range(5):
+            test_2019 = tmp_path / "tests" / f"test_2019_{model_name}_{i}.yaml"
+            test_2019.write_text(
+                f"""
+    test_2019_{model_name}_{i}:
+      model: {model_name}
+      vars:
+        execution_time: '2019-12-01'
+      outputs:
+        query:
+          rows:
+            - col_exec_ds_time: '2019-12-01'
+              col_current_date: '2019-12-01'
+              """
+            )
+
+            test_2025 = tmp_path / "tests" / f"test_2025_{model_name}_{i}.yaml"
+            test_2025.write_text(
+                f"""
+    test_2025_{model_name}_{i}:
+      model: {model_name}
+      vars:
+        execution_time: '2025-12-01'
+      outputs:
+        query:
+          rows:
+            - col_exec_ds_time: '2025-12-01'
+              col_current_date: '2025-12-01'
+              """
+            )
+
+    ctx = Context(
+        paths=tmp_path,
+        config=Config(default_test_connection=DuckDBConnectionConfig(concurrent_tasks=8)),
+    )
+
+    @model(
+        "py_model",
+        columns={"col_exec_ds_time": "timestamp_ntz", "col_current_date": "timestamp_ntz"},
+    )
+    def execute(context, start, end, execution_time, **kwargs):
+        datetime_now_utc = datetime.datetime.now(tz=datetime.timezone.utc)
+
+        context.engine_adapter.execute(exp.select("CURRENT_DATE()"))
+        current_date = context.engine_adapter.cursor.fetchone()[0]
+
+        return pd.DataFrame(
+            [{"col_exec_ds_time": datetime_now_utc, "col_current_date": current_date}]
+        )
+
+    python_model = model.get_registry()["py_model"].model(module_path=Path("."), path=Path("."))
+
+    ctx.upsert_model(
+        _create_model(
+            meta="MODEL(NAME sql_model)",
+            query="SELECT @execution_ds::timestamp_ntz AS col_exec_ds_time, current_date()::date AS col_current_date",
+            default_catalog=ctx.default_catalog,
+        )
+    )
+
+    ctx.upsert_model(python_model)
+
+    results = ctx.test()
+    assert len(results.successes) == 20