Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 88 additions & 9 deletions sqlmesh/core/renderer.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,65 @@
logger = logging.getLogger(__name__)


class TableMapping(t.Dict[str, str]):
"""A table name mapping that caches the dialect-normalized form of its keys.

`exp.replace_tables` normalizes every key of the mapping it's given, so resolving a single
table against a mapping of every model in an environment costs O(N). Resolving it against
this mapping costs a dictionary lookup, since each key is normalized once per dialect.
"""

def __init__(self, *args: t.Any, **kwargs: t.Any):
super().__init__(*args, **kwargs)
self._normalized_keys: t.Dict[DialectType, t.Dict[str, str]] = {}

def normalized_keys(self, dialect: DialectType) -> t.Dict[str, str]:
"""Returns a mapping from each normalized key to the last key that normalizes to it."""
normalized_keys = self._normalized_keys.get(dialect)
if normalized_keys is None:
normalized_keys = {exp.normalize_table_name(key, dialect=dialect): key for key in self}
self._normalized_keys[dialect] = normalized_keys
return normalized_keys

def __setitem__(self, key: str, value: str) -> None:
self._normalized_keys.clear()
super().__setitem__(key, value)

def __delitem__(self, key: str) -> None:
self._normalized_keys.clear()
super().__delitem__(key)

def __ior__(self, other: t.Any) -> TableMapping: # type: ignore[override,misc]
self._normalized_keys.clear()
return super().__ior__(other)

def update(self, *args: t.Any, **kwargs: t.Any) -> None:
self._normalized_keys.clear()
super().update(*args, **kwargs)

def setdefault(self, key: str, default: str) -> str: # type: ignore[override]
self._normalized_keys.clear()
return super().setdefault(key, default)

def pop(self, key: str, *args: t.Any) -> t.Any:
self._normalized_keys.clear()
return super().pop(key, *args)

def popitem(self) -> t.Tuple[str, str]:
self._normalized_keys.clear()
return super().popitem()

def clear(self) -> None:
self._normalized_keys.clear()
super().clear()


def _normalize_keys(mapping: t.Dict[str, str], dialect: DialectType) -> t.Dict[str, str]:
if isinstance(mapping, TableMapping):
return mapping.normalized_keys(dialect)
return {exp.normalize_table_name(key, dialect=dialect): key for key in mapping}


class BaseExpressionRenderer:
def __init__(
self,
Expand Down Expand Up @@ -325,20 +384,35 @@ def update_cache(self, expression: t.Optional[exp.Expr]) -> None:

def _resolve_table(
self,
table_name: str | exp.Expr,
table_name: str,
snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
table_mapping: t.Optional[t.Dict[str, str]] = None,
deployability_index: t.Optional[DeployabilityIndex] = None,
) -> exp.Table:
table = exp.replace_tables(
t.cast(exp.Table, exp.maybe_parse(table_name, into=exp.Table, dialect=self._dialect)),
{
**self._to_table_mapping((snapshots or {}).values(), deployability_index),
**(table_mapping or {}),
},
dialect=self._dialect,
copy=False,
table = t.cast(
exp.Table, exp.maybe_parse(table_name, into=exp.Table, dialect=self._dialect)
)

mapping: t.Dict[str, str] = {}
if table_mapping:
# An explicit mapping takes precedence over snapshots, so when one of its keys matches
# the table, that key alone decides the result. Among equivalent keys, the last wins.
key = _normalize_keys(table_mapping, self._dialect).get(
exp.normalize_table_name(table, dialect=self._dialect)
)
if key is not None:
mapping = {key: table_mapping[key]}

if not mapping and snapshots:
# An exact FQN match avoids scanning unrelated snapshots.
snapshot = snapshots.get(table_name)
# Keys normalized under different dialects may differ in casing or quoting.
# Fall back to the full mapping so exp.replace_tables can reconcile them.
mapping = self._to_table_mapping(
[snapshot] if snapshot else snapshots.values(), deployability_index
)

table = exp.replace_tables(table, mapping, dialect=self._dialect, copy=False)
# We quote the table here to mimic the behavior of _resolve_tables, otherwise we may end
# up normalizing twice, because _to_table_mapping returns the mapped names unquoted.
return (
Expand All @@ -363,6 +437,11 @@ def _resolve_tables(

expression = expression.copy()
with self._normalize_and_quote(expression) as expression:
# An expression with no table (e.g. most session or virtual properties) has nothing
# to expand or replace, so skip building the O(N) expand set and mapping.
if not expression.find(exp.Table):
return expression

snapshots = snapshots or {}
table_mapping = table_mapping or {}
mapping = {
Expand Down
14 changes: 9 additions & 5 deletions sqlmesh/core/snapshot/definition.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
from sqlmesh.core.model import Model, ModelKindMixin, ModelKindName, ViewKind, CustomKind
from sqlmesh.core.model.definition import _Model
from sqlmesh.core.node import IntervalUnit, NodeType
from sqlmesh.core.renderer import TableMapping
from sqlmesh.utils import sanitize_name, unique
from sqlmesh.utils.dag import DAG
from sqlmesh.utils.date import (
Expand Down Expand Up @@ -2007,14 +2008,17 @@ def to_view_mapping(
environment_naming_info: EnvironmentNamingInfo,
default_catalog: t.Optional[str] = None,
dialect: t.Optional[str] = None,
) -> t.Dict[str, str]:
return {
snapshot.name: snapshot.display_name(
environment_naming_info, default_catalog=default_catalog, dialect=dialect
) -> TableMapping:
return TableMapping(
(
snapshot.name,
snapshot.display_name(
environment_naming_info, default_catalog=default_catalog, dialect=dialect
),
)
for snapshot in snapshots
if snapshot.is_model
}
)


def has_paused_forward_only(
Expand Down
12 changes: 7 additions & 5 deletions sqlmesh/core/snapshot/evaluator.py
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,8 @@ def promote(
self._get_virtual_data_objects(target_snapshots, environment_naming_info)

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()}
with self.concurrent_context():
concurrent_apply_to_snapshots(
target_snapshots,
Expand All @@ -320,7 +322,7 @@ def promote(
start=start,
end=end,
execution_time=execution_time,
snapshots=snapshots,
snapshots=snapshots_by_name,
table_mapping=table_mapping,
environment_naming_info=environment_naming_info,
deployability_index=deployability_index, # type: ignore
Expand Down Expand Up @@ -1260,7 +1262,7 @@ def _promote_snapshot(
start: t.Optional[TimeLike] = None,
end: t.Optional[TimeLike] = None,
execution_time: t.Optional[TimeLike] = None,
snapshots: t.Optional[t.Dict[SnapshotId, Snapshot]] = None,
snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
table_mapping: t.Optional[t.Dict[str, str]] = None,
) -> None:
if not snapshot.is_model:
Expand Down Expand Up @@ -1299,9 +1301,9 @@ def _promote_snapshot(
**render_kwargs,
)

snapshot_by_name = {s.name: s for s in (snapshots or {}).values()}
render_kwargs["snapshots"] = snapshot_by_name
adapter.execute(snapshot.model.render_on_virtual_update(**render_kwargs))
adapter.execute(
snapshot.model.render_on_virtual_update(snapshots=snapshots, **render_kwargs)
)

if on_complete is not None:
on_complete(snapshot)
Expand Down
Loading
Loading