Fix: Support on-virtual-update statements in sqlmesh format (#3998)

Author: themisvaltinos

sqlmesh/core/dialect.py

@@ -644,6 +644,16 @@ def _props_sql(self: Generator, expressions: t.List[exp.Expression]) -> str:
     return "\n".join(props)
 
 
+def _on_virtual_update_sql(self: Generator, expressions: t.List[exp.Expression]) -> str:
+    statements = "\n".join(
+        self.sql(expression)
+        if isinstance(expression, JinjaStatement)
+        else f"{self.sql(expression)};"
+        for expression in expressions
+    )
+    return f"{ON_VIRTUAL_UPDATE_BEGIN};\n{statements}\n{ON_VIRTUAL_UPDATE_END};"
+
+
 def _sqlmesh_ddl_sql(self: Generator, expression: Model | Audit | Metric, name: str) -> str:
     return "\n".join([f"{name} (", _props_sql(self, expression.expressions), ")"])
 
@@ -1004,6 +1014,7 @@ def extend_sqlglot() -> None:
                     JinjaQuery: lambda self, e: f"{JINJA_QUERY_BEGIN};\n{e.name}\n{JINJA_END};",
                     JinjaStatement: lambda self,
                     e: f"{JINJA_STATEMENT_BEGIN};\n{e.name}\n{JINJA_END};",
+                    VirtualUpdateStatement: lambda self, e: _on_virtual_update_sql(self, e),
                     MacroDef: lambda self, e: f"@DEF({self.sql(e.this)}, {self.sql(e.expression)})",
                     MacroFunc: _macro_func_sql,
                     MacroStrReplace: lambda self, e: f"@{self.sql(e.this)}",

tests/core/test_dialect.py

@@ -220,6 +220,43 @@ def test_format_model_expressions():
   CAST(1 AS INT) AS bla"""
     )
 
+    x = format_model_expressions(
+        parse(
+            """MODEL(name foo);
+SELECT CAST(1 AS INT) AS bla;
+            on_virtual_update_begin;
+CREATE OR REPLACE VIEW test_view FROM demo_db.table;GRANT SELECT ON VIEW @this_model TO ROLE owner_name;
+JINJA_STATEMENT_BEGIN; GRANT SELECT ON VIEW {{this_model}} TO ROLE admin;        JINJA_END;
+    GRANT REFERENCES, SELECT ON FUTURE VIEWS IN DATABASE demo_db TO ROLE owner_name;
+@resolve_parent_name('parent');GRANT SELECT ON VIEW demo_db.table /* sqlglot.meta replace=false */ TO ROLE admin;
+ON_VIRTUAL_update_end;"""
+        )
+    )
+
+    assert (
+        x
+        == """MODEL (
+  name foo
+);
+
+SELECT
+  1::INT AS bla;
+
+ON_VIRTUAL_UPDATE_BEGIN;
+CREATE OR REPLACE VIEW test_view AS
+SELECT
+  *
+FROM demo_db.table;
+GRANT SELECT ON VIEW @this_model TO ROLE owner_name;
+JINJA_STATEMENT_BEGIN;
+GRANT SELECT ON VIEW {{this_model}} TO ROLE admin;
+JINJA_END;
+GRANT REFERENCES, SELECT ON FUTURE VIEWS IN DATABASE demo_db TO ROLE owner_name;
+@resolve_parent_name('parent');
+GRANT SELECT ON VIEW demo_db.table /* sqlglot.meta replace=false */ TO ROLE admin;
+ON_VIRTUAL_UPDATE_END;"""
+    )
+
 
 def test_macro_format():
     assert parse_one("@EACH(ARRAY(1,2), x -> x)").sql() == "@EACH(ARRAY(1, 2), x -> x)"