Fix: Allow deserialization of models using custom materializations (#3886)

Author: erindru

docs/guides/custom_materializations.md

@@ -167,7 +167,6 @@ In many cases, the above usage of a custom materialization will suffice.
 
 However, you may still want tighter integration with SQLMesh's internals:
 
-- You may want more control over what is considered a metadata change vs a data change
 - You may want to validate custom properties are correct before any database connections are made
 - You may want to leverage existing functionality of SQLMesh that relies on specific properties being present
 
@@ -176,43 +175,52 @@ During project load, SQLMesh will instantiate your *subclass* instead of `Custom
 
 This allows you to run custom validators at load time rather than having to perform extra validation when `insert()` is invoked on your `CustomMaterialization`.
 
-This approach also allows you set "top-level" properties directly in the `kind (...)` block rather than nesting them under `materialization_properties`.
+You can also define standard Python `@property` methods to "hoist" properties declared inside `materialization_properties` to the top level on your `Kind` object. This can make using them from within your custom materialization easier.
 
 To extend `CustomKind`, first you define a subclass like so:
 
 ```python linenums="1" hl_lines="7"
-from sqlmesh import CustomKind
+from typing_extensions import Self
 from pydantic import field_validator, ValidationInfo
+from sqlmesh import CustomKind
 from sqlmesh.utils.pydantic import list_of_fields_validator
+from sqlmesh.utils.errors import ConfigError
 
 class MyCustomKind(CustomKind):
 
-    primary_key: t.List[exp.Expression]
+    _primary_key: t.List[exp.Expression]
 
-    @field_validator("primary_key", mode="before")
-    @classmethod
-    def _validate_primary_key(cls, value: t.Any, info: ValidationInfo) -> t.Any:
-        return list_of_fields_validator(value, info.data)
+    @model_validator(mode="after")
+    def _validate_model(self) -> Self:
+        self._primary_key = list_of_fields_validator(
+            self.materialization_properties.get("primary_key"),
+            { "dialect": self.dialect }
+        )
+        if not self.primary_key:
+            raise ConfigError("primary_key must be specified")
+        return self
 
-```
+    @property
+    def primary_key(self) -> t.List[exp.Expression]:
+        return self._primary_key
 
-In this example, we define a field called `primary_key` that takes a list of fields. Notice that the field validation is just a simple Pydantic `@field_validator` with the [exact same usage](https://github.com/TobikoData/sqlmesh/blob/ade5f7245950822f3cfe5a68a0c243f91ceca600/sqlmesh/core/model/kind.py#L470) as the standard SQLMesh model kinds.
+```
 
 To use it within a model, we can do something like:
 
-```sql linenums="1" hl_lines="5"
+```sql linenums="1" hl_lines="4"
 MODEL (
   name my_db.my_model,
   kind CUSTOM (
     materialization 'my_custom_full',
-    primary_key (col1, col2)
+    materialization_properties (
+        primary_key = (col1, col2)
+    )
   )
 );
 ```
 
-Notice that the `primary_key` field we declared is top-level within the `kind` block instead of being nested under `materialization_properties`.
-
-To indicate to SQLMesh that it should use this subclass, specify it as a generic type parameter on your custom materialization class like so:
+To indicate to SQLMesh that it should use the `MyCustomKind` subclass instead of `CustomKind`, specify it as a generic type parameter on your custom materialization class like so:
 
 ```python linenums="1" hl_lines="1 16"
 class CustomFullMaterialization(CustomMaterialization[MyCustomKind]):
@@ -238,21 +246,9 @@ When SQLMesh loads your custom materialization, it will inspect the Python type
 
 In this example, this means that:
 
-- Validation for `primary_key` happens at load time instead of evaluation time.
+- Validation for `primary_key` happens at load time instead of evaluation time. So if there is an issue, you can abort early rather than halfway through applying a plan.
 - When your custom materialization is called to load data into tables, `model.kind` will resolve to your custom kind object so you can access the extra properties you defined without first needing to validate them / coerce them to a usable type.
 
-### Data vs Metadata changes
-
-Subclasses of `CustomKind` that add extra properties can also decide if they are data properties (changes may trigger the creation of new snapshots) or metadata properties (changes just update metadata about the model).
-
-They can also decide if they are relevant for text diffing when SQLMesh detects changes to a model.
-
-You can opt in to SQLMesh's change tracking by overriding the following methods:
-
- - If changing the property should change the data fingerprint, add it to [data_hash_values()](https://github.com/TobikoData/sqlmesh/blob/ade5f7245950822f3cfe5a68a0c243f91ceca600/sqlmesh/core/model/kind.py#L858)
- - If changing the property should change the metadata fingerprint, add it to [metadata_hash_values()](https://github.com/TobikoData/sqlmesh/blob/ade5f7245950822f3cfe5a68a0c243f91ceca600/sqlmesh/core/model/kind.py#L867)
- - If the property should show up in context diffs, add it to [to_expression()](https://github.com/TobikoData/sqlmesh/blob/ade5f7245950822f3cfe5a68a0c243f91ceca600/sqlmesh/core/model/kind.py#L880)
-
 
 ## Sharing custom materializations
 

examples/custom_materializations/custom_materializations/custom_kind.py

@@ -4,19 +4,15 @@
 
 from sqlmesh import CustomMaterialization, CustomKind, Model
 from sqlmesh.utils.pydantic import validate_string
-from pydantic import field_validator
 
 if t.TYPE_CHECKING:
     from sqlmesh import QueryOrDF
 
 
 class ExtendedCustomKind(CustomKind):
-    custom_property: t.Optional[str] = None
-
-    @field_validator("custom_property", mode="before")
-    @classmethod
-    def _validate_custom_property(cls, v: t.Any) -> str:
-        return validate_string(v)
+    @property
+    def custom_property(self) -> str:
+        return validate_string(self.materialization_properties.get("custom_property"))
 
 
 class CustomFullWithCustomKindMaterialization(CustomMaterialization[ExtendedCustomKind]):

examples/sushi/models/latest_order.sql

@@ -2,7 +2,9 @@ MODEL (
   name sushi.latest_order,
   kind CUSTOM (
     materialization 'custom_full_with_custom_kind',
-    custom_property 'sushi!!!'
+    materialization_properties (
+      custom_property = 'sushi!!!'
+    )
   ),
   cron '@daily'
 );

pytest.ini

@@ -39,6 +39,8 @@ markers =
     trino_delta: test for Trino (Delta connector)
 addopts = -n 0 --dist=loadgroup
 
+asyncio_default_fixture_loop_scope = session
+
 # Set this to True to enable logging during tests
 log_cli = False
 log_cli_format = %(asctime)s.%(msecs)03d %(filename)s:%(lineno)d %(levelname)s %(message)s

sqlmesh/core/model/definition.py

@@ -32,7 +32,14 @@
     single_value_or_tuple,
 )
 from sqlmesh.core.model.meta import ModelMeta, FunctionCall
-from sqlmesh.core.model.kind import ModelKindName, SeedKind, ModelKind, FullKind, create_model_kind
+from sqlmesh.core.model.kind import (
+    ModelKindName,
+    SeedKind,
+    ModelKind,
+    FullKind,
+    create_model_kind,
+    CustomKind,
+)
 from sqlmesh.core.model.seed import CsvSeedReader, Seed, create_seed
 from sqlmesh.core.renderer import ExpressionRenderer, QueryRenderer
 from sqlmesh.core.signal import SignalRegistry
@@ -979,6 +986,12 @@ def validate_definition(self) -> None:
                     self._path,
                 )
 
+        if isinstance(self.kind, CustomKind):
+            from sqlmesh.core.snapshot.evaluator import get_custom_materialization_type_or_raise
+
+            # Will raise if the custom materialization points to an invalid class
+            get_custom_materialization_type_or_raise(self.kind.materialization)
+
     def is_breaking_change(self, previous: Model) -> t.Optional[bool]:
         """Determines whether this model is a breaking change in relation to the `previous` model.
 

sqlmesh/core/model/kind.py

@@ -2,6 +2,7 @@
 
 import typing as t
 from enum import Enum
+from typing_extensions import Self
 
 from pydantic import Field
 from sqlglot import exp
@@ -240,43 +241,7 @@ class TimeColumn(PydanticModel):
     @classmethod
     def validator(cls) -> classmethod:
         def _time_column_validator(v: t.Any, info: ValidationInfo) -> TimeColumn:
-            dialect = get_dialect(info.data)
-
-            if isinstance(v, exp.Tuple):
-                column_expr = v.expressions[0]
-                column = (
-                    exp.column(column_expr)
-                    if isinstance(column_expr, exp.Identifier)
-                    else column_expr
-                )
-                format = v.expressions[1].name if len(v.expressions) > 1 else None
-            elif isinstance(v, exp.Expression):
-                column = exp.column(v) if isinstance(v, exp.Identifier) else v
-                format = None
-            elif isinstance(v, str):
-                column = d.parse_one(v, dialect=dialect)
-                column.meta.pop("sql")
-                format = None
-            elif isinstance(v, dict):
-                column_raw = v["column"]
-                column = (
-                    d.parse_one(column_raw, dialect=dialect)
-                    if isinstance(column_raw, str)
-                    else column_raw
-                )
-                format = v.get("format")
-            elif isinstance(v, TimeColumn):
-                column = v.column
-                format = v.format
-            else:
-                raise ConfigError(f"Invalid time_column: '{v}'.")
-
-            column = quote_identifiers(
-                normalize_identifiers(column, dialect=dialect), dialect=dialect
-            )
-            column.meta["dialect"] = dialect
-
-            return TimeColumn(column=column, format=format)
+            return TimeColumn.create(v, get_dialect(info.data))
 
         return field_validator("time_column", mode="before")(_time_column_validator)
 
@@ -314,6 +279,40 @@ def to_expression(self, dialect: str) -> exp.Expression:
     def to_property(self, dialect: str = "") -> exp.Property:
         return exp.Property(this="time_column", value=self.to_expression(dialect))
 
+    @classmethod
+    def create(cls, v: t.Any, dialect: str) -> Self:
+        if isinstance(v, exp.Tuple):
+            column_expr = v.expressions[0]
+            column = (
+                exp.column(column_expr) if isinstance(column_expr, exp.Identifier) else column_expr
+            )
+            format = v.expressions[1].name if len(v.expressions) > 1 else None
+        elif isinstance(v, exp.Expression):
+            column = exp.column(v) if isinstance(v, exp.Identifier) else v
+            format = None
+        elif isinstance(v, str):
+            column = d.parse_one(v, dialect=dialect)
+            column.meta.pop("sql")
+            format = None
+        elif isinstance(v, dict):
+            column_raw = v["column"]
+            column = (
+                d.parse_one(column_raw, dialect=dialect)
+                if isinstance(column_raw, str)
+                else column_raw
+            )
+            format = v.get("format")
+        elif isinstance(v, TimeColumn):
+            column = v.column
+            format = v.format
+        else:
+            raise ConfigError(f"Invalid time_column: '{v}'.")
+
+        column = quote_identifiers(normalize_identifiers(column, dialect=dialect), dialect=dialect)
+        column.meta["dialect"] = dialect
+
+        return cls(column=column, format=format)
+
 
 def _kind_dialect_validator(cls: t.Type, v: t.Optional[str]) -> str:
     if v is None:
@@ -836,17 +835,16 @@ class CustomKind(_ModelKind):
     auto_restatement_cron: t.Optional[SQLGlotCron] = None
     auto_restatement_intervals: t.Optional[SQLGlotPositiveInt] = None
 
+    # so that CustomKind subclasses know the dialect when validating / normalizing / interpreting values in `materialization_properties`
+    dialect: str = Field(exclude=True)
+
     _properties_validator = properties_validator
 
     @field_validator("materialization", mode="before")
     @classmethod
     def _validate_materialization(cls, v: t.Any) -> str:
-        from sqlmesh.core.snapshot.evaluator import get_custom_materialization_type
-
-        materialization = validate_string(v)
-        # The below call fails if a materialization with the given name doesn't exist.
-        get_custom_materialization_type(materialization)
-        return materialization
+        # note: create_model_kind() validates the custom materialization class
+        return validate_string(v)
 
     @property
     def materialization_properties(self) -> CustomMaterializationProperties:
@@ -985,11 +983,15 @@ def create_model_kind(v: t.Any, dialect: str, defaults: t.Dict[str, t.Any]) -> M
                     "The 'materialization' property is required for models of the CUSTOM kind"
                 )
 
-            actual_kind_type, _ = get_custom_materialization_type(
-                validate_string(props.get("materialization"))
-            )
-
-            return actual_kind_type(**props)
+            # The below call will print a warning if a materialization with the given name doesn't exist
+            # we dont want to throw an error here because we still want Models with a CustomKind to be able
+            # to be serialized / deserialized in contexts where the custom materialization class may not be available,
+            # such as in HTTP request handlers
+            if custom_materialization := get_custom_materialization_type(
+                validate_string(props.get("materialization")), raise_errors=False
+            ):
+                actual_kind_type, _ = custom_materialization
+                return actual_kind_type(**props)
 
         return kind_type(**props)