Fix: Concurrent dialect patching in model testing

Author: VaggelisD

sqlmesh/core/test/definition.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import datetime
+import threading
 import typing as t
 import unittest
 from collections import Counter
-from contextlib import AbstractContextManager, nullcontext
+from contextlib import nullcontext, contextmanager, AbstractContextManager
 from itertools import chain
 from pathlib import Path
 from unittest.mock import patch
@@ -46,6 +47,8 @@
 class ModelTest(unittest.TestCase):
     __test__ = False
 
+    CONCURRENT_RENDER_LOCK = threading.Lock()
+
     def __init__(
         self,
         body: t.Dict[str, t.Any],
@@ -57,6 +60,7 @@ def __init__(
         path: Path | None = None,
         preserve_fixtures: bool = False,
         default_catalog: str | None = None,
+        concurrency: bool = False,
     ) -> None:
         """ModelTest encapsulates a unit test for a model.
 
@@ -79,6 +83,7 @@ def __init__(
         self.preserve_fixtures = preserve_fixtures
         self.default_catalog = default_catalog
         self.dialect = dialect
+        self.concurrency = concurrency
 
         self._fixture_table_cache: t.Dict[str, exp.Table] = {}
         self._normalized_column_name_cache: t.Dict[str, str] = {}
@@ -310,6 +315,7 @@ def create_test(
         path: Path | None,
         preserve_fixtures: bool = False,
         default_catalog: str | None = None,
+        concurrency: bool = False,
     ) -> t.Optional[ModelTest]:
         """Create a SqlModelTest or a PythonModelTest.
 
@@ -353,6 +359,7 @@ def create_test(
             path,
             preserve_fixtures,
             default_catalog,
+            concurrency,
         )
 
     def __str__(self) -> str:
@@ -512,10 +519,40 @@ def _normalize_column_name(self, name: str) -> str:
 
         return normalized_name
 
+    @contextmanager
+    def _concurrent_render_context(self) -> t.Iterator[None]:
+        """
+        Context manager that ensures that the tests are executed safely in a concurrent environment.
+        This is needed in case `execution_time` is set, as we'd then have to:
+        - Freeze time through `time_machine` (not thread safe)
+        - Globally patch the SQLGlot dialect so that any date/time nodes are evaluated at the `execution_time` during generation
+        """
+        import time_machine
+
+        lock_ctx: AbstractContextManager = (
+            self.CONCURRENT_RENDER_LOCK
+            if (self.concurrency and self._execution_time)
+            else nullcontext()
+        )
+        time_ctx: AbstractContextManager = nullcontext()
+        dialect_patch_ctx: AbstractContextManager = nullcontext()
+
+        if self._execution_time:
+            time_ctx = time_machine.travel(self._execution_time, tick=False)
+            dialect_patch_ctx = patch.dict(
+                self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms
+            )
+
+        with lock_ctx, time_ctx, dialect_patch_ctx:
+            yield
+
     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)
+
+        with self._concurrent_render_context():
+            sql = query.sql(self._test_adapter_dialect, pretty=self.engine_adapter._pretty_sql)
+
+        return self.engine_adapter.fetchdf(sql)
 
     def _create_df(
         self,
@@ -626,6 +663,7 @@ def __init__(
         path: Path | None = None,
         preserve_fixtures: bool = False,
         default_catalog: str | None = None,
+        concurrency: bool = False,
     ) -> None:
         """PythonModelTest encapsulates a unit test for a Python model.
 
@@ -651,6 +689,7 @@ def __init__(
             path,
             preserve_fixtures,
             default_catalog,
+            concurrency,
         )
 
         self.context = TestExecutionContext(
@@ -674,22 +713,13 @@ def runTest(self) -> None:
 
     def _execute_model(self) -> pd.DataFrame:
         """Executes the python model and returns a DataFrame."""
-        if self._execution_time:
-            import time_machine
-
-            time_ctx: AbstractContextManager = time_machine.travel(self._execution_time, tick=False)
-        else:
-            time_ctx = nullcontext()
+        with self._concurrent_render_context():
+            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))
 
-        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()
+        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

@@ -120,6 +120,9 @@ def run_tests(
         default_catalog_dialect=default_catalog_dialect,
     )
 
+    # Ensure workers are not greater than the number of tests
+    num_workers = min(len(model_test_metadata) or 1, default_test_connection.concurrent_tasks)
+
     def _run_single_test(
         metadata: ModelTestMetadata, engine_adapter: EngineAdapter
     ) -> t.Optional[ModelTextTestResult]:
@@ -132,6 +135,7 @@ def _run_single_test(
             path=metadata.path,
             default_catalog=default_catalog,
             preserve_fixtures=preserve_fixtures,
+            concurrency=num_workers > 1,
         )
 
         if not test:
@@ -159,9 +163,6 @@ def _run_single_test(
 
     test_results = []
 
-    # Ensure workers are not greater than the number of tests
-    num_workers = min(len(model_test_metadata) or 1, default_test_connection.concurrent_tasks)
-
     start_time = time.perf_counter()
     try:
         with ThreadPoolExecutor(max_workers=num_workers) as pool:

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