Merge branch 'main' into themis/dbttrace

Author: themisvaltinos

examples/sushi_dbt/models/schema.yml

@@ -36,6 +36,8 @@ models:
               field: waiter_id
       - name: revenue
         description: Revenue from orders served by this waiter
+      - name: unused_column
+        data_type: int
   - name: waiters
     columns:
       - name: waiter_id

examples/sushi_dbt/models/top_waiters.sql

@@ -6,7 +6,8 @@
 
 SELECT
   waiter_id::INT AS waiter_id,
-  revenue::DOUBLE AS revenue
+  revenue::DOUBLE AS revenue,
+  1 AS unused_column
 FROM {{ ref('waiter_revenue_by_day', version=1) }}
 WHERE
   ds = (

sqlmesh/dbt/basemodel.py

@@ -328,12 +328,9 @@ def sqlmesh_model_kwargs(
             dependencies.macros, package=self.package_name
         )
         jinja_macros.add_globals(self._model_jinja_context(model_context, dependencies))
-        return {
+
+        model_kwargs = {
             "audits": [(test.name, {}) for test in self.tests],
-            "columns": column_types_to_sqlmesh(
-                column_types_override or self.columns, self.dialect(context)
-            )
-            or None,
             "column_descriptions": column_descriptions_to_sqlmesh(self.columns) or None,
             "depends_on": {
                 model.canonical_name(context) for model in model_context.refs.values()
@@ -349,6 +346,23 @@ def sqlmesh_model_kwargs(
             **self.sqlmesh_config_kwargs,
         }
 
+        # dbt doesn't respect the data_type field for DDL statements– instead, it optionally uses
+        # it to validate the actual data types at runtime through contracts or external plugins.
+        # Only the `columns_types` config of seed models is actually respected. We don't set the
+        # columns attribute to self.columns intentionally in all other cases, as that could result
+        # in unfaithful types when models are materialized.
+        #
+        # See:
+        # - https://docs.getdbt.com/reference/resource-properties/columns
+        # - https://docs.getdbt.com/reference/resource-configs/contract
+        # - https://docs.getdbt.com/reference/resource-configs/column_types
+        if column_types_override:
+            model_kwargs["columns"] = (
+                column_types_to_sqlmesh(column_types_override, self.dialect(context)) or None
+            )
+
+        return model_kwargs
+
     @abstractmethod
     def to_sqlmesh(
         self,

sqlmesh/dbt/seed.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import copy
 import typing as t
 
 import agate
@@ -50,15 +49,11 @@ def to_sqlmesh(
         """Converts the dbt seed into a SQLMesh model."""
         seed_path = self.path.absolute().as_posix()
 
-        if column_types := self.column_types:
-            column_types_override = copy.deepcopy(self.columns)
-            for name, data_type in column_types.items():
-                column = column_types_override.setdefault(name, ColumnConfig(name=name))
-                column.data_type = data_type
-                column.quote = self.quote_columns or column.quote
-                kwargs = self.sqlmesh_model_kwargs(context, column_types_override)
-        else:
-            kwargs = self.sqlmesh_model_kwargs(context)
+        column_types_override = {
+            name: ColumnConfig(name=name, data_type=data_type, quote=self.quote_columns)
+            for name, data_type in (self.column_types or {}).items()
+        }
+        kwargs = self.sqlmesh_model_kwargs(context, column_types_override)
 
         columns = kwargs.get("columns") or {}
 

sqlmesh/utils/jinja.py

@@ -200,6 +200,14 @@ def _extract(node: nodes.Node, parent: t.Optional[nodes.Node] = None) -> None:
     return extracted
 
 
+def is_variable_node(n: nodes.Node) -> bool:
+    return (
+        isinstance(n, nodes.Call)
+        and isinstance(n.node, nodes.Name)
+        and n.node.name in (c.VAR, c.BLUEPRINT_VAR)
+    )
+
+
 def extract_macro_references_and_variables(
     *jinja_strs: str, dbt_target_name: t.Optional[str] = None
 ) -> t.Tuple[t.Set[MacroReference], t.Set[str]]:
@@ -230,7 +238,15 @@ def extract_macro_references_and_variables(
 
         for call_name, node in extract_call_names(jinja_str):
             if call_name[0] in (c.VAR, c.BLUEPRINT_VAR):
-                assert isinstance(node, nodes.Call)
+                if not is_variable_node(node):
+                    # Find the variable node which could be nested
+                    for n in node.find_all(nodes.Call):
+                        if is_variable_node(n):
+                            node = n
+                            break
+                    else:
+                        raise ValueError(f"Could not find variable name in {jinja_str}")
+                node = t.cast(nodes.Call, node)
                 args = [jinja_call_arg_name(arg) for arg in node.args]
                 if args and args[0]:
                     variable_name = args[0].lower()

tests/dbt/converter/test_jinja.py

@@ -1,5 +1,9 @@
 import pytest
-from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroExtractor
+from sqlmesh.utils.jinja import (
+    JinjaMacroRegistry,
+    MacroExtractor,
+    extract_macro_references_and_variables,
+)
 from sqlmesh.dbt.converter.jinja import JinjaGenerator, convert_jinja_query, convert_jinja_macro
 import sqlmesh.dbt.converter.jinja_transforms as jt
 from pathlib import Path
@@ -437,3 +441,10 @@ def test_convert_jinja_macro(input: str, expected: str, sushi_dbt_context: Conte
     result = convert_jinja_macro(sushi_dbt_context, input.strip())
 
     assert " ".join(result.split()) == " ".join(expected.strip().split())
+
+
+def test_extract_macro_references_and_variables() -> None:
+    input = """JINJA_QUERY('{%- set something = "'"~var("variable").split("|") -%}"""
+    _, variables = extract_macro_references_and_variables(input)
+    assert len(variables) == 1
+    assert variables == {"variable"}

tests/dbt/test_config.py

@@ -7,11 +7,13 @@
 from dbt.adapters.base import BaseRelation, Column
 from pytest_mock import MockerFixture
 
+from sqlglot import exp
 from sqlmesh.core.audit import StandaloneAudit
 from sqlmesh.core.config import Config, ModelDefaultsConfig
 from sqlmesh.core.dialect import jinja_query
 from sqlmesh.core.model import SqlModel
 from sqlmesh.core.model.kind import OnDestructiveChange, OnAdditiveChange
+from sqlmesh.dbt.column import ColumnConfig
 from sqlmesh.dbt.common import Dependencies
 from sqlmesh.dbt.context import DbtContext
 from sqlmesh.dbt.loader import sqlmesh_config
@@ -1076,3 +1078,15 @@ def test_on_schema_change_properties(
 
     assert model.on_additive_change == expected_additive
     assert model.on_destructive_change == expected_destructive
+
+
+def test_sqlmesh_model_kwargs_columns_override():
+    context = DbtContext()
+    context.project_name = "Foo"
+    context.target = DuckDbConfig(name="target", schema="foo")
+
+    kwargs = ModelConfig(dialect="duckdb").sqlmesh_model_kwargs(
+        context,
+        {"c": ColumnConfig(name="c", data_type="uinteger")},
+    )
+    assert kwargs.get("columns") == {"c": exp.DataType.build(exp.DataType.Type.UINT)}

tests/dbt/test_transformation.py

@@ -608,7 +608,10 @@ def test_model_columns():
         name="target", schema="test", database="test", account="foo", user="bar", password="baz"
     )
     sqlmesh_model = model.to_sqlmesh(context)
-    assert sqlmesh_model.columns_to_types == expected_column_types
+
+    # Columns being present in a schema.yaml are not respected in DDLs, so SQLMesh doesn't
+    # set the corresponding columns_to_types_ attribute either to match dbt's behavior
+    assert sqlmesh_model.columns_to_types == None
     assert sqlmesh_model.column_descriptions == expected_column_descriptions
 
 
@@ -623,8 +626,11 @@ def test_seed_columns():
         },
     )
 
+    # dbt doesn't respect the data_type field in the DDLs– instead, it optionally uses it to
+    # validate the actual data types at runtime through contracts or external plugins. Thus,
+    # the actual data type is int, because that is what is inferred from the seed file.
     expected_column_types = {
-        "id": exp.DataType.build("text"),
+        "id": exp.DataType.build("int"),
         "name": exp.DataType.build("text"),
     }
     expected_column_descriptions = {
@@ -671,6 +677,27 @@ def test_seed_column_types():
     assert sqlmesh_seed.columns_to_types == expected_column_types
     assert sqlmesh_seed.column_descriptions == expected_column_descriptions
 
+    seed = SeedConfig(
+        name="foo",
+        package="package",
+        path=Path("examples/sushi_dbt/seeds/waiter_names.csv"),
+        column_types={
+            "name": "text",
+        },
+        columns={
+            # The `data_type` field does not affect the materialized seed's column type
+            "id": ColumnConfig(name="name", data_type="text"),
+        },
+        quote_columns=True,
+    )
+
+    expected_column_types = {
+        "id": exp.DataType.build("int"),
+        "name": exp.DataType.build("text"),
+    }
+    sqlmesh_seed = seed.to_sqlmesh(context)
+    assert sqlmesh_seed.columns_to_types == expected_column_types
+
 
 def test_seed_column_inference(tmp_path):
     seed_csv = tmp_path / "seed.csv"