Fix: render execution_time through a per-test generator, not a shared patch

ModelTest froze time by patching the dialect's generator TRANSFORMS and
SQLGlot's cached dispatch table, which are process-global, under a lock
that only covers the model render. setUp (fixture CREATE VIEW), tearDown
and the CTE checks render outside that lock, so with concurrent_tasks > 1
another test could observe the dispatch table mid-restore and fail with
"Unsupported expression type Create", or silently render CURRENT_* at the
other test's frozen time.

Build the frozen-time transforms into a per-test generator subclass and
render SQL model tests through it; keep the lock for time_machine only.
Python model tests still patch the shared dialect under the lock because
their SQL goes through the engine adapter. CTE checks now render with the
frozen time too, which they previously did not.

Signed-off-by: Niklas Dohmen <niklas@enam.co>

Author: niklasdohmen

sqlmesh/core/test/definition.py

@@ -34,6 +34,7 @@
     import pandas as pd
 
     from sqlglot.dialects.dialect import DialectType
+    from sqlglot.generator import Generator
 
     Row = t.Dict[str, t.Any]
 
@@ -48,6 +49,54 @@
 }
 
 
+_FROZEN_TIME_GENERATORS: t.Dict[
+    t.Tuple[t.Type[Generator], str, t.Optional[str]], t.Type[Generator]
+] = {}
+_FROZEN_TIME_GENERATORS_LOCK = threading.Lock()
+
+
+def _frozen_time_generator_class(
+    generator_class: t.Type[Generator], execution_time: str, dialect: t.Optional[str]
+) -> t.Type[Generator]:
+    """Returns a subclass of `generator_class` whose CURRENT_* transforms render `execution_time`.
+
+    SQLGlot caches one dispatch table per generator class, so a subclass gets its own table and
+    the dialect's shared generator class is never modified. Subclasses are cached per
+    (generator, execution time, dialect) so tests that share an execution time share a table.
+    """
+    key = (generator_class, execution_time, dialect)
+    with _FROZEN_TIME_GENERATORS_LOCK:
+        klass = _FROZEN_TIME_GENERATORS.get(key)
+        if klass is None:
+            exec_time = exp.Literal.string(execution_time)
+            klass = t.cast(
+                t.Type["Generator"],
+                type(
+                    f"{generator_class.__name__}FrozenTime",
+                    (generator_class,),
+                    {
+                        "TRANSFORMS": {
+                            **generator_class.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)
+                            ),
+                        }
+                    },
+                ),
+            )
+            _FROZEN_TIME_GENERATORS[key] = klass
+        return klass
+
+
 class ModelTest(unittest.TestCase):
     __test__ = False
 
@@ -116,31 +165,22 @@ def __init__(
         )
         self._qualified_fixture_schema = schema_(self._fixture_schema, self._fixture_catalog)
 
-        self._transforms = self._test_adapter_dialect.generator_class.TRANSFORMS
         self._execution_time = str(self.body.get("vars", {}).get("execution_time") or "")
 
         if self._execution_time:
             # Normalizes the execution time by converting it into UTC timezone
             self._execution_time = str(to_datetime(self._execution_time))
 
-        # When execution_time is set, we mock the CURRENT_* SQL expressions so they always return it
+        # When execution_time is set, the CURRENT_* SQL expressions must render as that time. The
+        # overrides live on a per-test generator subclass rather than on the dialect's shared
+        # generator class: SQLGlot caches one dispatch table per generator class, so patching the
+        # shared one is visible to every other thread that renders SQL (e.g. a concurrent test
+        # creating its fixture views) and races with the patch's restoration.
+        self._generator_class = self._test_adapter_dialect.generator_class
         if self._execution_time:
-            exec_time = exp.Literal.string(self._execution_time)
-            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)
-                ),
-            }
+            self._generator_class = _frozen_time_generator_class(
+                self._generator_class, self._execution_time, dialect
+            )
 
         super().__init__()
 
@@ -603,15 +643,18 @@ def _normalize_column_name(self, name: str) -> str:
         return normalized_name
 
     @contextmanager
