Dbt columns to sqlmesh (#391)

* Move dbt model creation from load_model to create_sql_model

* Rename parse_model to parse

* Add DBT project columns to SQLMesh models

* Fix test

* typo

* Use dict comprehension

Author: crericha

sqlmesh/core/context.py

@@ -49,7 +49,7 @@
 from sqlmesh.core.config import Config, load_config_from_paths
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.context_diff import ContextDiff
-from sqlmesh.core.dialect import format_model_expressions, pandas_to_sql, parse_model
+from sqlmesh.core.dialect import format_model_expressions, pandas_to_sql, parse
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.hooks import hook
@@ -552,7 +552,7 @@ def format(self) -> None:
             if not model.is_sql:
                 continue
             with open(model._path, "r+", encoding="utf-8") as file:
-                expressions = parse_model(file.read(), default_dialect=self.dialect)
+                expressions = parse(file.read(), default_dialect=self.dialect)
                 file.seek(0)
                 file.write(format_model_expressions(expressions, model.dialect))
                 file.truncate()

sqlmesh/core/dialect.py

@@ -383,9 +383,10 @@ def text_diff(
 JINJA_PATTERN = re.compile(r"{{|{%|{#")
 
 
-def parse_model(sql: str, default_dialect: str | None = None) -> t.List[exp.Expression]:
-    """Parse a sql string containing a model definition.
+def parse(sql: str, default_dialect: str | None = None) -> t.List[exp.Expression]:
+    """Parse a sql string.
 
+    Supports parsing model definition.
     If a jinja block is detected, the query is stored as raw string in a Jinja node.
 
     Args:

sqlmesh/core/loader.py

@@ -15,7 +15,7 @@
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.audit import Audit
-from sqlmesh.core.dialect import parse_model
+from sqlmesh.core.dialect import parse
 from sqlmesh.core.hooks import HookRegistry, hook
 from sqlmesh.core.macros import MacroRegistry, macro
 from sqlmesh.core.model import Model, SeedModel, load_model
@@ -190,7 +190,7 @@ def _load_sql_models(
             self._track_file(path)
             with open(path, "r", encoding="utf-8") as file:
                 try:
-                    expressions = parse_model(file.read(), default_dialect=self._context.dialect)
+                    expressions = parse(file.read(), default_dialect=self._context.dialect)
                 except SqlglotError as ex:
                     raise ConfigError(f"Failed to parse a model definition at '{path}': {ex}")
                 model = load_model(
@@ -240,7 +240,7 @@ def _load_audits(self) -> UniqueKeyDict[str, Audit]:
         for path in self._context.glob_path(self._context.audits_directory_path, ".sql"):
             self._track_file(path)
             with open(path, "r", encoding="utf-8") as file:
-                expressions = parse_model(file.read(), default_dialect=self._context.dialect)
+                expressions = parse(file.read(), default_dialect=self._context.dialect)
                 audits = Audit.load_multiple(
                     expressions=expressions,
                     path=path,

sqlmesh/core/model/common.py

@@ -7,7 +7,7 @@
 from sqlglot.expressions import split_num_words
 from sqlglot.helper import seq_get
 
-from sqlmesh.core.dialect import parse_model
+from sqlmesh.core.dialect import parse
 from sqlmesh.utils.errors import ConfigError
 
 
@@ -37,7 +37,7 @@ def parse_expression(
         return [e for e in (maybe_parse(i) for i in v) if e]
 
     if isinstance(v, str):
-        return seq_get(parse_model(v), 0)
+        return seq_get(parse(v), 0)
 
     if not v:
         raise ConfigError(f"Could not parse {v}")

sqlmesh/core/model/definition.py

@@ -382,11 +382,6 @@ def depends_on(self) -> t.Set[str]:
             self._depends_on = _find_tables(self.render_query()) - {self.name}
         return self._depends_on
 
-    @property
-    def column_descriptions(self) -> t.Dict[str, str]:
-        """A dictionary of column names to annotation comments."""
-        return {}
-
     @property
     def columns_to_types(self) -> t.Dict[str, exp.DataType]:
         """Returns the mapping of column names to types of this model."""
@@ -583,6 +578,9 @@ def columns_to_types(self) -> t.Dict[str, exp.DataType]:
 
     @property
     def column_descriptions(self) -> t.Dict[str, str]:
+        if self.column_descriptions_ is not None:
+            return self.column_descriptions_
+
         if self._column_descriptions is None:
             self._column_descriptions = {
                 select.alias: "\n".join(comment.strip() for comment in select.comments)
@@ -847,7 +845,7 @@ def load_model(
         return create_sql_model(
             name,
             query,
-            statements,
+            statements=statements,
             defaults=defaults,
             path=path,
             module_path=module_path,
@@ -880,8 +878,8 @@ def load_model(
 def create_sql_model(
     name: str,
     query: exp.Expression,
-    statements: t.List[exp.Expression],
     *,
+    statements: t.Optional[t.List[exp.Expression]] = None,
     defaults: t.Optional[t.Dict[str, t.Any]] = None,
     path: Path = Path(),
     module_path: Path = Path(),
@@ -936,7 +934,7 @@ def create_sql_model(
         time_column_format=time_column_format,
         python_env=python_env,
         dialect=dialect,
-        expressions=statements,
+        expressions=statements or [],
         query=query,
         **kwargs,
     )

sqlmesh/core/model/meta.py

@@ -55,6 +55,7 @@ class ModelMeta(PydanticModel):
     post: t.List[HookCall] = []
     depends_on_: t.Optional[t.Set[str]] = Field(default=None, alias="depends_on")
     columns_to_types_: t.Optional[t.Dict[str, exp.DataType]] = Field(default=None, alias="columns")
+    column_descriptions_: t.Optional[t.Dict[str, str]]
     audits: t.List[AuditReference] = []
 
     _croniter: t.Optional[croniter] = None
@@ -193,6 +194,11 @@ def partitioned_by(self) -> t.List[str]:
         time_column = [self.time_column.column] if self.time_column else []
         return unique([*time_column, *self.partitioned_by_])
 
+    @property
+    def column_descriptions(self) -> t.Dict[str, str]:
+        """A dictionary of column names to annotation comments."""
+        return self.column_descriptions_ or {}
+
     def interval_unit(self, sample_size: int = 10) -> IntervalUnit:
         """Returns the IntervalUnit of the model
 

sqlmesh/dbt/column.py

@@ -3,6 +3,7 @@
 import typing as t
 
 from pydantic import validator
+from sqlglot import exp, parse_one
 from sqlglot.helper import ensure_list
 
 from sqlmesh.dbt.common import GeneralConfig
@@ -21,6 +22,30 @@ def yaml_to_columns(
     return columns
 
 
+def column_types_to_sqlmesh(columns: t.Dict[str, ColumnConfig]) -> t.Dict[str, exp.DataType]:
+    """
+    Get the sqlmesh column types
+
+    Returns:
+        A dict of column name to exp.DataType
+    """
+    return {
+        name: parse_one(column.data_type, into=exp.DataType)
+        for name, column in columns.items()
+        if column.data_type
+    }
+
+
+def column_descriptions_to_sqlmesh(columns: t.Dict[str, ColumnConfig]) -> t.Dict[str, str]:
+    """
+    Get the sqlmesh column types
+
+    Returns:
+        A dict of column name to description
+    """
+    return {name: column.description for name, column in columns.items() if column.description}
+
+
 class ColumnConfig(GeneralConfig):
     """
     Column configuration for a DBT project

sqlmesh/dbt/model.py

@@ -9,8 +9,20 @@
 
 from sqlmesh.core import dialect as d
 from sqlmesh.core.config.base import UpdateStrategy
-from sqlmesh.core.model import Model, ModelKindName, load_model
-from sqlmesh.dbt.column import ColumnConfig, yaml_to_columns
+from sqlmesh.core.model import (
+    IncrementalByTimeRangeKind,
+    IncrementalByUniqueKeyKind,
+    Model,
+    ModelKind,
+    ModelKindName,
+    create_sql_model,
+)
+from sqlmesh.dbt.column import (
+    ColumnConfig,
+    column_descriptions_to_sqlmesh,
+    column_types_to_sqlmesh,
+    yaml_to_columns,
+)
 from sqlmesh.dbt.common import Dependencies, GeneralConfig
 from sqlmesh.dbt.macros import MacroConfig, ref_method, source_method, var_method
 from sqlmesh.dbt.seed import SeedConfig
@@ -108,6 +120,9 @@ def _validate_grants(cls, v: t.Dict[str, str]) -> t.Dict[str, t.List[str]]:
 
     @validator("columns", pre=True)
     def _validate_columns(cls, v: t.Any) -> t.Dict[str, ColumnConfig]:
+        if not isinstance(v, dict) or all(isinstance(col, ColumnConfig) for col in v.values()):
+            return v
+
         return yaml_to_columns(v)
 
     _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {
@@ -132,22 +147,6 @@ def to_sqlmesh(
         macros: UniqueKeyDict[str, MacroConfig],
     ) -> Model:
         """Converts the dbt model into a SQLMesh model."""
-        expressions = d.parse_model(
-            f"""
-            MODEL (
-              name {self.model_name},
-              kind {self.model_kind},
-            );
-            """
-            + self.sql,
-            default_dialect="",
-        )
-
-        for jinja in expressions[1:]:
-            # find all the refs here and filter the python env?
-            if isinstance(jinja, d.Jinja):
-                pass
-
         source_mapping = {config.config_name: config.source_name for config in sources.values()}
         model_mapping = {name: config.model_name for name, config in models.items()}
         model_mapping.update({name: config.seed_name for name, config in seeds.items()})
@@ -166,9 +165,15 @@ def to_sqlmesh(
             for ref in dependencies.refs
         }
 
-        return load_model(
-            expressions,
-            path=self.path,
+        expressions = d.parse(self.sql)
+
+        return create_sql_model(
+            self.model_name,
+            expressions[-1],
+            kind=self.model_kind,
+            statements=expressions[0:-1],
+            columns=column_types_to_sqlmesh(self.columns) or None,
+            column_descriptions_=column_descriptions_to_sqlmesh(self.columns) or None,
             python_env=python_env,
             depends_on=depends_on,
             start=self.start,
@@ -186,7 +191,7 @@ def model_name(self) -> str:
         return ".".join(part for part in (schema, self.alias or self.table_name) if part)
 
     @property
-    def model_kind(self) -> str:
+    def model_kind(self) -> ModelKind:
         """
         Get the sqlmesh ModelKind
 
@@ -195,20 +200,20 @@ def model_kind(self) -> str:
         """
         materialization = self.materialized
         if materialization == Materialization.TABLE:
-            return ModelKindName.FULL.value
+            return ModelKind(name=ModelKindName.FULL)
         if materialization == Materialization.VIEW:
-            return ModelKindName.VIEW.value
+            return ModelKind(name=ModelKindName.VIEW)
         if materialization == Materialization.INCREMENTAL:
             if self.time_column:
-                return f"{ModelKindName.INCREMENTAL_BY_TIME_RANGE.value} (TIME_COLUMN {self.time_column})"
+                return IncrementalByTimeRangeKind(time_column=self.time_column)
             if self.unique_key:
-                return f"{ModelKindName.INCREMENTAL_BY_UNIQUE_KEY.value} (UNIQUE_KEY ({','.join(self.unique_key)}))"
+                return IncrementalByUniqueKeyKind(unique_key=self.unique_key)
             raise ConfigError(
                 "SQLMesh ensures idempotent incremental loads and thus does not support append."
                 " Add either an unique key (merge) or a time column (insert-overwrite)."
             )
         if materialization == Materialization.EPHEMERAL:
-            return ModelKindName.EMBEDDED.value
+            return ModelKind(name=ModelKindName.EMBEDDED)
         raise ConfigError(f"{materialization.value} materialization not supported.")
 
     def _all_dependencies(

sqlmesh/dbt/seed.py

@@ -8,7 +8,12 @@
 
 from sqlmesh.core.config.base import UpdateStrategy
 from sqlmesh.core.model import Model, SeedKind, create_seed_model
-from sqlmesh.dbt.column import ColumnConfig, yaml_to_columns
+from sqlmesh.dbt.column import (
+    ColumnConfig,
+    column_descriptions_to_sqlmesh,
+    column_types_to_sqlmesh,
+    yaml_to_columns,
+)
 from sqlmesh.dbt.common import GeneralConfig
 from sqlmesh.utils.conversions import ensure_bool
 
@@ -66,6 +71,9 @@ def _validate_grants(cls, v: t.Dict[str, str]) -> t.Dict[str, t.List[str]]:
 
     @validator("columns", pre=True)
     def _validate_columns(cls, v: t.Any) -> t.Dict[str, ColumnConfig]:
+        if not isinstance(v, dict) or all(isinstance(col, ColumnConfig) for col in v.values()):
+            return v
+
         return yaml_to_columns(v)
 
     _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {
@@ -82,7 +90,11 @@ def _validate_columns(cls, v: t.Any) -> t.Dict[str, ColumnConfig]:
     def to_sqlmesh(self) -> Model:
         """Converts the dbt seed into a SQLMesh model."""
         return create_seed_model(
-            self.seed_name, SeedKind(path=self.path.absolute()), path=self.path
+            self.seed_name,
+            SeedKind(path=self.path.absolute()),
+            path=self.path,
+            columns=column_types_to_sqlmesh(self.columns) or None,
+            column_descriptions_=column_descriptions_to_sqlmesh(self.columns) or None,
         )
 
     @property

sqlmesh/magics.py

@@ -9,7 +9,7 @@
 
 from sqlmesh.core.console import get_console
 from sqlmesh.core.context import Context
-from sqlmesh.core.dialect import format_model_expressions, parse_model
+from sqlmesh.core.dialect import format_model_expressions, parse
 from sqlmesh.core.model import load_model
 from sqlmesh.core.test import ModelTestMetadata, get_all_model_tests
 from sqlmesh.utils.errors import MagicError, MissingContextException, SQLMeshError
@@ -69,7 +69,7 @@ def model(self, line: str, sql: t.Optional[str] = None) -> None:
 
         if sql:
             loaded = load_model(
-                parse_model(sql, default_dialect=self._context.dialect),
+                parse(sql, default_dialect=self._context.dialect),
                 macros=self._context._macros,
                 path=model._path,
                 dialect=self._context.dialect,