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'`

Discussion:

Patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 3528d110..77a06740 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 28a8d4d7..38b7c6dd 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -1196,7 +1196,11 @@ class Generator(metaclass=_Generator):
         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 @@ class Generator(metaclass=_Generator):
         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 @@ class Generator(metaclass=_Generator):
             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 87a1579d..b69cfcb7 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -3030,15 +3030,20 @@ class Parser(metaclass=_Parser):
         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 07a915d3..e7f1665f 100644
--- a/tests/test_transpile.py
+++ b/tests/test_transpile.py
@@ -563,7 +563,36 @@ FROM x""",
         )
         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):