From 1f60c83bce17c531b02d0d72da13493fc171a2ce Mon Sep 17 00:00:00 2001
From: Yuya Ebihara <ebyhry@gmail.com>
Date: Sat, 14 Sep 2024 21:26:49 +0900
Subject: [PATCH] Fix round trip of DROP ROLE statement

Also, migrate assertStatement to ParserAssert.
---
 .../main/java/io/trino/sql/SqlFormatter.java   |  6 +++++-
 .../main/java/io/trino/sql/tree/DropRole.java  |  6 ++++--
 .../io/trino/sql/parser/TestSqlParser.java     | 18 ++++++++++++------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/core/trino-parser/src/main/java/io/trino/sql/SqlFormatter.java b/core/trino-parser/src/main/java/io/trino/sql/SqlFormatter.java
index c4bef618c9fc..a88e09346789 100644
--- a/core/trino-parser/src/main/java/io/trino/sql/SqlFormatter.java
+++ b/core/trino-parser/src/main/java/io/trino/sql/SqlFormatter.java
@@ -2066,7 +2066,11 @@ protected Void visitCreateRole(CreateRole node, Integer indent)
         @Override
         protected Void visitDropRole(DropRole node, Integer indent)
         {
-            builder.append("DROP ROLE ").append(formatName(node.getName()));
+            builder.append("DROP ROLE ");
+            if (node.isExists()) {
+                builder.append("IF EXISTS ");
+            }
+            builder.append(formatName(node.getName()));
             node.getCatalog().ifPresent(catalog -> builder
                     .append(" IN ")
                     .append(formatName(catalog)));
diff --git a/core/trino-parser/src/main/java/io/trino/sql/tree/DropRole.java b/core/trino-parser/src/main/java/io/trino/sql/tree/DropRole.java
index 325aa7f4d746..4f32521133b9 100644
--- a/core/trino-parser/src/main/java/io/trino/sql/tree/DropRole.java
+++ b/core/trino-parser/src/main/java/io/trino/sql/tree/DropRole.java
@@ -63,13 +63,14 @@ public boolean equals(Object o)
         }
         DropRole dropRole = (DropRole) o;
         return Objects.equals(name, dropRole.name) &&
-                Objects.equals(catalog, dropRole.catalog);
+                Objects.equals(catalog, dropRole.catalog) &&
+                exists == dropRole.exists;
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(name, catalog);
+        return Objects.hash(name, catalog, exists);
     }
 
     @Override
@@ -78,6 +79,7 @@ public String toString()
         return toStringHelper(this)
                 .add("name", name)
                 .add("catalog", catalog)
+                .add("exists", exists)
                 .toString();
     }
 
diff --git a/core/trino-parser/src/test/java/io/trino/sql/parser/TestSqlParser.java b/core/trino-parser/src/test/java/io/trino/sql/parser/TestSqlParser.java
index 0079883aec6f..b8b78253d417 100644
--- a/core/trino-parser/src/test/java/io/trino/sql/parser/TestSqlParser.java
+++ b/core/trino-parser/src/test/java/io/trino/sql/parser/TestSqlParser.java
@@ -4754,12 +4754,18 @@ public void testCreateRole()
     @Test
     public void testDropRole()
     {
-        assertStatement("DROP ROLE role", new DropRole(location(1, 1), new Identifier("role"), Optional.empty(), false));
-        assertStatement("DROP ROLE IF EXISTS role", new DropRole(location(1, 1), new Identifier("role"), Optional.empty(), false));
-        assertStatement("DROP ROLE \"role\"", new DropRole(location(1, 1), new Identifier("role"), Optional.empty(), false));
-        assertStatement("DROP ROLE \"ro le\"", new DropRole(location(1, 1), new Identifier("ro le"), Optional.empty(), false));
-        assertStatement("DROP ROLE \"!@#$%^&*'ад\"\"мін\"", new DropRole(location(1, 1), new Identifier("!@#$%^&*'ад\"мін"), Optional.empty(), false));
-        assertStatement("DROP ROLE role IN my_catalog", new DropRole(location(1, 1), new Identifier("role"), Optional.of(new Identifier("my_catalog")), false));
+        assertThat(statement("DROP ROLE role")).isEqualTo(
+                new DropRole(location(1, 1), new Identifier(location(1, 11), "role", false), Optional.empty(), false));
+        assertThat(statement("DROP ROLE IF EXISTS role")).isEqualTo(
+                new DropRole(location(1, 1), new Identifier(location(1, 21), "role", false), Optional.empty(), true));
+        assertThat(statement("DROP ROLE \"role\"")).isEqualTo(
+                new DropRole(location(1, 1), new Identifier(location(1, 11), "role", true), Optional.empty(), false));
+        assertThat(statement("DROP ROLE \"ro le\"")).isEqualTo(
+                new DropRole(location(1, 1), new Identifier(location(1, 11), "ro le", true), Optional.empty(), false));
+        assertThat(statement("DROP ROLE \"!@#$%^&*'ад\"\"мін\"")).isEqualTo(
+                new DropRole(location(1, 1), new Identifier(location(1, 11), "!@#$%^&*'ад\"мін", true), Optional.empty(), false));
+        assertThat(statement("DROP ROLE role IN my_catalog")).isEqualTo(
+                new DropRole(location(1, 1), new Identifier(location(1, 11), "role", false), Optional.of(new Identifier(location(1, 19), "my_catalog", false)), false));
     }
 
     @Test