-    def _concurrent_render_context(self) -> t.Iterator[None]:
+    def _concurrent_render_context(self, patch_shared_dialect: bool = False) -> 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
+        This is needed in case `execution_time` is set, as we'd then have to freeze time through
+        `time_machine`, which is not thread safe.
+
+        SQL model tests render through `self._generator_class`, so the shared dialect is never
+        modified. Python model tests may run arbitrary SQL through the engine adapter, whose
+        generator cannot be swapped per test, so they additionally patch the shared generator's
+        transforms (`patch_shared_dialect=True`) while holding the lock.
         """
         import time_machine
-        from sqlglot.generator import _DISPATCH_CACHE
 
         lock_ctx: AbstractContextManager = (
             self.CONCURRENT_RENDER_LOCK if self.concurrency else nullcontext()
@@ -621,19 +664,30 @@ def _concurrent_render_context(self) -> t.Iterator[None]:
         dispatch_patch_ctx: AbstractContextManager = nullcontext()
 
         if self._execution_time:
-            generator_class = self._test_adapter_dialect.generator_class
             time_ctx = time_machine.travel(self._execution_time, tick=False)
-            dialect_patch_ctx = patch.dict(generator_class.TRANSFORMS, self._transforms)
+
+        if self._execution_time and patch_shared_dialect:
+            from sqlglot.generator import _DISPATCH_CACHE
+
+            generator_class = self._test_adapter_dialect.generator_class
+            transforms = self._generator_class.TRANSFORMS
+            dialect_patch_ctx = patch.dict(generator_class.TRANSFORMS, transforms)
 
             # sqlglot caches a dispatch table per generator class, so we need to patch
             # it as well to ensure the overridden transforms are actually used
             dispatch = _DISPATCH_CACHE.get(generator_class)
             if dispatch is not None:
-                dispatch_patch_ctx = patch.dict(dispatch, self._transforms)
+                dispatch_patch_ctx = patch.dict(dispatch, transforms)
 
         with lock_ctx, time_ctx, dialect_patch_ctx, dispatch_patch_ctx:
             yield
 
+    def _generate_sql(self, expression: exp.Expr) -> str:
+        """Generates SQL for the testing engine, rendering CURRENT_* at `execution_time` when set."""
+        return self._generator_class(
+            dialect=self._test_adapter_dialect, pretty=self.engine_adapter._pretty_sql
+        ).generate(expression)
+
     def _execute(self, query: exp.Query | str) -> pd.DataFrame:
         """Executes the given query using the testing engine adapter and returns a DataFrame."""
         return self.engine_adapter.fetchdf(query)
@@ -701,9 +755,7 @@ def test_ctes(self, ctes: t.Dict[str, exp.Expr], recursive: bool = False) -> Non
                 with self._concurrent_render_context():
                     # Similar to the model's query, we render the CTE query under the locked context
                     # so that the execution (fetchdf) can continue concurrently between the threads
-                    sql = cte_query.sql(
-                        self._test_adapter_dialect, pretty=self.engine_adapter._pretty_sql
-                    )
+                    sql = self._generate_sql(cte_query)
 
                 actual = self._execute(sql)
                 expected = self._create_df(values, columns=cte_query.named_selects, partial=partial)
@@ -715,7 +767,7 @@ def runTest(self) -> None:
             # Render the model's query and generate the SQL under the locked context so that
             # execution (fetchdf) can continue concurrently between the threads
             query = self._render_model_query()
-            sql = query.sql(self._test_adapter_dialect, pretty=self.engine_adapter._pretty_sql)
+            sql = self._generate_sql(query)
 
         with_clause = query.args.get("with_")
 
@@ -820,7 +872,7 @@ def _execute_model(self) -> pd.DataFrame:
         """Executes the python model and returns a DataFrame."""
         import pandas as pd
 
-        with self._concurrent_render_context():
+        with self._concurrent_render_context(patch_shared_dialect=True):
             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, variables=variables, **time_kwargs))

tests/core/test_test.py

@@ -3,6 +3,7 @@
 import datetime
 import typing as t
 import io
+import threading
 from pathlib import Path
 import unittest
 from unittest.mock import call, patch
@@ -1360,6 +1361,68 @@ def test_nested_data_types(sushi_context: Context) -> None:
     )
 
 
+def test_freeze_time_does_not_patch_shared_generator() -> None:
+    from sqlglot.dialects.dialect import Dialect
+    from sqlglot.generator import _DISPATCH_CACHE
+
+    duckdb = Dialect.get_or_raise("duckdb")
+    exp.select("1").sql("duckdb")  # make sure the shared dispatch table exists
+    shared_transforms = dict(duckdb.generator_class.TRANSFORMS)
+    shared_dispatch = dict(_DISPATCH_CACHE[duckdb.generator_class])
+
+    test = _create_test(
+        body=load_yaml(
+            """
+test_foo:
+  model: xyz
+  outputs:
+    query:
+      - cur_date: 2023-01-01
+  vars:
+    execution_time: "2023-01-01 12:05:03+00:00"
+            """
+        ),
+        test_name="test_foo",
+        model=_create_model("SELECT CURRENT_DATE AS cur_date"),
+        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+    )
+    test.concurrency = True
+
+    rendered: t.List[str] = []
+    errors: t.List[BaseException] = []
+
+    def render_ddl_through_shared_dialect() -> None:
+        # Mimics another test creating its fixture views while this test's frozen render
+        # context is active; both go through the dialect's shared generator class.
+        try:
+            for _ in range(200):
+                rendered.append(
+                    exp.Create(
+                        this=exp.to_table("s.v"), kind="VIEW", expression=exp.select("1")
+                    ).sql("duckdb")
+                )
+                rendered.append(exp.CurrentDate().sql("duckdb"))
+        except BaseException as e:  # pragma: no cover
+            errors.append(e)
+
+    with test._concurrent_render_context():
+        other = threading.Thread(target=render_ddl_through_shared_dialect)
+        other.start()
+        other.join()
+
+        # this test renders the frozen time...
+        assert test._generate_sql(exp.CurrentDate()) == "CAST('2023-01-01 12:05:03+00:00' AS DATE)"
+        # ...while the shared dialect is untouched, even inside the frozen context
+        assert exp.CurrentDate().sql("duckdb") == "CURRENT_DATE"
+
+    assert not errors
+    assert set(rendered) == {"CREATE VIEW s.v AS SELECT 1", "CURRENT_DATE"}
+    assert duckdb.generator_class.TRANSFORMS == shared_transforms
+    assert _DISPATCH_CACHE[duckdb.generator_class] == shared_dispatch
+
+    _check_successful_or_raise(test.run())
+
+
 def test_freeze_time(mocker: MockerFixture) -> None:
     mocker.patch("sqlmesh.core.test.definition.random_id", return_value="jzngz56a")
     test = _create_test(