sqlmesh.core.selector

  1from __future__ import annotations
  2
  3import fnmatch
  4import typing as t
  5from pathlib import Path
  6from itertools import zip_longest
  7import abc
  8
  9from sqlglot import exp
 10from sqlglot.errors import ParseError
 11from sqlglot.tokens import Token, TokenType, Tokenizer as BaseTokenizer
 12from sqlglot.dialects.dialect import Dialect, DialectType
 13from sqlglot.helper import seq_get
 14
 15from sqlmesh.core import constants as c
 16from sqlmesh.core.dialect import normalize_model_name
 17from sqlmesh.core.environment import Environment
 18from sqlmesh.core.model import update_model_schemas
 19from sqlmesh.core.audit import StandaloneAudit
 20from sqlmesh.utils import UniqueKeyDict
 21from sqlmesh.utils.dag import DAG
 22from sqlmesh.utils.git import GitClient
 23from sqlmesh.utils.errors import SQLMeshError
 24
 25
 26if t.TYPE_CHECKING:
 27    from typing_extensions import Literal as Lit  # noqa
 28    from sqlmesh.core.model import Model
 29    from sqlmesh.core.node import Node
 30    from sqlmesh.core.state_sync import StateReader
 31
 32
 33class Selector(abc.ABC):
 34    def __init__(
 35        self,
 36        state_reader: StateReader,
 37        models: UniqueKeyDict[str, Model],
 38        context_path: Path = Path("."),
 39        dag: t.Optional[DAG[str]] = None,
 40        default_catalog: t.Optional[str] = None,
 41        dialect: t.Optional[str] = None,
 42        cache_dir: t.Optional[Path] = None,
 43    ):
 44        self._state_reader = state_reader
 45        self._models = models
 46        self._context_path = context_path
 47        self._cache_dir = cache_dir if cache_dir else context_path / c.CACHE
 48        self._default_catalog = default_catalog
 49        self._dialect = dialect
 50        self._git_client = GitClient(context_path)
 51
 52        if dag is None:
 53            self._dag: DAG[str] = DAG()
 54            for fqn, model in models.items():
 55                self._dag.add(fqn, model.depends_on)
 56        else:
 57            self._dag = dag
 58
 59    def select_models(
 60        self,
 61        model_selections: t.Iterable[str],
 62        target_env_name: str,
 63        fallback_env_name: t.Optional[str] = None,
 64        ensure_finalized_snapshots: bool = False,
 65    ) -> t.Tuple[UniqueKeyDict[str, Model], t.Set[str]]:
 66        """Given a set of selections returns models from the current state with names matching the
 67        selection while sourcing the remaining models from the target environment.
 68
 69        Args:
 70            model_selections: A set of selections.
 71            target_env_name: The name of the target environment.
 72            fallback_env_name: The name of the fallback environment that will be used if the target
 73                environment doesn't exist.
 74            ensure_finalized_snapshots: Whether to source environment snapshots from the latest finalized
 75                environment state, or to use whatever snapshots are in the current environment state even if
 76                the environment is not finalized.
 77
 78        Returns:
 79            A tuple of (models dict, set of all matched FQNs including env models).
 80        """
 81        env_models = self._load_env_models(
 82            target_env_name, fallback_env_name, ensure_finalized_snapshots
 83        )
 84
 85        all_selected_models = self.expand_model_selections(
 86            model_selections, models={**env_models, **self._models}
 87        )
 88
 89        dag: DAG[str] = DAG()
 90        subdag = set()
 91
 92        for fqn in all_selected_models:
 93            if fqn not in subdag:
 94                subdag.add(fqn)
 95                subdag.update(self._dag.downstream(fqn))
 96
 97        models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
 98        all_model_fqns = set(self._models) | set(env_models)
 99        needs_update = False
