Merge branch 'main' into fix/dont-ser-engine

Author: georgesittas

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "requests",
     "rich[jupyter]",
     "ruamel.yaml",
-    "sqlglot[rs]~=27.19.0",
+    "sqlglot[rs]~=27.20.0",
     "tenacity",
     "time-machine",
     "json-stream"

sqlmesh/core/model/kind.py

@@ -119,6 +119,10 @@ def is_custom(self) -> bool:
     def is_managed(self) -> bool:
         return self.model_kind_name == ModelKindName.MANAGED
 
+    @property
+    def is_dbt_custom(self) -> bool:
+        return self.model_kind_name == ModelKindName.DBT_CUSTOM
+
     @property
     def is_symbolic(self) -> bool:
         """A symbolic model is one that doesn't execute at all."""
@@ -170,6 +174,7 @@ class ModelKindName(str, ModelKindMixin, Enum):
     EXTERNAL = "EXTERNAL"
     CUSTOM = "CUSTOM"
     MANAGED = "MANAGED"
+    DBT_CUSTOM = "DBT_CUSTOM"
 
     @property
     def model_kind_name(self) -> t.Optional[ModelKindName]:
@@ -887,6 +892,46 @@ def supports_python_models(self) -> bool:
         return False
 
 
+class DbtCustomKind(_ModelKind):
+    name: t.Literal[ModelKindName.DBT_CUSTOM] = ModelKindName.DBT_CUSTOM
+    materialization: str
+    adapter: str = "default"
+    definition: str
+    dialect: t.Optional[str] = Field(None, validate_default=True)
+
+    _dialect_validator = kind_dialect_validator
+
+    @field_validator("materialization", "adapter", "definition", mode="before")
+    @classmethod
+    def _validate_fields(cls, v: t.Any) -> str:
+        return validate_string(v)
+
+    @property
+    def data_hash_values(self) -> t.List[t.Optional[str]]:
+        return [
+            *super().data_hash_values,
+            self.materialization,
+            self.definition,
+            self.adapter,
+            self.dialect,
+        ]
+
+    def to_expression(
+        self, expressions: t.Optional[t.List[exp.Expression]] = None, **kwargs: t.Any
+    ) -> d.ModelKind:
+        return super().to_expression(
+            expressions=[
+                *(expressions or []),
+                *_properties(
+                    {
+                        "materialization": exp.Literal.string(self.materialization),
+                        "adapter": exp.Literal.string(self.adapter),
+                    }
+                ),
+            ],
+        )
+
+
 class EmbeddedKind(_ModelKind):
     name: t.Literal[ModelKindName.EMBEDDED] = ModelKindName.EMBEDDED
 
@@ -992,6 +1037,7 @@ def to_expression(
         SCDType2ByColumnKind,
         CustomKind,
         ManagedKind,
+        DbtCustomKind,
     ],
     Field(discriminator="name"),
 ]
@@ -1011,6 +1057,7 @@ def to_expression(
     ModelKindName.SCD_TYPE_2_BY_COLUMN: SCDType2ByColumnKind,
     ModelKindName.CUSTOM: CustomKind,
     ModelKindName.MANAGED: ManagedKind,
+    ModelKindName.DBT_CUSTOM: DbtCustomKind,
 }
 
 

sqlmesh/core/selector.py

@@ -16,6 +16,7 @@
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.model import update_model_schemas
+from sqlmesh.core.audit import StandaloneAudit
 from sqlmesh.utils import UniqueKeyDict
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.git import GitClient
@@ -25,6 +26,7 @@
 if t.TYPE_CHECKING:
     from typing_extensions import Literal as Lit  # noqa
     from sqlmesh.core.model import Model
+    from sqlmesh.core.node import Node
     from sqlmesh.core.state_sync import StateReader
 
 
@@ -167,7 +169,7 @@ def get_model(fqn: str) -> t.Optional[Model]:
         return models
 
     def expand_model_selections(
-        self, model_selections: t.Iterable[str], models: t.Optional[t.Dict[str, Model]] = None
+        self, model_selections: t.Iterable[str], models: t.Optional[t.Dict[str, Node]] = None
     ) -> t.Set[str]:
         """Expands a set of model selections into a set of model fqns that can be looked up in the Context.
 
@@ -180,7 +182,7 @@ def expand_model_selections(
 
         node = parse(" | ".join(f"({s})" for s in model_selections))
 
-        all_models = models or self._models
+        all_models: t.Dict[str, Node] = models or dict(self._models)
         models_by_tags: t.Dict[str, t.Set[str]] = {}
 
         for fqn, model in all_models.items():
@@ -226,6 +228,13 @@ def evaluate(node: exp.Expression) -> t.Set[str]:
                         if fnmatch.fnmatchcase(tag, pattern)
                     }
                 return models_by_tags.get(pattern, set())
+            if isinstance(node, ResourceType):
+                resource_type = node.name.lower()
+                return {
+                    fqn
+                    for fqn, model in all_models.items()
+                    if self._matches_resource_type(resource_type, model)
+                }
             if isinstance(node, Direction):
                 selected = set()
 
@@ -243,36 +252,49 @@ def evaluate(node: exp.Expression) -> t.Set[str]:
         return evaluate(node)
 
     @abc.abstractmethod
-    def _model_name(self, model: Model) -> str:
+    def _model_name(self, model: Node) -> str:
         """Given a model, return the name that a selector pattern contining wildcards should be fnmatch'd on"""
         pass
 
     @abc.abstractmethod
