Repository: TobikoData/sqlmesh
Pull request #2208 (2024-03-03 04:50:29)

fix: snowflake column level lineage issue due to normalization


Patch:
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index 1641a8ba..2dc328ed 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -13,6 +13,7 @@ from sqlglot.dialects.dialect import DialectType
 from sqlglot.dialects.snowflake import Snowflake
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.scope import traverse_scope
+from sqlglot.schema import MappingSchema
 from sqlglot.tokens import Token
 
 from sqlmesh.core.constants import MAX_MODEL_DEFINITION_SIZE
@@ -965,3 +966,14 @@ def schema_(
         db=exp.to_identifier(db, quoted=quoted) if db else None,
         catalog=exp.to_identifier(catalog, quoted=quoted) if catalog else None,
     )
+
+
+def normalize_mapping_schema(schema: t.Dict, dialect: DialectType) -> MappingSchema:
+    return MappingSchema(_unquote_schema(schema), dialect=dialect, normalize=False)
+
+
+def _unquote_schema(schema: t.Dict) -> t.Dict:
+    """SQLGlot schema expects unquoted normalized keys."""
+    return {
+        k.strip('"'): _unquote_schema(v) if isinstance(v, dict) else v for k, v in schema.items()
+    }
diff --git a/sqlmesh/core/lineage.py b/sqlmesh/core/lineage.py
index 02cbd01a..837f4d75 100644
--- a/sqlmesh/core/lineage.py
+++ b/sqlmesh/core/lineage.py
@@ -5,9 +5,10 @@ from collections import defaultdict
 
 from sqlglot import exp
 from sqlglot.helper import first
-from sqlglot.lineage import lineage
+from sqlglot.lineage import Node
+from sqlglot.lineage import lineage as sqlglot_lineage
 
-from sqlmesh.core.dialect import normalize_model_name
+from sqlmesh.core.dialect import normalize_mapping_schema, normalize_model_name
 
 if t.TYPE_CHECKING:
     from sqlmesh.core.context import Context
@@ -30,13 +31,28 @@ def _render_query(model: Model) -> exp.Query:
     return query
 
 
+def lineage(
+    column: str | exp.Column,
+    model: Model,
+    **kwargs: t.Any,
+) -> Node:
+    return sqlglot_lineage(
+        column,
+        sql=_render_query(model),
+        schema=normalize_mapping_schema(model.mapping_schema, dialect=model.dialect),
+        dialect=model.dialect,
+        **{
+            "infer_schema": True,
+            **kwargs,
+        },
+    )
+
+
 def column_dependencies(context: Context, model_name: str, column: str) -> t.Dict[str, t.Set[str]]:
     model = context.get_model(model_name)
     parents = defaultdict(set)
 
-    for node in lineage(
-        column, sql=_render_query(model), schema=model.mapping_schema, infer_schema=True
-    ).walk():
+    for node in lineage(column, model).walk():
         if node.downstream:
             continue
 
diff --git a/sqlmesh/core/renderer.py b/sqlmesh/core/renderer.py
index 706acc6f..05e9f4ba 100644
--- a/sqlmesh/core/renderer.py
+++ b/sqlmesh/core/renderer.py
@@ -14,7 +14,6 @@ from sqlglot.optimizer.qualify import qualify
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 from sqlglot.optimizer.qualify_tables import qualify_tables
 from sqlglot.optimizer.simplify import simplify
-from sqlglot.schema import MappingSchema
 
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
@@ -66,7 +65,7 @@ class BaseExpressionRenderer:
         self._cache: t.Dict[CacheKey, t.List[t.Optional[exp.Expression]]] = {}
 
     def update_schema(self, schema: t.Dict[str, t.Any]) -> None:
