Add SHOW CURRENT ROLES

Instead of select * from information_schema.roles, SHOW CURRENT ROLES rewrites to select * from information_schema.enabled_roles. All users can see what roles they're currently using, so no need for access control checks. Extracted-From: prestodb/presto#10904
trinodb · Jan 29, 2019 · 5d9615c · 5d9615c
1 parent 0481f18
commit 5d9615c
Show file tree

Hide file tree

Showing 9 changed files with 79 additions and 20 deletions.
diff --git a/presto-docs/src/main/sphinx/sql/show-roles.rst b/presto-docs/src/main/sphinx/sql/show-roles.rst
@@ -7,9 +7,14 @@ Synopsis
 
 .. code-block:: none
 
-    SHOW ROLES [ FROM catalog ]
+    SHOW [CURRENT] ROLES [ FROM catalog ]
 
 Description
 -----------
 
-List the roles in ``catalog`` or in the current catalog.
+``SHOW ROLES`` lists all the roles in ``catalog`` or in the
+current catalog if ``catalog`` is not specified.
+
+``SHOW CURRENT ROLES`` lists the enabled roles for the session
+in ``catalog, or in the current catalog if ``catalog`` is not
+specified.
diff --git a/presto-main/src/main/java/io/prestosql/sql/rewrite/ShowQueriesRewrite.java b/presto-main/src/main/java/io/prestosql/sql/rewrite/ShowQueriesRewrite.java
@@ -83,6 +83,7 @@
 import static com.google.common.base.Strings.nullToEmpty;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.prestosql.connector.informationSchema.InformationSchemaMetadata.TABLE_COLUMNS;
+import static io.prestosql.connector.informationSchema.InformationSchemaMetadata.TABLE_ENABLED_ROLES;
 import static io.prestosql.connector.informationSchema.InformationSchemaMetadata.TABLE_ROLES;
 import static io.prestosql.connector.informationSchema.InformationSchemaMetadata.TABLE_SCHEMATA;
 import static io.prestosql.connector.informationSchema.InformationSchemaMetadata.TABLE_TABLES;
@@ -265,10 +266,18 @@ protected Node visitShowRoles(ShowRoles node, Void context)
             }
 
             String catalog = node.getCatalog().map(c -> c.getValue().toLowerCase(ENGLISH)).orElseGet(() -> session.getCatalog().get());
-            accessControl.checkCanShowRoles(session.getRequiredTransactionId(), session.getIdentity(), catalog);
-            return simpleQuery(
-                    selectList(aliasedName("role_name", "Role")),
-                    from(catalog, TABLE_ROLES));
+
+            if (node.isCurrent()) {
+                return simpleQuery(
+                        selectList(aliasedName("role_name", "Role")),
+                        from(catalog, TABLE_ENABLED_ROLES));
+            }
+            else {
+                accessControl.checkCanShowRoles(session.getRequiredTransactionId(), session.getIdentity(), catalog);
+                return simpleQuery(
+                        selectList(aliasedName("role_name", "Role")),
+                        from(catalog, TABLE_ROLES));
+            }
         }
 
         @Override

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
@@ -99,7 +99,7 @@ statement
     | SHOW COLUMNS (FROM | IN) qualifiedName                           #showColumns
     | SHOW STATS FOR qualifiedName                                     #showStats
     | SHOW STATS FOR '(' querySpecification ')'                        #showStatsForQuery
-    | SHOW ROLES ((FROM | IN) identifier)?                             #showRoles
+    | SHOW CURRENT? ROLES ((FROM | IN) identifier)?                    #showRoles
     | DESCRIBE qualifiedName                                           #showColumns
     | DESC qualifiedName                                               #showColumns
     | SHOW FUNCTIONS                                                   #showFunctions

