diff --git a/sqlmesh/core/renderer.py b/sqlmesh/core/renderer.py index a2660a435e..80972dbfe5 100644 --- a/sqlmesh/core/renderer.py +++ b/sqlmesh/core/renderer.py @@ -149,6 +149,7 @@ def _render( table_mapping: t.Optional[t.Dict[str, str]] = None, deployability_index: t.Optional[DeployabilityIndex] = None, runtime_stage: RuntimeStage = RuntimeStage.LOADING, + render_table_mapping: t.Optional[TableMapping] = None, **kwargs: t.Any, ) -> t.List[t.Optional[exp.Expr]]: """Renders a expression, expanding macros with provided kwargs @@ -161,6 +162,7 @@ def _render( table_mapping: Table mapping of physical locations. Takes precedence over snapshot mappings. deployability_index: Determines snapshots that are deployable in the context of this evaluation. runtime_stage: Indicates the current runtime stage, for example if we're still loading the project, etc. + render_table_mapping: Shared snapshot mapping and overrides prepared for this promotion. kwargs: Additional kwargs to pass to the renderer. Returns: @@ -232,6 +234,7 @@ def _resolve_table(table: str | exp.Table) -> str: end=end, execution_time=execution_time, runtime_stage=runtime_stage, + render_table_mapping=render_table_mapping, ), snapshots=snapshots, default_catalog=self._default_catalog, @@ -430,6 +433,7 @@ def _resolve_tables( table_mapping: t.Optional[t.Dict[str, str]] = None, expand: t.Iterable[str] = tuple(), deployability_index: t.Optional[DeployabilityIndex] = None, + render_table_mapping: t.Optional[TableMapping] = None, **kwargs: t.Any, ) -> E: if not snapshots and not table_mapping and not expand: @@ -444,10 +448,14 @@ def _resolve_tables( snapshots = snapshots or {} table_mapping = table_mapping or {} - mapping = { - **self._to_table_mapping(snapshots.values(), deployability_index), - **table_mapping, - } + mapping: t.Dict[str, str] = ( + render_table_mapping + if render_table_mapping is not None + else { + **self._to_table_mapping(snapshots.values(), deployability_index), + **table_mapping, + } + ) expand = set(expand) | { name for name, snapshot in snapshots.items() if snapshot.is_embedded } @@ -477,6 +485,7 @@ def _expand(node: exp.Expr) -> exp.Expr: table_mapping=table_mapping, expand=expand, deployability_index=deployability_index, + render_table_mapping=render_table_mapping, **kwargs, ) if nested_query is not None: @@ -489,6 +498,20 @@ def _expand(node: exp.Expr) -> exp.Expr: expression = expression.transform(_expand, copy=False) # type: ignore + if isinstance(mapping, TableMapping): + # Normalize the shared mapping once, then replace only tables in this expression. + normalized_keys = mapping.normalized_keys(self._dialect) + mapping = { + key: mapping[key] + for table in expression.find_all(exp.Table) + if ( + key := normalized_keys.get( + exp.normalize_table_name(table, dialect=self._dialect) + ) + ) + is not None + } + if mapping: expression = exp.replace_tables( expression, mapping, dialect=self._dialect, copy=False diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py index 9d9ecf7672..4340bcedcf 100644 --- a/sqlmesh/core/snapshot/evaluator.py +++ b/sqlmesh/core/snapshot/evaluator.py @@ -41,6 +41,7 @@ from sqlmesh.core.engine_adapter.shared import InsertOverwriteStrategy, DataObjectType, DataObject from sqlmesh.core.model.meta import GrantsTargetLayer from sqlmesh.core.macros import RuntimeStage +from sqlmesh.core.renderer import TableMapping from sqlmesh.core.model import ( AuditResult, IncrementalUnmanagedKind, @@ -286,6 +287,8 @@ def promote( on_complete: A callback to call on each successfully promoted snapshot. """ + from sqlmesh.core.snapshot import to_table_mapping + tables_by_gateway: t.Dict[t.Union[str, None], t.List[exp.Table]] = defaultdict(list) for snapshot in target_snapshots: if snapshot.is_model and not snapshot.is_symbolic: @@ -314,6 +317,16 @@ def promote( deployability_index = deployability_index or DeployabilityIndex.all_deployable() # Renderers look snapshots up by model name, not by SnapshotId. snapshots_by_name = {s.name: s for s in (snapshots or {}).values()} + # Share the physical mapping and overrides across all promotion renders and workers. + render_table_mapping = TableMapping( + { + **to_table_mapping(snapshots_by_name.values(), deployability_index), + **(table_mapping or {}), + } + ) + for snapshot in target_snapshots: + if snapshot.is_model: + render_table_mapping.normalized_keys(snapshot.model.dialect) with self.concurrent_context(): concurrent_apply_to_snapshots( target_snapshots, @@ -324,6 +337,7 @@ def promote( execution_time=execution_time, snapshots=snapshots_by_name, table_mapping=table_mapping, + render_table_mapping=render_table_mapping, environment_naming_info=environment_naming_info, deployability_index=deployability_index, # type: ignore on_complete=on_complete, @@ -1264,6 +1278,7 @@ def _promote_snapshot( execution_time: t.Optional[TimeLike] = None, snapshots: t.Optional[t.Dict[str, Snapshot]] = None, table_mapping: t.Optional[t.Dict[str, str]] = None, + render_table_mapping: t.Optional[TableMapping] = None, ) -> None: if not snapshot.is_model: return @@ -1298,11 +1313,16 @@ def _promote_snapshot( environment=environment_naming_info.name, snapshots=snapshots, snapshot=snapshot, + render_table_mapping=render_table_mapping, **render_kwargs, ) adapter.execute( - snapshot.model.render_on_virtual_update(snapshots=snapshots, **render_kwargs) + snapshot.model.render_on_virtual_update( + snapshots=snapshots, + render_table_mapping=render_table_mapping, + **render_kwargs, + ) ) if on_complete is not None: @@ -2016,6 +2036,7 @@ def promote( snapshots=kwargs.get("snapshots"), deployability_index=kwargs.get("deployability_index"), table_mapping=kwargs.get("table_mapping"), + render_table_mapping=kwargs.get("render_table_mapping"), runtime_stage=kwargs.get("runtime_stage"), ) self.adapter.create_view( diff --git a/tests/core/test_model.py b/tests/core/test_model.py index c5d1a71bee..30c6b8156c 100644 --- a/tests/core/test_model.py +++ b/tests/core/test_model.py @@ -15,7 +15,7 @@ from sqlglot.schema import MappingSchema from sqlmesh.cli.project_init import init_example_project, ProjectTemplate from sqlmesh.core.environment import EnvironmentNamingInfo -from sqlmesh.core.renderer import TableMapping +from sqlmesh.core.renderer import ExpressionRenderer, TableMapping from sqlmesh.core.model.kind import TimeColumn, ModelKindName, SeedKind from sqlmesh import CustomMaterialization, CustomKind @@ -9900,7 +9900,10 @@ def resolve_named(evaluator, name): ) -def test_resolve_tables_expand_reveals_table_after_find_check(make_snapshot: t.Callable): +@pytest.mark.parametrize("use_cached_mapping", [False, True]) +def test_resolve_tables_expand_reveals_table_after_find_check( + make_snapshot: t.Callable, use_cached_mapping: bool +): """Embedded-model expansion (`expand=`) runs as an `expression.transform` *before* the new `expression.find(exp.Table)` short-circuit in `_resolve_tables`, so a table reference that only exists *after* inlining an embedded model's query must still be seen by `find()` and @@ -9909,6 +9912,8 @@ def test_resolve_tables_expand_reveals_table_after_find_check(make_snapshot: t.C must still resolve to its physical table name, not be silently skipped because it wasn't present at the time `_resolve_tables` was first called.""" + from sqlmesh.core.snapshot import to_table_mapping + grandparent = load_sql_based_model(d.parse("MODEL (name grandparent); SELECT 1 AS c")) grandparent_snapshot = make_snapshot(grandparent) grandparent_snapshot.categorize_as(SnapshotChangeCategory.BREAKING) @@ -9922,7 +9927,12 @@ def test_resolve_tables_expand_reveals_table_after_find_check(make_snapshot: t.C child = load_sql_based_model(d.parse("MODEL (name child); SELECT c FROM mid;")) snapshots = {'"grandparent"': grandparent_snapshot, '"mid"': mid_snapshot} - query = child.render_query(snapshots=snapshots) + query = child.render_query( + snapshots=snapshots, + render_table_mapping=( + TableMapping(to_table_mapping(snapshots.values(), None)) if use_cached_mapping else None + ), + ) assert query is not None rendered_sql = query.sql() @@ -10071,6 +10081,28 @@ def items(self): assert ItemsCountingDict.items_call_count == 0 +@pytest.mark.parametrize( + "dialect", + [ + "duckdb,normalization_strategy=case_insensitive", + "snowflake,normalization_strategy=uppercase", + ], +) +def test_cached_table_mapping_preserves_key_order(dialect): + # Updating an existing raw key must not move it after a later dialect-equivalent key. + table_mapping = {'"DB"."T"': "early_table", "DB.T": "last_table"} + table_mapping.update({'"DB"."T"': "override_table"}) + renderer = ExpressionRenderer(parse_one("SELECT * FROM DB.T", dialect=dialect), dialect, []) + + uncached = renderer.render(table_mapping=table_mapping) + cached = renderer.render( + table_mapping=table_mapping, render_table_mapping=TableMapping(table_mapping) + ) + + assert cached == uncached + assert exp.table_name(cached[0].find(exp.Table)) == "last_table" + + def test_table_mapping_normalized_keys(): table_mapping = TableMapping({'"db"."a"': "view_a", "db.A": "view_a_upper"}) @@ -10414,7 +10446,8 @@ def dummy_model_entry(evaluator: MacroEvaluator) -> exp.Select: assert len(context.engine_adapters) == 2 -def test_model_on_virtual_update(make_snapshot: t.Callable): +@pytest.mark.parametrize("use_cached_mapping", [False, True]) +def test_model_on_virtual_update(make_snapshot: t.Callable, use_cached_mapping: bool): # Macro to test resolution within virtual statement @macro() def resolve_parent_name(evaluator, name): @@ -10495,6 +10528,7 @@ def resolve_parent_name(evaluator, name): rendered_on_virtual_update = model.render_on_virtual_update( snapshots=snapshots, table_mapping=table_mapping, + render_table_mapping=TableMapping(table_mapping) if use_cached_mapping else None, this_model=model_snapshot.qualified_view_name.table_for_environment( environment_naming_info, dialect=dialect ), @@ -10526,6 +10560,7 @@ def resolve_parent_name(evaluator, name): rendered_parent_on_virtual_update = parent.render_on_virtual_update( snapshots=snapshots, table_mapping=table_mapping, + render_table_mapping=TableMapping(table_mapping) if use_cached_mapping else None, this_model=parent_snapshot.qualified_view_name.table_for_environment( environment_naming_info, dialect=dialect ), diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py index a9ca86496a..abcd2712d6 100644 --- a/tests/core/test_snapshot_evaluator.py +++ b/tests/core/test_snapshot_evaluator.py @@ -58,6 +58,7 @@ SnapshotTableCleanupTask, ) from sqlmesh.core.snapshot.definition import to_view_mapping +from sqlmesh.core import snapshot as snapshot_module from sqlmesh.core.snapshot.evaluator import ( CustomMaterialization, EngineManagedStrategy, @@ -5766,3 +5767,95 @@ def test_promote_resolves_this_model_with_single_mapping_entry( assert spy.call_count == 20 for call in spy.call_args_list: assert len(call.args[1]) == 1 + + +@pytest.mark.parametrize("ddl_concurrent_tasks", [1, 4]) +@pytest.mark.parametrize("use_view_mapping", [False, True]) +def test_promote_reuses_table_mapping( + mocker: MockerFixture, make_snapshot, ddl_concurrent_tasks, use_view_mapping +): + @macro() + def resolve_upstream_query(evaluator): + return evaluator.resolve_tables( + exp.select("COUNT(*)").from_("test_schema.model_0").subquery() + ) + + adapter_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter") + adapter_mock.dialect = "duckdb" + adapter_mock.with_settings.return_value = adapter_mock + evaluator = SnapshotEvaluator(adapter_mock, ddl_concurrent_tasks=ddl_concurrent_tasks) + + snapshots = {} + for i in range(10): + model = load_sql_based_model( + d.parse( + f""" + MODEL ( + name test_schema.model_{i}, + kind FULL, + virtual_properties ( + upstream_count = (SELECT COUNT(*) FROM test_schema.model_0), + macro_count = @resolve_upstream_query() + ), + ); + SELECT 1 AS a; + ON_VIRTUAL_UPDATE_BEGIN; + GRANT SELECT ON VIEW test_schema.model_0 TO ROLE owner; + ON_VIRTUAL_UPDATE_END; + """ + ) + ) + snapshot = make_snapshot(model) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + snapshots[snapshot.snapshot_id] = snapshot + + upstream = next(iter(snapshots.values())) + upstream.dev_version_ = "preview" + # Representative snapshots use the production table even when they are not deployable. + deployability_index = DeployabilityIndex( + indexed_ids=frozenset(), representative_shared_version_ids={upstream.snapshot_id} + ) + mapping_spy = mocker.spy(snapshot_module, "to_table_mapping") + replacement_spy = mocker.spy(exp, "replace_tables") + + for environment_name in ("first", "second"): + environment_naming_info = EnvironmentNamingInfo(name=environment_name) + view_mapping = ( + to_view_mapping(snapshots.values(), environment_naming_info) + if use_view_mapping + else None + ) + expected_table = ( + view_mapping[upstream.name] + if view_mapping is not None + else upstream.table_name(deployability_index.is_representative(upstream)) + ) + mapping_spy.reset_mock() + replacement_spy.reset_mock() + adapter_mock.create_view.reset_mock() + adapter_mock.execute.reset_mock() + + evaluator.promote( + list(snapshots.values()), + environment_naming_info, + snapshots=snapshots, + table_mapping=view_mapping, + deployability_index=deployability_index, + ) + + full_mapping_calls = [ + call for call in mapping_spy.call_args_list if len(call.args[0]) == len(snapshots) + ] + assert len(full_mapping_calls) == 1 + assert all(len(call.args[1]) <= 1 for call in replacement_spy.call_args_list) + assert adapter_mock.create_view.call_count == len(snapshots) + for call in adapter_mock.create_view.call_args_list: + for query in call.kwargs["view_properties"].values(): + assert exp.table_name(query.find(exp.Table)) == expected_table + for call in adapter_mock.execute.call_args_list: + for statement in call.args[0]: + assert exp.table_name(statement.find(exp.Table)) == expected_table + + # A later promotion must see changed versions, deployability and environment overrides. + upstream.version = "updated" + deployability_index = DeployabilityIndex.none_deployable()