Repository: tobymao/sqlglot

Issue:
Comments lose their position in queries 
I'm seeing an issue where comments end up moving to different positions in the tree during transpilation.  For example:

```
>>> query = """
... select a from b join c on c.id = b.id
... where 1 = 1
... -- comment 1
... and foo
... -- comment 2
... and bar
... """
>>> sqlglot.transpile(query)
['SELECT a FROM b JOIN c ON c.id = b.id WHERE /* comment 2 */ /* comment 1 */ 1 = 1 AND foo AND bar']
```

The comment2 move to the beginning of the WHERE clause.  If there is no condition after the comment, the comment is dropped entirely.

If we transpile that query again, the comments move again:

```
'SELECT a FROM b JOIN c ON c.id = b.id/* comment 2 */ /* comment 1 */  WHERE 1 = 1 AND foo AND bar'
```


Discussion:
This is related to https://github.com/tobymao/sqlglot/commit/1f9a4f7f4 (related discussion: https://github.com/tobymao/sqlglot/pull/758#discussion_r1031702139).

cc: @mpf82 I think we may need to get rid of `exp.Binary` from `WITH_SEPARATED_COMMENTS`. I'll play around a bit and see if I can fix both this and the next issue I mentioned.
There's actually another issue, because connectors (`AND`, `OR`) have special sqlgen logic. I'm looking into this now.

Example:
```python
import sqlglot

query = """
select a from b
where foo
-- comment 1
and bar
-- comment 2
and bla
"""

print(sqlglot.parse_one(query).sql(pretty=True))
```

Which outputs:
```
➜  sqlglot git:(main) ✗ python test_comments.py
SELECT
  a
FROM b
WHERE
  foo AND bar AND bla /* comment 2 */
```

Thanks for the report @joshmarcus!

Resolving patch:
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -183,7 +183,7 @@ class Generator:
         exp.WithJournalTableProperty: exp.Properties.Location.POST_NAME,
     }
 
-    WITH_SEPARATED_COMMENTS = (exp.Select, exp.From, exp.Where, exp.Binary)
+    WITH_SEPARATED_COMMENTS = (exp.Select, exp.From, exp.Where)
     SENTINEL_LINE_BREAK = "__SQLGLOT__LB__"
 
     __slots__ = (
@@ -322,10 +322,15 @@ def pad_comment(self, comment: str) -> str:
         comment = comment + " " if comment[-1].strip() else comment
         return comment
 
-    def maybe_comment(self, sql: str, expression: exp.Expression) -> str:
-        comments = expression.comments if self._comments else None
+    def maybe_comment(
+        self,
+        sql: str,
+        expression: t.Optional[exp.Expression] = None,
+        comments: t.Optional[t.List[str]] = None,
+    ) -> str:
+        comments = (comments or (expression and expression.comments)) if self._comments else None  # type: ignore
 
-        if not comments:
+        if not comments or isinstance(expression, exp.Binary):
             return sql
 
         sep = "\n" if self.pretty else " "
@@ -1722,7 +1727,18 @@ def connector_sql(self, expression: exp.Connector, op: str) -> str:
         if not self.pretty:
             return self.binary(expression, op)
 
-        sqls = tuple(self.sql(e) for e in expression.flatten(unnest=False))
+        sqls = []
+        comments = set()
+
+        for e in tuple(expression.flatten(unnest=False)):
+            sql = self.sql(e)
+
+            if id(e.parent) not in comments:
+                comments.add(id(e.parent))
+                sql = self.maybe_comment(sql, e, e.parent.comments)
+
+            sqls.append(sql)
+
         sep = "\n" if self.text_width(sqls) > self._max_text_width else " "
         return f"{sep}{op} ".join(sqls)
 
@@ -1945,6 +1961,7 @@ def use_sql(self, expression: exp.Use) -> str:
         return f"USE{kind}{this}"
 
     def binary(self, expression: exp.Binary, op: str) -> str:
+        op = self.maybe_comment(op, comments=expression.comments)
         return f"{self.sql(expression, 'this')} {op} {self.sql(expression, 'expression')}"
 
     def function_fallback_sql(self, expression: exp.Func) -> str:
diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
--- a/sqlglot/tokens.py
+++ b/sqlglot/tokens.py
@@ -810,6 +810,9 @@ def _scan(self, until: t.Optional[t.Callable] = None) -> None:
             if until and until():
                 break
 
+        if self.tokens:
+            self.tokens[-1].comments.extend(self._comments)
+
     def _chars(self, size: int) -> str:
         if size == 1:
             return self._char  # type: ignore
@@ -949,7 +952,7 @@ def _scan_comment(self, comment_start: str) -> bool:
 
         # Leading comment is attached to the succeeding token, whilst trailing comment to the preceding.
         # Multiple consecutive comments are preserved by appending them to the current comments list.
-        if comment_start_line == self._prev_token_line or self._end:
+        if comment_start_line == self._prev_token_line:
             self.tokens[-1].comments.extend(self._comments)
             self._comments = []
             self._prev_token_line = self._line


Test patch:
diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql
--- a/tests/fixtures/identity.sql
+++ b/tests/fixtures/identity.sql
@@ -719,6 +719,7 @@ SELECT a /* x */ /* y */ /* z */, b /* k */ /* m */
 SELECT * FROM foo /* x */, bla /* x */
 SELECT 1 /* comment */ + 1
 SELECT 1 /* c1 */ + 2 /* c2 */
+SELECT 1 /* c1 */ + /* c2 */ 2 /* c3 */
 SELECT 1 /* c1 */ + 2 /* c2 */ + 3 /* c3 */
 SELECT 1 /* c1 */ + 2 /* c2 */, 3 /* c3 */
 SELECT x FROM a.b.c /* x */, e.f.g /* x */
diff --git a/tests/test_transpile.py b/tests/test_transpile.py
--- a/tests/test_transpile.py
+++ b/tests/test_transpile.py
@@ -200,6 +200,40 @@ def test_comments(self):
             read="mysql",
             pretty=True,
         )
+        self.validate(
+            """
+            SELECT a FROM b
+            WHERE foo
+            -- comment 1
+            AND bar
+            -- comment 2
+            AND bla;
+            """,
+            "SELECT a FROM b WHERE foo AND /* comment 1 */ bar AND /* comment 2 */ bla",
+        )
+        self.validate(
+            """
+            SELECT a FROM b WHERE foo
+            -- comment 1
+            """,
+            "SELECT a FROM b WHERE foo /* comment 1 */",
+        )
+        self.validate(
+            """
+            select a from b
+            where foo
+            -- comment 1
+            and bar
+            -- comment 2
+            and bla
+            """,
+            """SELECT
+  a
+FROM b
+WHERE
+  foo /* comment 1 */ AND bar AND bla /* comment 2 */""",
+            pretty=True,
+        )
 
     def test_types(self):
         self.validate("INT 1", "CAST(1 AS INT)")