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