Repository: TobikoData/sqlmesh (python)

Issue:
`sqlmesh format` strips comment
```sql
SELECT
  *
FROM underlying
WHERE
  @MY_CUSTOM_MACRO() /* this comment gets removed by sqlmesh format */
```

Pull request description:
fix: macro function comments closes #3088

Patch:
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index d3d26139..d7a0cc05 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -184,6 +184,8 @@ Options:
                               trailing.
   --max-text-width INTEGER    The max number of characters in a segment before
                               creating new lines in pretty mode.
+  --check                     Whether or not to check formatting (but not
+                              actually format anything).
   --help                      Show this message and exit.
 ```
 
diff --git a/docs/reference/notebook.md b/docs/reference/notebook.md
index 02be8a84..b000d1d5 100644
--- a/docs/reference/notebook.md
+++ b/docs/reference/notebook.md
@@ -438,3 +438,32 @@ options:
   --read READ    The input dialect of the sql string.
   --write WRITE  The output dialect of the sql string.
 ```
+
+#### format
+```
+%format [--transpile TRANSPILE] [--append-newline] [--normalize]
+              [--pad PAD] [--indent INDENT]
+              [--normalize-functions NORMALIZE_FUNCTIONS] [--leading-comma]
+              [--max-text-width MAX_TEXT_WIDTH] [--check]
+
+Format all SQL models and audits.
+
+options:
+  --transpile TRANSPILE, -t TRANSPILE
+                        Transpile project models to the specified dialect.
+  --append-newline      Whether or not to append a newline to the end of the
+                        file.
+  --normalize           Whether or not to normalize identifiers to lowercase.
+  --pad PAD             Determines the pad size in a formatted string.
+  --indent INDENT       Determines the indentation size in a formatted string.
+  --normalize-functions NORMALIZE_FUNCTIONS
+                        Whether or not to normalize all function names.
+                        Possible values are: 'upper', 'lower'
+  --leading-comma       Determines whether or not the comma is leading or
+                        trailing in select expressions. Default is trailing.
+  --max-text-width MAX_TEXT_WIDTH
+                        The max number of characters in a segment before
+                        creating new lines in pretty mode.
+  --check               Whether or not to check formatting (but not actually
+                        format anything).
+```
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index dff6f194..29daf997 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -255,12 +255,19 @@ def evaluate(
     type=int,
     help="The max number of characters in a segment before creating new lines in pretty mode.",
 )
+@click.option(
+    "--check",
+    is_flag=True,
+    help="Whether or not to check formatting (but not actually format anything).",
+    default=None,
+)
 @click.pass_context
 @error_handler
 @cli_analytics
 def format(ctx: click.Context, **kwargs: t.Any) -> None:
     """Format all SQL models and audits."""
-    ctx.obj.format(**{k: v for k, v in kwargs.items() if v is not None})
+    if not ctx.obj.format(**{k: v for k, v in kwargs.items() if v is not None}):
+        ctx.exit(1)
 
 
 @cli.command("diff")
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index d3ed084a..add5f099 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -863,17 +863,18 @@ class GenericContext(BaseContext, t.Generic[C]):
         self,
         transpile: t.Optional[str] = None,
         append_newline: t.Optional[bool] = None,
+        *,
+        check: t.Optional[bool] = None,
         **kwargs: t.Any,
-    ) -> None:
+    ) -> bool:
         """Format all SQL models and audits."""
         format_targets = {**self._models, **self._audits}
         for target in format_targets.values():
             if not target._path.suffix == ".sql":
                 continue
             with open(target._path, "r+", encoding="utf-8") as file:
-                expressions = parse(
-                    file.read(), default_dialect=self.config_for_node(target).dialect
-                )
+                before = file.read()
+                expressions = parse(before, default_dialect=self.config_for_node(target).dialect)
                 if transpile and is_meta_expression(expressions[0]):
                     for prop in expressions[0].expressions:
                         if prop.name.lower() == "dialect":
@@ -885,15 +886,18 @@ class GenericContext(BaseContext, t.Generic[C]):
                             )
                 format = self.config_for_node(target).format
                 opts = {**format.generator_options, **kwargs}
