From dd9b64cafc8ba47028aee72ddcd3594a691da3e1 Mon Sep 17 00:00:00 2001 From: Andrii Rosa Date: Tue, 28 Nov 2017 14:50:51 +0100 Subject: [PATCH] Introduce SET ROLE statement Extracted-From: https://github.com/prestodb/presto/pull/10904 --- .../antlr4/io/prestosql/sql/parser/SqlBase.g4 | 4 +- .../java/io/prestosql/sql/SqlFormatter.java | 23 ++++ .../io/prestosql/sql/parser/AstBuilder.java | 15 +++ .../io/prestosql/sql/tree/AstVisitor.java | 5 + .../java/io/prestosql/sql/tree/SetRole.java | 112 ++++++++++++++++++ .../prestosql/sql/parser/TestSqlParser.java | 13 ++ 6 files changed, 171 insertions(+), 1 deletion(-) create mode 100644 presto-parser/src/main/java/io/prestosql/sql/tree/SetRole.java 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 index e2e5ccbc60934..776d0152a37fc 100644 --- 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 index 7764caad88cc7..c31bb350cdb98 100644 --- 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 index 3c83996f82b15..f576d5c087960 100644 --- 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 index 816266b288ec3..c86d3babccb83 100644 --- 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 new file mode 100644 index 0000000000000..bca68fb14485b --- /dev/null +++ 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 role; + private final Optional catalog; + + public SetRole(Type type, Optional role, Optional catalog) + { + this(Optional.empty(), type, role, catalog); + } + + public SetRole(NodeLocation location, Type type, Optional role, Optional catalog) + { + this(Optional.of(location), type, role, catalog); + } + + private SetRole(Optional location, Type type, Optional role, Optional 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 getRole() + { + return role; + } + + public Optional getCatalog() + { + return catalog; + } + + @Override + public List getChildren() + { + return ImmutableList.of(); + } + + @Override + public R accept(AstVisitor 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 index 7852b3fe18b91..e60826c8468a9 100644 --- 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);