Repository: TobikoData/sqlmesh (python)

Issue:
Improve `sqlmesh format` behaviour
Trying to capture a few concerns in a single place:

## Audits

As it stands, the `audits` property on `MODEL` will get re-written into a single line (more or less). This can be quite problematic for models that have complex or a lot of audits.

For example,

```sql
MODEL(
  audits [
    not_null(columns=[
      foo_id,
      foo_normalised,
      bar_normalised,
      total_weight,
      overall_rank,
      cumulative_total_weight_share,
      market_rank,
      market_cumulative_total_weight_share,
      total_weight_decile,
      tier,
    ]),

    unique_values(columns=[foo_id]),

    accepted_range(column=foo_normalised, min_v=0, max_v=100),
    accepted_range(column=bar_normalised, min_v=0, max_v=100),
    accepted_range(column=total_weight, min_v=0, max_v=100),
    accepted_range(column=cumulative_total_weight_share, min_v=0, max_v=1),
    accepted_range(column=market_cumulative_total_weight_share, min_v=0, max_v=1),

    accepted_values(column=tier, is_in=['Tier 1', 'Tier 2', 'Tier 3', 'Long Tail']),
    accepted_values(column=total_weight_decile, is_in=['Decile_01', 'Decile_02','Decile_03','Decile_04','Decile_05','Decile_06','Decile_07','Decile_08','Decile_09','Decile_10']),
  ],
);
```

Will be re-written to

```sql
MODEL(
  audits [not_null(
    columns = [foo_id, foo_normalised, bar_normalised, total_weight, overall_rank, cumulative_total_weight_share, market_rank, market_cumulative_total_weight_share, total_weight_decile, tier]
  ), unique_values(columns = [foo_id]), accepted_range(column = foo_normalised, min_v = 0, max_v = 100), accepted_range(column = bar_normalised, min_v = 0, max_v = 100), accepted_range(column = total_weight, min_v = 0, max_v = 100), accepted_range(column = cumulative_total_weight_share, min_v = 0, max_v = 1), accepted_range(column = market_cumulative_total_weight_share, min_v = 0, max_v = 1), accepted_values(column = tier, is_in = ['Tier 1', 'Tier 2', 'Tier 3', 'Long Tail']), accepted_values(
    column = total_weight_decile,
    is_in = ['Decile_01', 'Decile_02', 'Decile_03', 'Decile_04', 'Decile_05', 'Decile_06', 'Decile_07', 'Decile_08', 'Decile_09', 'Decile_10']
  )],
);
```

Similarly, we often apply comments to audits to document their purpose, e.g.:

```sql
  audits [
    forall(criteria=[
      /* This is a comment about this complicated audit */
      foo IS NOT NULL OR (
        some_date < '2023-11-30' OR
        something_else IN('a', 'b') OR
        /* See this conversation https://slack.com/archives/ABCD/p777 */
        (some_date BETWEEN '2024-02-24' AND '2024-02-25') OR
        /* See this conversation https://slack.com/archives/ABCD/p777 */
        (some_date BETWEEN '2024-03-14' AND '2024-03-16')
      )
    ]),
    /* ensure the this column is well formed, but don't fail since we expect this to change in time */
    accepted_values_non_blocking(column=hs_tenant, is_in=('foo', 'bar', 'baz')),
    /* ensure that the interval has some data */
    number_of_rows(threshold=100),
  ],
```

this gets re-written as

```sql
  audits [not_null(
    columns = [hs_shopped_on_time_utc, hs_user_id, hs_availability_id, hs_check_in_date, hs_check_out_date, hs_tenant, hs_lodging_id, hs_rate_count, hs_rate_provider_count]
  ), forall(
    criteria = [NOT foo /* This is a comment about this complicated audit */ IS NULL
    OR (
      some_date < '2023-11-30'
      OR something_else IN ('a', 'b')
      OR (
        some_date BETWEEN '2024-02-24' AND '2024-02-25'
      ) /* See this conversation https://slack.com/archives/ABCD/p777 */
      OR (
        some_date BETWEEN '2024-03-14' AND '2024-03-16'
      ) /* See this conversation https://slack.com/archives/ABCD/p777 */
    )]
  ), accepted_values_non_blocking(column = hs_tenant, is_in = ('foo', 'bar', 'baz')) /* ensure the this column is well formed, but don't fail since we expect this to change in time */, number_of_rows(threshold = 100) /* ensure that the interval has some data */],
```

