refactor on method in modelmeta

Author: themisvaltinos

sqlmesh/core/context.py

@@ -1647,25 +1647,9 @@ def table_diff(
                         target = target_snapshot.qualified_view_name.for_environment(
                             target_env.naming_info, adapter.dialect
                         )
-
-                        model_on = []
-                        if not on:
-                            for expr in [
-                                ref.expression for ref in model.all_references if ref.unique
-                            ]:
-                                if isinstance(expr, exp.Tuple):
-                                    model_on.extend(
-                                        [
-                                            key.this.sql(dialect=adapter.dialect)
-                                            for key in expr.expressions
-                                        ]
-                                    )
-                                else:
-                                    # Handle a single Column or Paren expression
-                                    model_on.append(expr.this.sql(dialect=adapter.dialect))
-
-                        models_to_diff.append((model, adapter, source, target, on or model_on))
-                        if not (on or model_on):
+                        model_on = on or model.on
+                        models_to_diff.append((model, adapter, source, target, model_on))
+                        if not model_on:
                             models_without_grain.append(model)
                     else:
                         models_no_diff.append(model_fqn)

sqlmesh/core/model/meta.py

@@ -450,6 +450,20 @@ def all_references(self) -> t.List[Reference]:
             Reference(model_name=self.name, expression=e, unique=True) for e in self.references
         ]
 
+    @property
+    def on(self) -> t.List[str]:
+        """The grains to be used as join condition in table_diff."""
+
+        on: t.List[str] = []
+        for expr in [ref.expression for ref in self.all_references if ref.unique]:
+            if isinstance(expr, exp.Tuple):
+                on.extend([key.this.sql(dialect=self.dialect) for key in expr.expressions])
+            else:
+                # Handle a single Column or Paren expression
+                on.append(expr.this.sql(dialect=self.dialect))
+
+        return on
+
     @property
     def managed_columns(self) -> t.Dict[str, exp.DataType]:
         return getattr(self.kind, "managed_columns", {})