Introduce SET ROLE statement

Extracted-From: prestodb/presto#10904
trinodb · Jan 29, 2019 · dd9b64c · dd9b64c
1 parent e8f81d9
commit dd9b64c
Show file tree

Hide file tree

Showing 6 changed files with 171 additions and 1 deletion.
diff --git a/presto-parser/src/main/antlr4/io/prestosql/sql/parser/SqlBase.g4 b/presto-parser/src/main/antlr4/io/prestosql/sql/parser/SqlBase.g4
@@ -76,6 +76,7 @@ statement
         FROM principal (',' principal)*
         (GRANTED BY grantor)?
         (IN catalog=identifier)?                                       #revokeRoles
+    | SET ROLE (ALL | NONE | role=identifier) (IN catalog=identifier)? #setRole
     | GRANT
         (privilege (',' privilege)* | ALL PRIVILEGES)
         ON TABLE? qualifiedName TO grantee=identifier
@@ -497,7 +498,7 @@ nonReserved
     | JSON
     | LAST | LATERAL | LEVEL | LIMIT | LOGICAL
     | MAP | MINUTE | MONTH
-    | NFC | NFD | NFKC | NFKD | NO | NULLIF | NULLS
+    | NFC | NFD | NFKC | NFKD | NO | NONE | NULLIF | NULLS
     | ONLY | OPTION | ORDINALITY | OUTPUT | OVER
     | PARTITION | PARTITIONS | PATH | POSITION | PRECEDING | PRIVILEGES | PROPERTIES
     | RANGE | READ | RENAME | REPEATABLE | REPLACE | RESET | RESTRICT | REVOKE | ROLE | ROLES | ROLLBACK | ROW | ROWS
@@ -616,6 +617,7 @@ NFD : 'NFD';
 NFKC : 'NFKC';
 NFKD : 'NFKD';
 NO: 'NO';
+NONE: 'NONE';
 NORMALIZE: 'NORMALIZE';
 NOT: 'NOT';
 NULL: 'NULL';

diff --git a/presto-parser/src/main/java/io/prestosql/sql/SqlFormatter.java b/presto-parser/src/main/java/io/prestosql/sql/SqlFormatter.java
@@ -79,6 +79,7 @@
 import io.prestosql.sql.tree.Select;
 import io.prestosql.sql.tree.SelectItem;
 import io.prestosql.sql.tree.SetPath;
+import io.prestosql.sql.tree.SetRole;
 import io.prestosql.sql.tree.SetSession;
 import io.prestosql.sql.tree.ShowCatalogs;
 import io.prestosql.sql.tree.ShowColumns;
@@ -1164,6 +1165,28 @@ protected Void visitRevokeRoles(RevokeRoles node, Integer context)
             return null;
         }
 
+        @Override
+        protected Void visitSetRole(SetRole node, Integer context)
+        {
+            builder.append("SET ROLE ");
+            SetRole.Type type = node.getType();
+            switch (type) {
+                case ALL:
+                case NONE:
+                    builder.append(type.toString());
+                    break;
+                case ROLE:
+                    builder.append(node.getRole().get());
+                    break;
+                default:
+                    throw new IllegalArgumentException("Unsupported type: " + type);
+            }
+            if (node.getCatalog().isPresent()) {
+                builder.append(" IN ").append(node.getCatalog().get());
+            }
+            return null;
+        }
+
         @Override
         public Void visitGrant(Grant node, Integer indent)
         {

diff --git a/presto-parser/src/main/java/io/prestosql/sql/parser/AstBuilder.java b/presto-parser/src/main/java/io/prestosql/sql/parser/AstBuilder.java
@@ -11,6 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 package io.prestosql.sql.parser;
 
 import com.google.common.collect.ImmutableList;
@@ -129,6 +130,7 @@
 import io.prestosql.sql.tree.Select;
 import io.prestosql.sql.tree.SelectItem;
 import io.prestosql.sql.tree.SetPath;
+import io.prestosql.sql.tree.SetRole;
 import io.prestosql.sql.tree.SetSession;
 import io.prestosql.sql.tree.ShowCatalogs;
 import io.prestosql.sql.tree.ShowColumns;
@@ -864,6 +866,19 @@ public Node visitRevokeRoles(SqlBaseParser.RevokeRolesContext context)
                 getIdentifierIfPresent(context.catalog));
     }
 