-    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Model]) -> t.Set[str]:
+    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Node]) -> t.Set[str]:
         """Given a pattern, return the keys of the matching models from :all_models"""
         pass
 
+    @abc.abstractmethod
+    def _matches_resource_type(self, resource_type: str, model: Node) -> bool:
+        """Indicate whether or not the supplied model matches the supplied resource type"""
+        pass
+
 
 class NativeSelector(Selector):
     """Implementation of selectors that matches objects based on SQLMesh native names"""
 
-    def _model_name(self, model: Model) -> str:
+    def _model_name(self, model: Node) -> str:
         return model.name
 
-    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Model]) -> t.Set[str]:
+    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Node]) -> t.Set[str]:
         fqn = normalize_model_name(pattern, self._default_catalog, self._dialect)
         return {fqn} if fqn in all_models else set()
 
+    def _matches_resource_type(self, resource_type: str, model: Node) -> bool:
+        if resource_type == "model":
+            return model.is_model
+        if resource_type == "audit":
+            return isinstance(model, StandaloneAudit)
+
+        raise SQLMeshError(f"Unsupported resource type: {resource_type}")
+
 
 class DbtSelector(Selector):
     """Implementation of selectors that matches objects based on the DBT names instead of the SQLMesh native names"""
 
-    def _model_name(self, model: Model) -> str:
+    def _model_name(self, model: Node) -> str:
         if dbt_fqn := model.dbt_fqn:
             return dbt_fqn
         raise SQLMeshError("dbt node information must be populated to use dbt selectors")
 
-    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Model]) -> t.Set[str]:
+    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Node]) -> t.Set[str]:
         # a pattern like "staging.customers" should match a model called "jaffle_shop.staging.customers"
         # but not a model called "jaffle_shop.customers.staging"
         # also a pattern like "aging" should not match "staging" so we need to consider components; not substrings
@@ -306,6 +328,40 @@ def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Model]) -
                 matches.add(fqn)
         return matches
 
+    def _matches_resource_type(self, resource_type: str, model: Node) -> bool:
+        """
+        ref: https://docs.getdbt.com/reference/node-selection/methods#resource_type
+
+        # supported by SQLMesh
+        "model"
+        "seed"
+        "source" # external model
+        "test" # standalone audit
+
+        # not supported by SQLMesh yet, commented out to throw an error if someone tries to use them
+        "analysis"
+        "exposure"
+        "metric"
+        "saved_query"
+        "semantic_model"
+        "snapshot"
+        "unit_test"
+        """
+        if resource_type not in ("model", "seed", "source", "test"):
+            raise SQLMeshError(f"Unsupported resource type: {resource_type}")
+
+        if isinstance(model, StandaloneAudit):
+            return resource_type == "test"
+
+        if resource_type == "model":
+            return model.is_model and not model.kind.is_external and not model.kind.is_seed
+        if resource_type == "source":
+            return model.kind.is_external
+        if resource_type == "seed":
+            return model.kind.is_seed
+
+        return False
+
 
 class SelectorDialect(Dialect):
     IDENTIFIERS_CAN_START_WITH_DIGIT = True
@@ -336,6 +392,10 @@ class Tag(exp.Expression):
     pass
 
 
+class ResourceType(exp.Expression):
+    pass
+
+
 class Direction(exp.Expression):
     pass
 
@@ -388,7 +448,8 @@ def _parse_var() -> exp.Expression:
         upstream = _match(TokenType.PLUS)
         downstream = None
         tag = _parse_kind("tag")
-        git = False if tag else _parse_kind("git")
+        resource_type = False if tag else _parse_kind("resource_type")
+        git = False if resource_type else _parse_kind("git")
         lstar = "*" if _match(TokenType.STAR) else ""
         directions = {}
 
@@ -414,6 +475,8 @@ def _parse_var() -> exp.Expression:
 
         if tag:
             this = Tag(this=this)
+        if resource_type:
+            this = ResourceType(this=this)
         if git:
             this = Git(this=this)
         if directions: