sqlmesh.core.reference

  1from __future__ import annotations
  2
  3import typing as t
  4from collections import deque
  5
  6from sqlglot import exp
  7
  8from sqlmesh.utils.errors import ConfigError, SQLMeshError
  9from sqlmesh.utils.pydantic import PydanticModel
 10
 11if t.TYPE_CHECKING:
 12    from sqlmesh.core.model import Model
 13
 14
 15class Reference(PydanticModel, frozen=True):
 16    model_name: str
 17    expression: exp.Expression
 18    unique: bool = False
 19    _name: str = ""
 20
 21    @property
 22    def columns(self) -> t.List[str]:
 23        expression = self.expression
 24        if isinstance(expression, exp.Alias):
 25            expression = expression.this
 26        expression = expression.unnest()
 27        if isinstance(expression, (exp.Array, exp.Tuple)):
 28            return [e.output_name for e in expression.expressions]
 29        return [expression.output_name]
 30
 31    @property
 32    def name(self) -> str:
 33        if not self._name:
 34            keys = []
 35
 36            if isinstance(self.expression, (exp.Tuple, exp.Array)):
 37                for e in self.expression.expressions:
 38                    if not e.output_name:
 39                        raise ConfigError(
 40                            f"Reference '{e}' must have an inferrable name or explicit alias."
 41                        )
 42                    keys.append(e.output_name)
 43            elif self.expression.output_name:
 44                keys.append(self.expression.output_name)
 45            else:
 46                raise ConfigError(
 47                    f"Reference '{self.expression}' must have an inferrable name or explicit alias."
 48                )
 49
 50            self._name = "__".join(keys)
 51        return self._name
 52
 53
 54class ReferenceGraph:
 55    def __init__(self, models: t.Iterable[Model]):
 56        self._model_refs: t.Dict[str, t.Dict[str, Reference]] = {}
 57        self._ref_models: t.Dict[str, t.Set[str]] = {}
 58        self._dim_models: t.Dict[str, t.Set[str]] = {}
 59
 60        for model in models:
 61            self.add_model(model)
 62
 63    def add_model(self, model: Model) -> None:
 64        """Add a model and its references to the graph.
 65
 66        Args:
 67            model: the model to add.
 68        """
 69        for column in model.columns_to_types or {}:
 70            self._dim_models.setdefault(column, set())
 71            self._dim_models[column].add(model.name)
 72
 73        for ref in model.all_references:
 74            self._model_refs.setdefault(model.name, {})
 75            self._model_refs[model.name][ref.name] = ref
 76            self._ref_models.setdefault(ref.name, set())
 77            self._ref_models[ref.name].add(model.name)
 78
 79    def models_for_column(self, source: str, column: str, max_depth: int = 3) -> t.List[str]:
 80        """Find all the models with a column that join to a source within max_depth.
 81
 82        Args:
 83            source: The source model.
 84            column: The column to look for.
 85            max_depth: The maximum number of models to join to find a path.
 86
 87        Returns:
 88            The list of models that fit the criteria of the search.
 89        """
 90        models = []
 91
 92        for model in self._dim_models[column]:
 93            try:
 94                if model != source:
 95                    self.find_path(source, model)
 96                models.append(model)
 97            except SQLMeshError:
 98                pass
 99
100        return sorted(models)
101
102    def find_path(self, source: str, target: str, max_depth: int = 3) -> t.List[Reference]:
103        """Find a path from source model to target model with max depth.
104
105        Args:
106            source: The source model.
107            target: The target model.
108            max_depth: The maximum number of models to join to find a path.
109
110        Returns:
111            The list of references representing the join path of source to target.
112        """
113        if source not in self._model_refs:
114            return []
115
116        queue = deque(([ref] for ref in self._model_refs[source].values()))
117
118        while queue:
119            path = queue.popleft()
120            visited = set()
121            many = False
122
123            for ref in path:
124                visited.add(ref.model_name)
125                many = many or not ref.unique
126
127            ref_name = path[-1].name
128
129            for model_name in sorted(self._ref_models[ref_name]):
130                for ref in self._model_refs[model_name].values():
131                    # paths cannot have loops or contain many to many refs
132                    if model_name in visited or (many and not ref.unique):
133                        continue
134
135                    new_path = path + [ref]
136
137                    if model_name == target:
138                        return new_path
139
140                    if len(new_path) < max_depth:
141                        queue.append(new_path)
142
143        raise SQLMeshError(
144            f"Cannot find path between '{source}' and '{target}'. Make sure that references/grains are configured and that a many to many join is not occurring."
145        )