Fix: Invalidate the in-memory cache for the columns-to-types mapping when updating the model's schema (#971)

Author: izeigerman

sqlmesh/core/loader.py

@@ -12,10 +12,9 @@
 from pathlib import Path
 
 from ruamel.yaml import YAML
-from sqlglot import exp
 from sqlglot.errors import SqlglotError
 from sqlglot.optimizer.qualify_columns import validate_qualify_columns
-from sqlglot.schema import MappingSchema, nested_set
+from sqlglot.schema import MappingSchema
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.audit import Audit
@@ -49,20 +48,8 @@ def update_model_schemas(dag: DAG[str], models: UniqueKeyDict[str, Model]) -> No
         if not model:
             continue
 
-        external = False
-
-        for dep in model.depends_on:
-            external = external or dep not in models
-            table = exp.to_table(dep, dialect=model.dialect)
-            mapping_schema = schema.find(table)
-
-            if mapping_schema:
-                nested_set(
-                    model.mapping_schema,
-                    tuple(str(part) for part in table.parts),
-                    {k: str(v) for k, v in mapping_schema.items()},
-                )
-
+        external = any(dep not in models for dep in model.depends_on)
+        model.update_schema(schema)
         schema.add_table(name, model.columns_to_types, dialect=model.dialect)
 
         if external:

sqlmesh/core/model/definition.py

@@ -16,6 +16,7 @@
 from sqlglot.diff import Insert, Keep
 from sqlglot.helper import ensure_list
 from sqlglot.optimizer.scope import traverse_scope
+from sqlglot.schema import MappingSchema, nested_set
 from sqlglot.time import format_time
 
 from sqlmesh.core import constants as c
@@ -405,6 +406,19 @@ def convert_to_time_column(self, time: TimeLike) -> exp.Expression:
                 return exp.cast(exp.Literal.string(time), time_column_type)
         return exp.convert(time)
 
+    def update_schema(self, schema: MappingSchema) -> None:
+        """Updates the schema for this model's dependencies based on the given mapping schema."""
+        for dep in self.depends_on:
+            table = exp.to_table(dep, dialect=self.dialect)
+            mapping_schema = schema.find(table)
+
+            if mapping_schema:
+                nested_set(
+                    self.mapping_schema,
+                    tuple(str(part) for part in table.parts),
+                    {k: str(v) for k, v in mapping_schema.items()},
+                )
+
     @property
     def depends_on(self) -> t.Set[str]:
         """All of the upstream dependencies referenced in the model's query, excluding self references.
@@ -701,6 +715,10 @@ def column_descriptions(self) -> t.Dict[str, str]:
             }
         return self._column_descriptions
 
+    def update_schema(self, schema: MappingSchema) -> None:
+        super().update_schema(schema)
+        self._columns_to_types = None
+
     def validate_definition(self) -> None:
         query = self._query_renderer.render()