diff --git a/sqlglot/generator.py b/sqlglot/generator.py index 857eff1459..7fbef6e312 100644 --- a/sqlglot/generator.py +++ b/sqlglot/generator.py @@ -260,8 +260,9 @@ class Generator: # Expressions whose comments are separated from them for better formatting WITH_SEPARATED_COMMENTS: t.Tuple[t.Type[exp.Expression], ...] = ( - exp.Select, + exp.Drop, exp.From, + exp.Select, exp.Where, exp.With, ) diff --git a/sqlglot/parser.py b/sqlglot/parser.py index db2ff69563..324b520321 100644 --- a/sqlglot/parser.py +++ b/sqlglot/parser.py @@ -1145,6 +1145,7 @@ def _parse_drop(self) -> exp.Drop | exp.Command: return self.expression( exp.Drop, + comments=start.comments, exists=self._parse_exists(), this=self._parse_table(schema=True), kind=kind, diff --git a/tests/test_transpile.py b/tests/test_transpile.py index 1138b4e4a9..7a5de31cba 100644 --- a/tests/test_transpile.py +++ b/tests/test_transpile.py @@ -285,6 +285,47 @@ def test_comments(self): "SELECT 1 /* hi this is a comment */", read="snowflake", ) + self.validate( + "-- comment\nDROP TABLE IF EXISTS foo", + "/* comment */ DROP TABLE IF EXISTS foo", + ) + self.validate( + """ + -- comment1 + -- comment2 + + -- comment3 + DROP TABLE IF EXISTS db.tba + """, + """/* comment1 */ +/* comment2 */ +/* comment3 */ +DROP TABLE IF EXISTS db.tba""", + pretty=True, + ) + self.validate( + """ + CREATE TABLE db.tba AS + SELECT a, b, c + FROM tb_01 + WHERE + -- comment5 + a = 1 AND b = 2 --comment6 + -- and c = 1 + -- comment7 + """, + """CREATE TABLE db.tba AS +SELECT + a, + b, + c +FROM tb_01 +WHERE + a /* comment5 */ = 1 AND b = 2 /* comment6 */ + /* and c = 1 */ + /* comment7 */""", + pretty=True, + ) def test_types(self): self.validate("INT 1", "CAST(1 AS INT)")