Repository: tobymao/sqlglot

Issue:
Dropping and reordering comments when using CTEs
I think I've found some corner cases where:

1. Comments are dropped
2. Comments before the final select statement after a sequence of CTEs gets moved to the top of the query
3. Some of the comments are reproduced out of order

## Example

Here's an example:

```python
import sqlglot
def foo(s):
    print(sqlglot.transpile(s, read='presto', write='spark', pretty=True)[0])

s1 = """
-- comment at the top
with
-- comment for tbl1
tbl1 as (select 1)

-- comment for tbl2
, tb12 as (select 2)

-- comment for tbl3
, tb13 as (select 3)

-- comment for final select
select * from tbl1
"""
foo(s1)
```

gives me the output:

```sql
/* comment for final select */
/* comment at the top */
WITH tbl1 /* comment for tbl1 */ AS (
  SELECT
    1
), tb12 AS (
  SELECT
    2
), tb13 AS (
  SELECT
    3
)
SELECT
  *
FROM tbl1
```

## Expected behavior

1. Preserve comments for `tbl2` and `tbl3` (they are missing in the example above)
2. If possible, keep the comment for the final select closer to where it was in the original query
4. Even if it can't be kept closer, the comments are out of order, as the "comment at the top" should be first.

## Alternative inputs

I played with the example a bit, and it seems that the dropped comments have to do with the comma placement. I can preserve the comments if I do something like the example below, but, ideally, they would also be preserved in the example above. With the example below, we still have the issue of the comment at the top and the final select comment showing up out of order, and the final select comment being moved to the top of the output (rather than near the final select statement).

```python
s2 = """
-- comment at the top
with
-- comment for tbl1
tbl1 as (select 1),

-- comment for tbl2
tb12 as (select 2),

-- comment for tbl3
tb13 as (select 3)

-- comment for final select
select * from tbl1
"""
foo(s2)
```

gives the output:

```sql
/* comment for final select */
/* comment at the top */
WITH tbl1 /* comment for tbl1 */ AS (
  SELECT
    1
), tb12 /* comment for tbl2 */ AS (
  SELECT
    2
), tb13 /* comment for tbl3 */ AS (
  SELECT
    3
)
SELECT
  *
FROM tbl1
```

## Version

This behavior is happening on the main branch. I ran the example using `sqlglot.__version__='25.24.5.dev2'`

Hints from issue comments:
@VaggelisD: Hey @ajfriend, thanks for reporting this. 

Note that preserving the comments exactly as they come is not always possible, our attempts are on a best-effort basis as is documented [here](https://github.com/tobymao/sqlglot?tab=readme-ov-file#examples). 

The issue you describe regarding the comma placement is precisely what happens; The comments are attached to nearby AST nodes but commas are not added to the AST, so it might be that their surrounding comments are either dropped or attached to other nodes (thus, appearing out of order in the result query).


Solution patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 3528d11060..77a067404b 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -7054,7 +7054,7 @@ def update(
         )
     if with_:
         cte_list = [
-            CTE(this=maybe_parse(qry, dialect=dialect, **opts), alias=alias)
+            alias_(CTE(this=maybe_parse(qry, dialect=dialect, **opts)), alias, table=True)
             for alias, qry in with_.items()
         ]
         update_expr.set(
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 28a8d4d73f..38b7c6ddcb 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -1196,7 +1196,11 @@ def with_sql(self, expression: exp.With) -> str:
         return f"WITH {recursive}{sql}"
 
     def cte_sql(self, expression: exp.CTE) -> str:
-        alias = self.sql(expression, "alias")
+        alias = expression.args.get("alias")
+        if alias:
+            alias.add_comments(expression.pop_comments())
+
+        alias_sql = self.sql(expression, "alias")
 
         materialized = expression.args.get("materialized")
         if materialized is False:
@@ -1204,7 +1208,7 @@ def cte_sql(self, expression: exp.CTE) -> str:
         elif materialized:
             materialized = "MATERIALIZED "
 
-        return f"{alias} AS {materialized or ''}{self.wrap(expression)}"
+        return f"{alias_sql} AS {materialized or ''}{self.wrap(expression)}"
 
     def tablealias_sql(self, expression: exp.TableAlias) -> str:
         alias = self.sql(expression, "this")
@@ -2517,6 +2521,11 @@ def select_sql(self, expression: exp.Select) -> str:
             self.sql(expression, "from", comment=False),
         )
 
+        # If both the CTE and SELECT clauses have comments, generate the latter earlier
+        if expression.args.get("with"):
+            sql = self.maybe_comment(sql, expression)
+            expression.pop_comments()
+
         sql = self.prepend_ctes(expression, sql)
 
         if not self.SUPPORTS_SELECT_INTO and into:
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 87a1579d88..b69cfcb721 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -3030,15 +3030,20 @@ def _parse_with(self, skip_with_token: bool = False) -> t.Optional[exp.With]:
         comments = self._prev_comments
         recursive = self._match(TokenType.RECURSIVE)
 
+        last_comments = None
         expressions = []
         while True:
             expressions.append(self._parse_cte())
+            if last_comments:
+                expressions[-1].add_comments(last_comments)
 
             if not self._match(TokenType.COMMA) and not self._match(TokenType.WITH):
                 break
             else:
                 self._match(TokenType.WITH)
 
+            last_comments = self._prev_comments
+
         return self.expression(
             exp.With, comments=comments, expressions=expressions, recursive=recursive
         )


Test patch:
diff --git a/tests/test_transpile.py b/tests/test_transpile.py
index 07a915d3a7..e7f1665fc5 100644
--- a/tests/test_transpile.py
+++ b/tests/test_transpile.py
@@ -563,7 +563,36 @@ def test_comments(self):
         )
         self.validate(
             """with a as /* comment */ ( select * from b) select * from a""",
-            """WITH a AS (SELECT * FROM b) /* comment */ SELECT * FROM a""",
+            """WITH a /* comment */ AS (SELECT * FROM b) SELECT * FROM a""",
+        )
+        self.validate(
+            """
+  -- comment at the top
+WITH
+-- comment for tbl1
+tbl1 AS (SELECT 1)
+-- comment for tbl2
+, tbl2 AS (SELECT 2)
+-- comment for tbl3
+, tbl3 AS (SELECT 3)
+-- comment for final select
+SELECT * FROM tbl1""",
+            """/* comment at the top */
+WITH tbl1 /* comment for tbl1 */ AS (
+  SELECT
+    1
+), tbl2 /* comment for tbl2 */ AS (
+  SELECT
+    2
+), tbl3 /* comment for tbl3 */ AS (
+  SELECT
+    3
+)
+/* comment for final select */
+SELECT
+  *
+FROM tbl1""",
+            pretty=True,
         )
 
     def test_types(self):