+    @Override
+    public Node visitSetRole(SqlBaseParser.SetRoleContext context)
+    {
+        SetRole.Type type = SetRole.Type.ROLE;
+        if (context.ALL() != null) {
+            type = SetRole.Type.ALL;
+        }
+        else if (context.NONE() != null) {
+            type = SetRole.Type.NONE;
+        }
+        return new SetRole(getLocation(context), type, getIdentifierIfPresent(context.role), getIdentifierIfPresent(context.catalog));
+    }
+
     @Override
     public Node visitGrant(SqlBaseParser.GrantContext context)
     {

diff --git a/presto-parser/src/main/java/io/prestosql/sql/tree/AstVisitor.java b/presto-parser/src/main/java/io/prestosql/sql/tree/AstVisitor.java
@@ -617,6 +617,11 @@ protected R visitRevokeRoles(RevokeRoles node, C context)
         return visitStatement(node, context);
     }
 
+    protected R visitSetRole(SetRole node, C context)
+    {
+        return visitStatement(node, context);
+    }
+
     protected R visitGrant(Grant node, C context)
     {
         return visitStatement(node, context);

diff --git a/presto-parser/src/main/java/io/prestosql/sql/tree/SetRole.java b/presto-parser/src/main/java/io/prestosql/sql/tree/SetRole.java
@@ -0,0 +1,112 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.sql.tree;
+
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+
+import static com.google.common.base.MoreObjects.toStringHelper;
+import static java.util.Objects.requireNonNull;
+
+public class SetRole
+        extends Statement
+{
+    public enum Type
+    {
+        ROLE, ALL, NONE
+    }
+
+    private final Type type;
+    private final Optional<Identifier> role;
+    private final Optional<Identifier> catalog;
+
+    public SetRole(Type type, Optional<Identifier> role, Optional<Identifier> catalog)
+    {
+        this(Optional.empty(), type, role, catalog);
+    }
+
+    public SetRole(NodeLocation location, Type type, Optional<Identifier> role, Optional<Identifier> catalog)
+    {
+        this(Optional.of(location), type, role, catalog);
+    }
+
+    private SetRole(Optional<NodeLocation> location, Type type, Optional<Identifier> role, Optional<Identifier> catalog)
+    {
+        super(location);
+        this.type = requireNonNull(type, "type is null");
+        this.role = requireNonNull(role, "role is null");
+        this.catalog = requireNonNull(catalog, "catalog is null");
+    }
+
+    public Type getType()
+    {
+        return type;
+    }
+
+    public Optional<Identifier> getRole()
+    {
+        return role;
+    }
+
+    public Optional<Identifier> getCatalog()
+    {
+        return catalog;
+    }
+
+    @Override
+    public List<? extends Node> getChildren()
+    {
+        return ImmutableList.of();
+    }
+
+    @Override
+    public <R, C> R accept(AstVisitor<R, C> visitor, C context)
+    {
+        return visitor.visitSetRole(this, context);
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        SetRole setRole = (SetRole) o;
+        return type == setRole.type &&
+                Objects.equals(role, setRole.role) &&
+                Objects.equals(catalog, setRole.catalog);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(type, role, catalog);
+    }
+
+    @Override
+    public String toString()
+    {
+        return toStringHelper(this)
+                .add("type", type)
+                .add("role", role)
+                .add("catalog", catalog)
+                .toString();
+    }
+}
diff --git a/presto-parser/src/test/java/io/prestosql/sql/parser/TestSqlParser.java b/presto-parser/src/test/java/io/prestosql/sql/parser/TestSqlParser.java
@@ -110,6 +110,7 @@
 import io.prestosql.sql.tree.Select;
 import io.prestosql.sql.tree.SelectItem;
 import io.prestosql.sql.tree.SetPath;
+import io.prestosql.sql.tree.SetRole;
 import io.prestosql.sql.tree.SetSession;
 import io.prestosql.sql.tree.ShowCatalogs;
 import io.prestosql.sql.tree.ShowColumns;
@@ -2213,6 +2214,18 @@ public void testRevokeRoles()
                         Optional.of(new Identifier("catalog"))));
     }
 
+    @Test
+    public void testSetRole()
+            throws Exception
+    {
+        assertStatement("SET ROLE ALL", new SetRole(SetRole.Type.ALL, Optional.empty(), Optional.empty()));
+        assertStatement("SET ROLE NONE", new SetRole(SetRole.Type.NONE, Optional.empty(), Optional.empty()));
+        assertStatement("SET ROLE role", new SetRole(SetRole.Type.ROLE, Optional.of(new Identifier("role")), Optional.empty()));
+        assertStatement("SET ROLE ALL IN catalog", new SetRole(SetRole.Type.ALL, Optional.empty(), Optional.of(new Identifier("catalog"))));
+        assertStatement("SET ROLE NONE IN catalog", new SetRole(SetRole.Type.NONE, Optional.empty(), Optional.of(new Identifier("catalog"))));
+        assertStatement("SET ROLE \"role\" IN \"catalog\"", new SetRole(SetRole.Type.ROLE, Optional.of(new Identifier("role")), Optional.of(new Identifier("catalog"))));
+    }
+
     private static void assertCast(String type)
     {
         assertCast(type, type);