Repository: TobikoData/sqlmesh
Pull request #2527 (2024-04-27 01:58:56)

fix!: add implicit dialect to pydantic models in python
python models allow users to create various pydantic models that contain sql queries. these sql queries are by default parsed as sqlglot dialect which has unintended consequences.

we now set the config default dialect while loading these python models which allows the pydantic validators to set the expected dialect

tobymao: will add more tests later

Patch:
diff --git a/sqlmesh/core/loader.py b/sqlmesh/core/loader.py
index 3ba600cc..856ef668 100644
--- a/sqlmesh/core/loader.py
+++ b/sqlmesh/core/loader.py
@@ -328,27 +328,33 @@ class SqlMeshLoader(Loader):
 
         for context_path, config in self._context.configs.items():
             variables = self._variables(config)
-            for path in self._glob_paths(context_path / c.MODELS, config=config, extension=".py"):
-                if not os.path.getsize(path):
-                    continue
+            model_registry._dialect = config.model_defaults.dialect
+            try:
+                for path in self._glob_paths(
+                    context_path / c.MODELS, config=config, extension=".py"
+                ):
+                    if not os.path.getsize(path):
+                        continue
 
-                self._track_file(path)
-                import_python_file(path, context_path)
-                new = registry.keys() - registered
-                registered |= new
-                for name in new:
-                    model = registry[name].model(
-                        path=path,
-                        module_path=context_path,
-                        defaults=config.model_defaults.dict(),
-                        dialect=config.model_defaults.dialect,
-                        time_column_format=config.time_column_format,
-                        physical_schema_override=config.physical_schema_override,
-                        project=config.project,
-                        default_catalog=self._context.default_catalog,
-                        variables=variables,
-                    )
-                    models[model.fqn] = model
+                    self._track_file(path)
+                    import_python_file(path, context_path)
+                    new = registry.keys() - registered
+                    registered |= new
+                    for name in new:
+                        model = registry[name].model(
+                            path=path,
+                            module_path=context_path,
+                            defaults=config.model_defaults.dict(),
+                            dialect=config.model_defaults.dialect,
+                            time_column_format=config.time_column_format,
+                            physical_schema_override=config.physical_schema_override,
+                            project=config.project,
+                            default_catalog=self._context.default_catalog,
+                            variables=variables,
+                        )
+                        models[model.fqn] = model
+            finally:
+                model_registry._dialect = None
 
         return models
 
diff --git a/sqlmesh/core/model/decorator.py b/sqlmesh/core/model/decorator.py
index ae189e1a..3320fb35 100644
--- a/sqlmesh/core/model/decorator.py
+++ b/sqlmesh/core/model/decorator.py
@@ -4,6 +4,7 @@ import typing as t
 from pathlib import Path
 
 from sqlglot import exp
+from sqlglot.dialects.dialect import DialectType
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.dialect import MacroFunc
@@ -17,6 +18,7 @@ class model(registry_decorator):
     """Specifies a function is a python based model."""
 
     registry_name = "python_models"
+    _dialect: DialectType = None
 
     def __init__(self, name: str, is_sql: bool = False, **kwargs: t.Any) -> None:
         if not name:
diff --git a/sqlmesh/core/model/kind.py b/sqlmesh/core/model/kind.py
index 4f672821..c5e01930 100644
--- a/sqlmesh/core/model/kind.py
+++ b/sqlmesh/core/model/kind.py
@@ -26,6 +26,7 @@ from sqlmesh.utils.pydantic import (
     column_validator,
     field_validator,
     field_validator_v1_args,
+    get_dialect,
 )
 
 if sys.version_info >= (3, 9):
@@ -177,8 +178,8 @@ class TimeColumn(PydanticModel):
     @classmethod
     def validator(cls) -> classmethod:
         def _time_column_validator(v: t.Any, values: t.Any) -> TimeColumn:
