diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py index 136de947a3..8326418e9d 100644 --- a/sqlmesh/core/test/definition.py +++ b/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)) diff --git a/tests/core/test_test.py b/tests/core/test_test.py index 6ac20cc2d8..76ecb4def1 100644 --- a/tests/core/test_test.py +++ b/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(