diff --git a/presto-parser/src/main/java/io/prestosql/sql/SqlFormatter.java b/presto-parser/src/main/java/io/prestosql/sql/SqlFormatter.java
@@ -1263,7 +1263,11 @@ public Void visitShowGrants(ShowGrants node, Integer indent)
         @Override
         protected Void visitShowRoles(ShowRoles node, Integer context)
         {
-            builder.append("SHOW ROLES");
+            builder.append("SHOW ");
+            if (node.isCurrent()) {
+                builder.append("CURRENT ");
+            }
+            builder.append("ROLES");
 
             if (node.getCatalog().isPresent()) {
                 builder.append(" FROM ")

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
@@ -942,7 +942,8 @@ public Node visitShowRoles(SqlBaseParser.ShowRolesContext context)
     {
         return new ShowRoles(
                 getLocation(context),
-                getIdentifierIfPresent(context.identifier()));
+                getIdentifierIfPresent(context.identifier()),
+                context.CURRENT() != null);
     }
 
     @Override

diff --git a/presto-parser/src/main/java/io/prestosql/sql/tree/ShowRoles.java b/presto-parser/src/main/java/io/prestosql/sql/tree/ShowRoles.java
@@ -26,28 +26,35 @@ public class ShowRoles
         extends Statement
 {
     private final Optional<Identifier> catalog;
+    private final boolean current;
 
-    public ShowRoles(Optional<Identifier> catalog)
+    public ShowRoles(Optional<Identifier> catalog, boolean current)
     {
-        this(Optional.empty(), catalog);
+        this(Optional.empty(), catalog, current);
     }
 
-    public ShowRoles(NodeLocation location, Optional<Identifier> catalog)
+    public ShowRoles(NodeLocation location, Optional<Identifier> catalog, boolean current)
     {
-        this(Optional.of(location), catalog);
+        this(Optional.of(location), catalog, current);
     }
 
-    public ShowRoles(Optional<NodeLocation> location, Optional<Identifier> catalog)
+    public ShowRoles(Optional<NodeLocation> location, Optional<Identifier> catalog, boolean current)
     {
         super(location);
         this.catalog = requireNonNull(catalog, "catalog is null");
+        this.current = current;
     }
 
     public Optional<Identifier> getCatalog()
     {
         return catalog;
     }
 
+    public boolean isCurrent()
+    {
+        return current;
+    }
+
     @Override
     public <R, C> R accept(AstVisitor<R, C> visitor, C context)
     {
@@ -63,7 +70,7 @@ public List<Node> getChildren()
     @Override
     public int hashCode()
     {
-        return Objects.hash(catalog);
+        return Objects.hash(catalog, current);
     }
 
     @Override
@@ -76,14 +83,16 @@ public boolean equals(Object obj)
             return false;
         }
         ShowRoles o = (ShowRoles) obj;
-        return Objects.equals(catalog, o.catalog);
+        return Objects.equals(catalog, o.catalog) &&
+                current == o.current;
     }
 
     @Override
     public String toString()
     {
         return toStringHelper(this)
                 .add("catalog", catalog)
+                .add("current", current)
                 .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
@@ -1448,11 +1448,18 @@ public void testShowRoles()
             throws Exception
     {
         assertStatement("SHOW ROLES",
-                new ShowRoles(Optional.empty()));
+                new ShowRoles(Optional.empty(), false));
         assertStatement("SHOW ROLES FROM foo",
-                new ShowRoles(Optional.of(new Identifier("foo"))));
+                new ShowRoles(Optional.of(new Identifier("foo")), false));
         assertStatement("SHOW ROLES IN foo",
-                new ShowRoles(Optional.of(new Identifier("foo"))));
+                new ShowRoles(Optional.of(new Identifier("foo")), false));
+
+        assertStatement("SHOW CURRENT ROLES",
+                new ShowRoles(Optional.empty(), true));
+        assertStatement("SHOW CURRENT ROLES FROM foo",
+                new ShowRoles(Optional.of(new Identifier("foo")), true));
+        assertStatement("SHOW CURRENT ROLES IN foo",
+                new ShowRoles(Optional.of(new Identifier("foo")), true));
     }
 
     @Test

diff --git a/presto-parser/src/test/java/io/prestosql/sql/parser/TestStatementBuilder.java b/presto-parser/src/test/java/io/prestosql/sql/parser/TestStatementBuilder.java
@@ -248,6 +248,8 @@ public void testStatementBuilder()
         printStatement("show grants");
         printStatement("show roles");
         printStatement("show roles from foo");
+        printStatement("show current roles");
+        printStatement("show current roles from foo");
 
         printStatement("prepare p from select * from (select * from T) \"A B\"");
 

diff --git a/presto-product-tests/src/main/java/io/prestosql/tests/hive/TestRoles.java b/presto-product-tests/src/main/java/io/prestosql/tests/hive/TestRoles.java
@@ -491,7 +491,7 @@ public void testSetAdminRole()
                         row("admin"));
     }
 
-    @Test(groups = {HIVE_CONNECTOR, ROLES, AUTHORIZATION, PROFILE_SPECIFIC_TESTS})
+    @Test(groups = {ROLES, AUTHORIZATION, PROFILE_SPECIFIC_TESTS})
     public void testShowRoles()
     {
         QueryAssert.assertThat(onPresto().executeQuery("SHOW ROLES"))
@@ -514,6 +514,28 @@ public void testShowRoles()
                         row("role1"));
     }
 
+    @Test(groups = {ROLES, AUTHORIZATION, PROFILE_SPECIFIC_TESTS})
+    public void testShowCurrentRoles()
+    {
+        QueryAssert.assertThat(onPresto().executeQuery("SHOW CURRENT ROLES"))
+                .containsOnly(
+                        row("public"));
+        onPresto().executeQuery("CREATE ROLE role1");
+        onPresto().executeQuery("CREATE ROLE role2");
+        onPresto().executeQuery("GRANT role1 TO alice");
+        onPresto().executeQuery("GRANT role2 TO alice");
+        QueryAssert.assertThat(onPrestoAlice().executeQuery("SHOW CURRENT ROLES"))
+                .containsOnly(
+                        row("public"),
+                        row("role1"),
+                        row("role2"));
+        onPrestoAlice().executeQuery("SET ROLE role2");
+        QueryAssert.assertThat(onPrestoAlice().executeQuery("SHOW CURRENT ROLES"))
+                .containsOnly(
+                        row("public"),
+                        row("role2"));
+    }
+
     private static QueryExecutor onPrestoAlice()
     {
         return connectToPresto("alice@presto");