-            values = values if isinstance(values, dict) else values.data
-            dialect = values.get("dialect")
+            dialect = get_dialect(values)
+
             if isinstance(v, exp.Tuple):
                 column_expr = v.expressions[0]
                 column = (
@@ -203,7 +204,8 @@ class TimeColumn(PydanticModel):
                 )
                 format = v.get("format")
             elif isinstance(v, TimeColumn):
-                return v
+                column = v.column
+                format = v.format
             else:
                 raise ConfigError(f"Invalid time_column: '{v}'.")
 
@@ -251,14 +253,27 @@ class TimeColumn(PydanticModel):
         return exp.Property(this="time_column", value=self.to_expression(dialect))
 
 
+def _kind_dialect_validator(cls: t.Type, v: t.Optional[str]) -> str:
+    if v is None:
+        return get_dialect({})
+    return v
+
+
+kind_dialect_validator = field_validator("dialect", mode="before", always=True)(
+    _kind_dialect_validator
+)
+
+
 class _Incremental(_ModelKind):
-    dialect: str = ""
+    dialect: t.Optional[str] = Field(None, validate_default=True)
     batch_size: t.Optional[SQLGlotPositiveInt] = None
     batch_concurrency: t.Optional[SQLGlotPositiveInt] = None
     lookback: t.Optional[SQLGlotPositiveInt] = None
     forward_only: SQLGlotBool = False
     disable_restatement: SQLGlotBool = False
 