-                file.seek(0)
-                file.write(
-                    format_model_expressions(expressions, transpile or target.dialect, **opts)
-                )
+                after = format_model_expressions(expressions, transpile or target.dialect, **opts)
                 if append_newline is None:
                     append_newline = format.append_newline
                 if append_newline:
-                    file.write("\n")
-                file.truncate()
+                    after += "\n"
+                if not check:
+                    file.seek(0)
+                    file.write(after)
+                    file.truncate()
+                elif before != after:
+                    return False
+        return True
 
     @python_api_analytics
     def plan(
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index 1f4dcb69..77b569cb 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -584,8 +584,10 @@ def _macro_func_sql(self: Generator, expression: MacroFunc) -> str:
     expression = expression.this
     name = expression.name
     if name in KEYWORD_MACROS:
-        return _macro_keyword_func_sql(self, expression)
-    return f"@{name}({self.format_args(*expression.expressions)})"
+        sql = _macro_keyword_func_sql(self, expression)
+    else:
+        sql = f"@{name}({self.format_args(*expression.expressions)})"
+    return self.maybe_comment(sql, expression)
 
 
 def _override(klass: t.Type[Tokenizer | Parser], func: t.Callable) -> None:
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index c2fe264f..d5fb095b 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -735,12 +735,18 @@ class SQLMeshMagics(Magics):
         type=int,
         help="The max number of characters in a segment before creating new lines in pretty mode.",
     )
+    @argument(
+        "--check",
+        action="store_true",
+        help="Whether or not to check formatting (but not actually format anything).",
+        default=None,
+    )
     @line_magic
     @pass_sqlmesh_context
-    def format(self, context: Context, line: str) -> None:
+    def format(self, context: Context, line: str) -> bool:
         """Format all SQL models and audits."""
         args = parse_argstring(self.format, line)
-        context.format(**{k: v for k, v in vars(args).items() if v is not None})
+        return context.format(**{k: v for k, v in vars(args).items() if v is not None})
 
     @magic_arguments()
     @argument("environment", type=str, help="The environment to diff local state against.")


Test patch:
diff --git a/tests/core/test_dialect.py b/tests/core/test_dialect.py
index 6c5a6da2..edbfdd68 100644
--- a/tests/core/test_dialect.py
+++ b/tests/core/test_dialect.py
@@ -207,7 +207,9 @@ def test_format_body_macros():
     Model ( name foo );
     @WITH(TRUE) x AS (SELECT 1)
     SELECT col::int
-    FROM foo @ORDER_BY(@include_order_by)
+    FROM foo
+    WHERE @MY_MACRO() /* my macro comment */
+@ORDER_BY(@include_order_by)
     @EACH( @columns,
     item -> @'@iteaoeuatnoehutoenahuoanteuhonateuhaoenthuaoentuhaeotnhaoem'),      @'@foo'
     """
@@ -224,6 +226,8 @@ def test_format_body_macros():
 SELECT
   col::INT
 FROM foo
+WHERE
+  @MY_MACRO() /* my macro comment */
 @ORDER_BY(@include_order_by)
   @EACH(@columns, item -> @'@iteaoeuatnoehutoenahuoanteuhonateuhaoenthuaoentuhaeotnhaoem'),
   @'@foo'"""
diff --git a/tests/core/test_format.py b/tests/core/test_format.py
index 907e0b4f..a75d8a5e 100644
--- a/tests/core/test_format.py
+++ b/tests/core/test_format.py
@@ -53,9 +53,14 @@ def test_format_files(tmp_path: pathlib.Path):
         == 'SELECT * FROM @this_model WHERE "CaseSensitive_item_id" < 0'
     )
 
+    assert not context.format(check=True)
+
     # Transpile project to BigQuery
     context.format(transpile="bigquery")
 
+    # Ensure format check is successful
+    assert context.format(transpile="bigquery", check=True)
+
     # Ensure transpilation success AND model specific dialect is mutated
     upd1 = f1.read_text(encoding="utf-8")
     assert (