We can probably live without the comments, but it would be important to keep the one-audit-per-line format for maintainability by people.

## Comments

### Stripped away

Some comments get stripped away, for example in this case, the comment will be lost after `sqlmesh format`

```sql
MODEL(name test.test);

/* Define a constant */
@DEF(foo, [a,b,c]);

SELECT @EACH(@foo, c -> c) FROM bar
```

re-wrote to

```sql
MODEL (
  name test.test
);

@DEF(foo, [a, b, c]);

SELECT
  @EACH(@foo, c -> c)
FROM bar
```

### Duplicated

In some cases, they get duplicated:

```sql
MODEL(name test.test);

-- CAREFUL: UNION ALL is **positional**
SELECT a,b,c FROM foo
UNION ALL
SELECT a,b,c FROM bar
```

rewritten to

```sql
MODEL (
  name test.test
);

/* CAREFUL: UNION ALL is **positional** */
SELECT
  a,
  b,
  c
FROM foo
/* CAREFUL: UNION ALL is **positional** */
UNION ALL
SELECT
  a,
  b,
  c
FROM bar
```

### Non-deterministic

Applying `sqlmesh format` will make (whitespace) changes each time:

```sql
MODEL(name test.test);

WITH foo AS(
    SELECT
        a,b
    FROM base
    /*
      NOTE: LEFT JOIN because blah blah blah
    */
    LEFT JOIN other ON other.id = a
)
SELECT * FROM foo
```

is re-written to

```sql
MODEL (
  name test.test
);

WITH foo AS (
  SELECT
    a,
    b
  FROM base
  /*
        NOTE: LEFT JOIN because blah blah blah
      */
  LEFT JOIN other
    ON other.id = a
)
SELECT
  *
FROM foo
```

running `sqlmesh format` again will result in

```sql
MODEL (
  name test.test
);

WITH foo AS (
  SELECT
    a,
    b
  FROM base
  /*
          NOTE: LEFT JOIN because blah blah blah
        */
  LEFT JOIN other
    ON other.id = a
)
SELECT
  *
FROM foo
```


Pull request description:
fix!: improve sqlmesh format closes #2469
Closes #2469