+    _dialect_validator = kind_dialect_validator
+
     @property
     def data_hash_values(self) -> t.List[t.Optional[str]]:
         return [
@@ -322,7 +337,7 @@ class IncrementalByUniqueKeyKind(_Incremental):
             return expression
 
         if isinstance(v, str):
-            return t.cast(exp.When, d.parse_one(v, into=exp.When))
+            return t.cast(exp.When, d.parse_one(v, into=exp.When, dialect=get_dialect(values)))
 
         if not v:
             return v
@@ -411,7 +426,7 @@ class FullKind(_ModelKind):
 
 
 class _SCDType2Kind(_ModelKind):
-    dialect: str = ""
+    dialect: t.Optional[str] = Field(None, validate_default=True)
     unique_key: SQLGlotListOfFields
     valid_from_name: SQLGlotColumn = Field(exp.column("valid_from"), validate_default=True)
     valid_to_name: SQLGlotColumn = Field(exp.column("valid_to"), validate_default=True)
@@ -421,6 +436,7 @@ class _SCDType2Kind(_ModelKind):
     forward_only: SQLGlotBool = True
     disable_restatement: SQLGlotBool = True
 
+    _dialect_validator = kind_dialect_validator
     # Remove once Pydantic 1 is deprecated
     _always_validate_column = field_validator(
         "valid_from_name", "valid_to_name", mode="before", always=True
@@ -432,10 +448,9 @@ class _SCDType2Kind(_ModelKind):
     def _time_data_type_validator(
         cls, v: t.Union[str, exp.Expression], values: t.Any
     ) -> exp.Expression:
-        values = values if isinstance(values, dict) else values.data
         if isinstance(v, exp.Expression) and not isinstance(v, exp.DataType):
             v = v.name
-        dialect = values.get("dialect")
+        dialect = get_dialect(values)
         data_type = exp.DataType.build(v, dialect=dialect)
         data_type.meta["dialect"] = dialect
         return data_type
@@ -552,7 +567,7 @@ def model_kind_type_from_name(name: t.Optional[str]) -> t.Type[ModelKind]:
 
 @field_validator_v1_args
 def _model_kind_validator(cls: t.Type, v: t.Any, values: t.Dict[str, t.Any]) -> ModelKind:
-    dialect = values.get("dialect")
+    dialect = get_dialect(values)
 
     if isinstance(v, _ModelKind):
         return t.cast(ModelKind, v)
diff --git a/sqlmesh/utils/pydantic.py b/sqlmesh/utils/pydantic.py
index f8465ff7..fae6bd35 100644
--- a/sqlmesh/utils/pydantic.py
+++ b/sqlmesh/utils/pydantic.py
@@ -75,6 +75,21 @@ def parse_obj_as(type_: T, obj: t.Any) -> T:
     return pydantic.tools.parse_obj_as(type_, obj)  # type: ignore
 
 
+def get_dialect(values: t.Any) -> str:
+    """Extracts dialect from a dict or pydantic obj, defaulting to the globally set dialect.
+
+    Python models allow users to instantiate pydantic models by hand. This is problematic
+    because the validators kick in with the SQLGLot dialect. To instantiate Pydantic Models used
+    in python models using the project default dialect, we set a class variable on the model
+    registry and use that here.
+    """
+
+    from sqlmesh.core.model import model
+
+    dialect = (values if isinstance(values, dict) else values.data).get("dialect")
+    return model._dialect if dialect is None else dialect
+
+
 def _expression_encoder(e: exp.Expression) -> str:
     return e.meta.get("sql") or e.sql(dialect=e.meta.get("dialect"))
 
@@ -294,8 +309,7 @@ def _get_field(
     v: t.Any,
     values: t.Any,
 ) -> exp.Expression:
-    values = values if isinstance(values, dict) else values.data
-    dialect = values.get("dialect")
+    dialect = get_dialect(values)
 
     if isinstance(v, exp.Expression):
         expression = v
@@ -315,8 +329,7 @@ def _get_fields(
     v: t.Any,
     values: t.Any,
 ) -> t.List[exp.Expression]:
-    values = values if isinstance(values, dict) else values.data
-    dialect = values.get("dialect")
+    dialect = get_dialect(values)
 
     if isinstance(v, (exp.Tuple, exp.Array)):
         expressions: t.List[exp.Expression] = v.expressions


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 70729d83..e0e02a2d 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -3897,3 +3897,43 @@ def test_macros_in_model_statement(sushi_context, assert_exp_eq):
     assert model.time_column
     assert model.time_column.column == exp.column("a", quoted=True)
     assert model.start == "2023-01-01"
+
+
+def test_python_model_dialect():
+    model._dialect = "snowflake"
+
+    @model(
+        name="a",
+        kind=IncrementalByTimeRangeKind(time_column=TimeColumn(column="x", format="YYMMDD")),
+        columns={},
+    )
+    def test(context, **kwargs):
+        return None
+
+    m = model.get_registry()["a"].model(
+        module_path=Path("."),
+        path=Path("."),
+        dialect="snowflake",
+    )
+
+    assert m.time_column.column.sql() == '"X"'
+    assert m.time_column.format == "%y%m%d"
+
+    @model(
+        name="b",
+        kind=IncrementalByTimeRangeKind(time_column="y"),
+        columns={},
+    )
+    def test(context, **kwargs):
+        return None
+
+    m = model.get_registry()["b"].model(
+        module_path=Path("."),
+        path=Path("."),
+        dialect="snowflake",
+    )
+
+    assert m.time_column.column.sql() == '"Y"'
+    assert m.time_column.format == "%Y-%m-%d"
+
+    model._dialect = None
diff --git a/tests/core/test_snapshot.py b/tests/core/test_snapshot.py
index 9be09127..6bd5e610 100644
--- a/tests/core/test_snapshot.py
+++ b/tests/core/test_snapshot.py
@@ -109,7 +109,6 @@ def test_json(snapshot: Snapshot):
                 "batch_size": 30,
                 "forward_only": False,
                 "disable_restatement": False,
-                "dialect": "",
             },
             "mapping_schema": {},
             "start": "2020-01-01",
@@ -730,7 +729,7 @@ def test_fingerprint_jinja_macros(model: Model):
         }
     )
     original_fingerprint = SnapshotFingerprint(
-        data_hash="2664333317",
+        data_hash="2973224250",
         metadata_hash="892368116",
     )