-        self.schema = MappingSchema(_unquote_schema(schema), dialect=self._dialect, normalize=False)
+        self.schema = d.normalize_mapping_schema(schema, dialect=self._dialect)
 
     def _render(
         self,
@@ -534,10 +533,3 @@ def _normalize_and_quote(query: E, dialect: str, default_catalog: t.Optional[str
     normalize_identifiers(query, dialect=dialect)
     yield query
     quote_identifiers(query, dialect=dialect)
-
-
-def _unquote_schema(schema: t.Dict) -> t.Dict:
-    """SQLGlot schema expects unquoted normalized keys."""
-    return {
-        k.strip('"'): _unquote_schema(v) if isinstance(v, dict) else v for k, v in schema.items()
-    }
diff --git a/sqlmesh/dbt/loader.py b/sqlmesh/dbt/loader.py
index 568c783b..df8da990 100644
--- a/sqlmesh/dbt/loader.py
+++ b/sqlmesh/dbt/loader.py
@@ -41,7 +41,7 @@ def sqlmesh_config(
     project_root = project_root or Path()
     context = DbtContext(project_root=project_root)
     profile = Profile.load(context, target_name=dbt_target_name)
-    model_defaults = kwargs.get("model_defaults", ModelDefaultsConfig())
+    model_defaults = kwargs.pop("model_defaults", ModelDefaultsConfig())
     model_defaults.dialect = profile.target.dialect
 
     loader_kwargs = kwargs.pop("loader_kwargs", {})
diff --git a/web/server/api/endpoints/lineage.py b/web/server/api/endpoints/lineage.py
index 639b5961..982661f7 100644
--- a/web/server/api/endpoints/lineage.py
+++ b/web/server/api/endpoints/lineage.py
@@ -5,11 +5,10 @@ from collections import defaultdict
 
 from fastapi import APIRouter, Depends
 from sqlglot import exp
-from sqlglot.lineage import lineage
 
 from sqlmesh.core.context import Context
 from sqlmesh.core.dialect import normalize_model_name
-from sqlmesh.core.lineage import _render_query, column_dependencies
+from sqlmesh.core.lineage import column_dependencies, lineage
 from web.server.exceptions import ApiException
 from web.server.models import LineageColumn
 from web.server.settings import get_loaded_context
@@ -54,9 +53,7 @@ def create_lineage_adjacency_list(
                 models={},
             )
             continue
-        root = lineage(
-            column, sql=_render_query(model), schema=model.mapping_schema, infer_schema=True
-        )
+        root = lineage(column, model)
 
         for node in root.walk():
             if root.name == "UNION" and node is root:


Test patch:
diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
index 13feed6f..33a95779 100644
--- a/sqlmesh/core/test/definition.py
+++ b/sqlmesh/core/test/definition.py
@@ -137,7 +137,9 @@ class ModelTest(unittest.TestCase):
                 return
 
             diff = expected.compare(actual).rename(columns={"self": "exp", "other": "act"})
-            e.args = (f"Data differs (exp: expected, act: actual)\n\n{diff}",)
+            description = self.body.get("description")
+            description = f"\n\n\nTest description: {description}" if description else ""
+            e.args = (f"Data differs (exp: expected, act: actual)\n\n{diff}{description}",)
             raise e
 
     def runTest(self) -> None:
diff --git a/tests/core/test_lineage.py b/tests/core/test_lineage.py
index 47bdf425..f15abae8 100644
--- a/tests/core/test_lineage.py
+++ b/tests/core/test_lineage.py
@@ -1,4 +1,9 @@
-from sqlmesh.core.lineage import column_dependencies, column_description
+from sqlmesh.core import dialect as d
+from sqlmesh.core.config import Config
+from sqlmesh.core.config.model import ModelDefaultsConfig
+from sqlmesh.core.context import Context
+from sqlmesh.core.lineage import column_dependencies, column_description, lineage
+from sqlmesh.core.model import load_sql_based_model
 
 
 def test_column_dependencies(sushi_context_pre_scheduling):
@@ -12,3 +17,24 @@ def test_column_dependencies(sushi_context_pre_scheduling):
 def test_column_description(sushi_context_pre_scheduling):
     context = sushi_context_pre_scheduling
     assert column_description(context, "sushi.top_waiters", "waiter_id") == "Waiter id"
+
+
+def test_lineage():
+    context = Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="snowflake")))
+
+    model = load_sql_based_model(
+        d.parse(
+            """
+        MODEL (name db.model1);
+
+        SELECT "A"
+        FROM (
+            SELECT 1 a
+        ) x
+        """
+        ),
+    )
+
+    context.upsert_model(model)
+    node = lineage('"A"', model)
+    assert node.name == "A"
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index bb50dcdb..4af267e1 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -410,6 +410,7 @@ def test_missing_column_failure(sushi_context: Context, full_model_without_ctes:
         """
 test_foo:
   model: sushi.foo
+  description: sushi.foo's output has a missing column (fails intentionally)
   inputs:
     raw:
       - id: 1
@@ -424,7 +425,15 @@ test_foo:
     result = _create_test(body, "test_foo", model, sushi_context).run()
     assert result and not result.wasSuccessful()
 
-    expected_msg = "AssertionError: Data differs (exp: expected, act: actual)\n\n  value      ds    \n    exp act exp act\n0   NaN   2 NaN   3\n"
+    expected_msg = """AssertionError: Data differs (exp: expected, act: actual)
+
+  value      ds    
+    exp act exp act
+0   NaN   2 NaN   3
+
+
+Test description: sushi.foo's output has a missing column (fails intentionally)
+"""
     assert expected_msg in result.failures[0][1]