Patch:
diff --git a/setup.py b/setup.py
index 249e5baa..3fd5a560 100644
--- a/setup.py
+++ b/setup.py
@@ -46,7 +46,7 @@ setup(
         "requests",
         "rich[jupyter]",
         "ruamel.yaml",
-        "sqlglot[rs]~=23.10.0",
+        "sqlglot[rs]~=23.11.2",
     ],
     extras_require={
         "bigquery": [
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index fe8b0d42..4006e8f1 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -195,6 +195,7 @@ def _parse_macro(self: Parser, keyword_macro: str = "") -> t.Optional[exp.Expres
     if self._prev.text != SQLMESH_MACRO_PREFIX:
         return self._parse_parameter()
 
+    comments = self._prev.comments
     index = self._index
     field = self._parse_primary() or self._parse_function(functions={}) or self._parse_id_var()
 
@@ -207,25 +208,37 @@ def _parse_macro(self: Parser, keyword_macro: str = "") -> t.Optional[exp.Expres
         if isinstance(field, exp.Anonymous):
             if macro_name == "DEF":
                 return self.expression(
-                    MacroDef, this=field.expressions[0], expression=field.expressions[1]
+                    MacroDef,
+                    this=field.expressions[0],
+                    expression=field.expressions[1],
+                    comments=comments,
                 )
             if macro_name == "SQL":
                 into = field.expressions[1].this.lower() if len(field.expressions) > 1 else None
-                return self.expression(MacroSQL, this=field.expressions[0], into=into)
+                return self.expression(
+                    MacroSQL, this=field.expressions[0], into=into, comments=comments
+                )
         else:
-            field = exp.Anonymous(this=field.sql_name(), expressions=list(field.args.values()))
+            field = self.expression(
+                exp.Anonymous,
+                this=field.sql_name(),
+                expressions=list(field.args.values()),
+                comments=comments,
+            )
 
-        return self.expression(MacroFunc, this=field)
+        return self.expression(MacroFunc, this=field, comments=comments)
 
     if field is None:
         return None
 
     if field.is_string or (isinstance(field, exp.Identifier) and field.quoted):
-        return self.expression(MacroStrReplace, this=exp.Literal.string(field.this))
+        return self.expression(
+            MacroStrReplace, this=exp.Literal.string(field.this), comments=comments
+        )
 
     if "@" in field.this:
         return field
-    return self.expression(MacroVar, this=field.this)
+    return self.expression(MacroVar, this=field.this, comments=comments)
 
 
 KEYWORD_MACROS = {"WITH", "JOIN", "WHERE", "GROUP_BY", "HAVING", "ORDER_BY", "LIMIT"}
@@ -778,7 +791,11 @@ def extend_sqlglot() -> None:
                 }
             )
 
-            generator.WITH_SEPARATED_COMMENTS = (*generator.WITH_SEPARATED_COMMENTS, Model)
+            generator.WITH_SEPARATED_COMMENTS = (
+                *generator.WITH_SEPARATED_COMMENTS,
+                Model,
+                MacroDef,
+            )
 
     _override(Parser, _parse_statement)
     _override(Parser, _parse_join)


Test patch:
diff --git a/tests/core/test_dialect.py b/tests/core/test_dialect.py
index b4770a30..4bbfb22d 100644
--- a/tests/core/test_dialect.py
+++ b/tests/core/test_dialect.py
@@ -25,10 +25,36 @@ def test_format_model_expressions():
     references (
      a,
      (b, c) as d,
-     )
+     ),
+     audits [
+    not_null(columns=[
+      foo_id,
+      foo_normalised,
+      bar_normalised,
+      total_weight,
+      overall_rank,
+      cumulative_total_weight_share,
+      market_rank,
+      market_cumulative_total_weight_share,
+      total_weight_decile,
+      tier,
+    ]),
+
+    unique_values(columns=[foo_id]),
+
+    accepted_range(column=foo_normalised, min_v=0, max_v=100),
+    accepted_range(column=bar_normalised, min_v=0, max_v=100),
+    accepted_range(column=total_weight, min_v=0, max_v=100),
+    accepted_range(column=cumulative_total_weight_share, min_v=0, max_v=1),
+    accepted_range(column=market_cumulative_total_weight_share, min_v=0, max_v=1),
+
+    accepted_values(column=tier, is_in=['Tier 1', 'Tier 2', 'Tier 3', 'Long Tail']),
+    accepted_values(column=total_weight_decile, is_in=['Decile_01', 'Decile_02','Decile_03','Decile_04','Decile_05','Decile_06','Decile_07','Decile_08','Decile_09','Decile_10']),
+  ],
     )
     ;
 
+    /* comment */
     @DEF(x
     , 1);
 
@@ -59,9 +85,48 @@ def test_format_model_expressions():
         == """MODEL (
   name a.b,
   kind FULL,
-  references (a, (b, c) AS d)
+  references (a, (b, c) AS d),
+  audits ARRAY(
+    NOT_NULL(
+      columns = ARRAY(
+        foo_id,
+        foo_normalised,
+        bar_normalised,
+        total_weight,
+        overall_rank,
+        cumulative_total_weight_share,
+        market_rank,
+        market_cumulative_total_weight_share,
+        total_weight_decile,
+        tier
+      )
+    ),
+    UNIQUE_VALUES(columns = ARRAY(foo_id)),
+    ACCEPTED_RANGE(column = foo_normalised, min_v = 0, max_v = 100),
+    ACCEPTED_RANGE(column = bar_normalised, min_v = 0, max_v = 100),
+    ACCEPTED_RANGE(column = total_weight, min_v = 0, max_v = 100),
+    ACCEPTED_RANGE(column = cumulative_total_weight_share, min_v = 0, max_v = 1),
+    ACCEPTED_RANGE(column = market_cumulative_total_weight_share, min_v = 0, max_v = 1),
+    ACCEPTED_VALUES(column = tier, is_in = ARRAY('Tier 1', 'Tier 2', 'Tier 3', 'Long Tail')),
+    ACCEPTED_VALUES(
+      column = total_weight_decile,
+      is_in = ARRAY(
+        'Decile_01',
+        'Decile_02',
+        'Decile_03',
+        'Decile_04',
+        'Decile_05',
+        'Decile_06',
+        'Decile_07',
+        'Decile_08',
+        'Decile_09',
+        'Decile_10'
+      )
+    )
+  )
 );
 
+/* comment */
 @DEF(x, 1);
 
 SELECT