Fix: render execution_time through a per-test generator instead of patching the shared dialect

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(