100
101        def get_model(fqn: str) -> t.Optional[Model]:
102            if fqn not in all_selected_models and fqn in env_models:
103                # Unselected modified or added model.
104                model_from_env = env_models[fqn]
105                try:
106                    # this triggers a render_query() which can throw an exception
107                    model_from_env.depends_on
108                    return model_from_env
109                except Exception as e:
110                    raise SQLMeshError(
111                        f"Model '{model_from_env.name}' sourced from state cannot be rendered "
112                        f"in the local environment due to:\n> {str(e)}"
113                    ) from e
114            if fqn in all_selected_models and fqn in self._models:
115                # Selected modified or removed model.
116                return self._models[fqn]
117            return None
118
119        for fqn in all_model_fqns:
120            model = get_model(fqn)
121
122            if not model:
123                continue
124
125            if model.fqn in subdag:
126                dag.add(model.fqn, model.depends_on)
127
128                for dep in model.depends_on:
129                    schema = model.mapping_schema
130
131                    for part in exp.to_table(dep).parts:
132                        schema = schema.get(part.sql()) or {}
133
134                    parent = get_model(dep)
135
136                    parent_schema = {
137                        c: t.sql(dialect=model.dialect)
138                        for c, t in ((parent and parent.columns_to_types) or {}).items()
139                    }
140
141                    if schema != parent_schema:
142                        model = model.copy(update={"mapping_schema": {}})
143                        needs_update = True
144                        break
145
146            models[model.fqn] = model
147
148        if needs_update:
149            update_model_schemas(dag, models=models, cache_dir=self._cache_dir)
150
151        return models, all_selected_models
152
153    def _load_env_models(
154        self,
155        target_env_name: str,
156        fallback_env_name: t.Optional[str] = None,
157        ensure_finalized_snapshots: bool = False,
158    ) -> t.Dict[str, "Model"]:
159        """Loads models from the target environment, falling back to the fallback environment if needed."""
160        target_env = self._state_reader.get_environment(Environment.sanitize_name(target_env_name))
161        if target_env and target_env.expired:
162            target_env = None
163
164        if not target_env and fallback_env_name:
165            target_env = self._state_reader.get_environment(
166                Environment.sanitize_name(fallback_env_name)
167            )
168
169        if not target_env:
170            return {}
171
172        environment_snapshot_infos = (
173            target_env.snapshots
174            if not ensure_finalized_snapshots
175            else target_env.finalized_or_current_snapshots
176        )
177        return {
178            s.name: s.model
179            for s in self._state_reader.get_snapshots(environment_snapshot_infos).values()
180            if s.is_model
181        }
182
183    def expand_model_selections(
184        self, model_selections: t.Iterable[str], models: t.Optional[t.Dict[str, Node]] = None
185    ) -> t.Set[str]:
186        """Expands a set of model selections into a set of model fqns that can be looked up in the Context.
187
188        Args:
189            model_selections: A set of model selections.
190
191        Returns:
192            A set of model fqns.
193        """
194
195        node = parse(" | ".join(f"({s})" for s in model_selections))
196
197        all_models: t.Dict[str, Node] = models or dict(self._models)
198        models_by_tags: t.Dict[str, t.Set[str]] = {}
199
200        for fqn, model in all_models.items():
201            for tag in model.tags:
202                tag = tag.lower()
203                models_by_tags.setdefault(tag, set())
204                models_by_tags[tag].add(model.fqn)
205
206        def evaluate(node: exp.Expr) -> t.Set[str]:
207            if isinstance(node, exp.Var):
208                pattern = node.this
209                if "*" in pattern:
210                    return {
211                        fqn
212                        for fqn, model in all_models.items()
213                        if fnmatch.fnmatchcase(self._model_name(model), node.this)
214                    }
215                return self._pattern_to_model_fqns(pattern, all_models)
216            if isinstance(node, exp.And):
217                return evaluate(node.left) & evaluate(node.right)
218            if isinstance(node, exp.Or):
219                return evaluate(node.left) | evaluate(node.right)
220            if isinstance(node, exp.Paren):
221                return evaluate(node.this)
222            if isinstance(node, exp.Not):
223                return set(all_models) - evaluate(node.this)
224            if isinstance(node, Git):
225                target_branch = node.name
226                git_modified_files = {
227                    *self._git_client.list_untracked_files(),
228                    *self._git_client.list_uncommitted_changed_files(),
229                    *self._git_client.list_committed_changed_files(target_branch=target_branch),
230                }
231                return {m.fqn for m in all_models.values() if m._path in git_modified_files}
232            if isinstance(node, Tag):
233                pattern = node.name.lower()
234
235                if "*" in pattern:
236                    return {
237                        model
238                        for tag, models in models_by_tags.items()
239                        for model in models
240                        if fnmatch.fnmatchcase(tag, pattern)
241                    }
242                return models_by_tags.get(pattern, set())
243            if isinstance(node, ResourceType):
244                resource_type = node.name.lower()
245                return {
246                    fqn
247                    for fqn, model in all_models.items()
248                    if self._matches_resource_type(resource_type, model)
249                }
250            if isinstance(node, Direction):
251                selected = set()
252
253                for model_name in evaluate(node.this):
254                    selected.add(model_name)
255                    if node.args.get("up"):
256                        for u in self._dag.upstream(model_name):
257                            if u in all_models:
258                                selected.add(u)
259                    if node.args.get("down"):
260                        selected.update(self._dag.downstream(model_name))
261                return selected
262            raise ParseError(f"Unexpected node {node}")
263
264        return evaluate(node)
265
266    @abc.abstractmethod
267    def _model_name(self, model: Node) -> str:
268        """Given a model, return the name that a selector pattern contining wildcards should be fnmatch'd on"""
269        pass
270
271    @abc.abstractmethod
272    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Node]) -> t.Set[str]:
273        """Given a pattern, return the keys of the matching models from :all_models"""
274        pass
275
276    @abc.abstractmethod
277    def _matches_resource_type(self, resource_type: str, model: Node) -> bool:
278        """Indicate whether or not the supplied model matches the supplied resource type"""
279        pass
280
281
282class NativeSelector(Selector):
283    """Implementation of selectors that matches objects based on SQLMesh native names"""
284
285    def _model_name(self, model: Node) -> str:
286        return model.name
287
288    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Node]) -> t.Set[str]:
289        fqn = normalize_model_name(pattern, self._default_catalog, self._dialect)
290        return {fqn} if fqn in all_models else set()
291
292    def _matches_resource_type(self, resource_type: str, model: Node) -> bool:
293        if resource_type == "model":
294            return model.is_model
295        if resource_type == "audit":
296            return isinstance(model, StandaloneAudit)
297
298        raise SQLMeshError(f"Unsupported resource type: {resource_type}")
299
300
301class DbtSelector(Selector):
302    """Implementation of selectors that matches objects based on the DBT names instead of the SQLMesh native names"""
303
304    def _model_name(self, model: Node) -> str:
305        if dbt_fqn := model.dbt_fqn:
306            return dbt_fqn
307        raise SQLMeshError("dbt node information must be populated to use dbt selectors")
308
309    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Node]) -> t.Set[str]:
310        # a pattern like "staging.customers" should match a model called "jaffle_shop.staging.customers"
311        # but not a model called "jaffle_shop.customers.staging"
312        # also a pattern like "aging" should not match "staging" so we need to consider components; not substrings
313        pattern_components = pattern.split(".")
314        first_pattern_component = pattern_components[0]
315        matches = set()
316        for fqn, model in all_models.items():
317            if not model.dbt_fqn:
318                continue
319
320            dbt_fqn_components = model.dbt_fqn.split(".")
321            try:
322                starting_idx = dbt_fqn_components.index(first_pattern_component)
323            except ValueError:
324                continue
325            for pattern_component, fqn_component in zip_longest(
326                pattern_components, dbt_fqn_components[starting_idx:]
327            ):
328                if pattern_component and not fqn_component:
329                    # the pattern still goes but we have run out of fqn components to match; no match
330                    break
331                if fqn_component and not pattern_component:
332                    # all elements of the pattern have matched elements of the fqn; match
333                    matches.add(fqn)
334                    break
335                if pattern_component != fqn_component:
336                    # the pattern explicitly doesnt match a component; no match
337                    break
338            else:
339                # called if no explicit break, indicating all components of the pattern matched all components of the fqn
340                matches.add(fqn)
341        return matches
342
343    def _matches_resource_type(self, resource_type: str, model: Node) -> bool:
344        """
345        ref: https://docs.getdbt.com/reference/node-selection/methods#resource_type
346
347        # supported by SQLMesh
348        "model"
349        "seed"
350        "source" # external model
351        "test" # standalone audit
352
353        # not supported by SQLMesh yet, commented out to throw an error if someone tries to use them
354        "analysis"
355        "exposure"
356        "metric"
357        "saved_query"
358        "semantic_model"
359        "snapshot"
360        "unit_test"
361        """
362        if resource_type not in ("model", "seed", "source", "test"):
363            raise SQLMeshError(f"Unsupported resource type: {resource_type}")
364
365        if isinstance(model, StandaloneAudit):
366            return resource_type == "test"
367
368        if resource_type == "model":
369            return model.is_model and not model.kind.is_external and not model.kind.is_seed
370        if resource_type == "source":
371            return model.kind.is_external
372        if resource_type == "seed":
373            return model.kind.is_seed
374
375        return False
376
377
378class SelectorDialect(Dialect):
379    IDENTIFIERS_CAN_START_WITH_DIGIT = True
380
381    class Tokenizer(BaseTokenizer):
382        SINGLE_TOKENS = {
383            "(": TokenType.L_PAREN,
384            ")": TokenType.R_PAREN,
385            "&": TokenType.AMP,
386            "|": TokenType.PIPE,
387            "^": TokenType.CARET,
388            "+": TokenType.PLUS,
389            "*": TokenType.STAR,
390            ":": TokenType.COLON,
391        }
392
393        KEYWORDS = {}
394        IDENTIFIERS = ["\\"]  # there are no identifiers but need to put something here
395        IDENTIFIER_START = ""
396        IDENTIFIER_END = ""
397
398
399class Git(exp.Expression):
400    pass
401
402
403class Tag(exp.Expression):
404    pass
405
406
407class ResourceType(exp.Expression):
408    pass
409
410
411class Direction(exp.Expression):
412    pass
413
414
415def parse(selector: str, dialect: DialectType = None) -> exp.Expr:
416    tokens = SelectorDialect().tokenize(selector)
417    i = 0
418
419    def _curr() -> t.Optional[Token]:
420        return seq_get(tokens, i)
421
422    def _prev() -> Token:
423        return tokens[i - 1]
424
425    def _advance(num: int = 1) -> Token:
426        nonlocal i
427        i += num
428        return _prev()
429
430    def _next() -> t.Optional[Token]:
431        return seq_get(tokens, i + 1)
432
433    def _error(msg: str) -> str:
434        return f"{msg} at index {i}: {selector}"
435
436    def _match(token_type: TokenType, raise_unmatched: bool = False) -> t.Optional[Token]:
437        token = _curr()
438        if token and token.token_type == token_type:
439            return _advance()
440        if raise_unmatched:
441            raise ParseError(_error(f"Expected {token_type}"))
442        return None
443
444    def _parse_kind(kind: str) -> bool:
445        token = _curr()
446        next_token = _next()
447
448        if (
449            token
450            and token.token_type == TokenType.VAR
451            and token.text.lower() == kind
452            and next_token
453            and next_token.token_type == TokenType.COLON
454        ):
455            _advance(2)
456            return True
457        return False
458
459    def _parse_var() -> exp.Expr:
460        upstream = _match(TokenType.PLUS)
461        downstream = None
462        tag = _parse_kind("tag")
463        resource_type = False if tag else _parse_kind("resource_type")
464        git = False if resource_type else _parse_kind("git")
465        lstar = "*" if _match(TokenType.STAR) else ""
466        directions = {}
467
468        if _match(TokenType.VAR) or _match(TokenType.NUMBER):
469            name = _prev().text
470            rstar = "*" if _match(TokenType.STAR) else ""
471            downstream = _match(TokenType.PLUS)
472            this: exp.Expr = exp.Var(this=f"{lstar}{name}{rstar}")
473
474        elif _match(TokenType.L_PAREN):
475            this = exp.Paren(this=_parse_conjunction())
476            downstream = _match(TokenType.PLUS)
477            _match(TokenType.R_PAREN, True)
478        elif lstar:
479            this = exp.var("*")
480        else:
481            raise ParseError(_error("Expected model name."))
482
483        if upstream:
484            directions["up"] = True
485        if downstream:
486            directions["down"] = True
487
488        if tag:
489            this = Tag(this=this)
490        if resource_type:
491            this = ResourceType(this=this)
492        if git:
493            this = Git(this=this)
494        if directions:
495            this = Direction(this=this, **directions)
496        return this
497
498    def _parse_unary() -> exp.Expr:
499        if _match(TokenType.CARET):
500            return exp.Not(this=_parse_unary())
501        return _parse_var()
502
503    def _parse_conjunction() -> exp.Expr:
504        this = _parse_unary()
505
506        if _match(TokenType.AMP):
507            this = exp.And(this=this, expression=_parse_unary())
508        if _match(TokenType.PIPE):
509            this = exp.Or(this=this, expression=_parse_conjunction())
510
511        return this
512
513    return _parse_conjunction()