[WIP] Support role management

presto-cli/pom.xml

@@ -24,6 +24,11 @@
             <artifactId>presto-client</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>com.facebook.presto</groupId>
+            <artifactId>presto-spi</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>com.facebook.presto</groupId>
             <artifactId>presto-parser</artifactId>

presto-cli/src/main/java/com/facebook/presto/cli/Console.java

@@ -15,6 +15,7 @@
 
 import com.facebook.presto.cli.ClientOptions.OutputFormat;
 import com.facebook.presto.client.ClientSession;
+import com.facebook.presto.spi.security.SelectedRole;
 import com.facebook.presto.sql.parser.IdentifierSymbol;
 import com.facebook.presto.sql.parser.ParsingException;
 import com.facebook.presto.sql.parser.SqlParser;
@@ -55,6 +56,7 @@
 import static com.facebook.presto.client.ClientSession.withCatalogAndSchema;
 import static com.facebook.presto.client.ClientSession.withPreparedStatements;
 import static com.facebook.presto.client.ClientSession.withProperties;
+import static com.facebook.presto.client.ClientSession.withRoles;
 import static com.facebook.presto.client.ClientSession.withTransactionId;
 import static com.facebook.presto.sql.parser.StatementSplitter.Statement;
 import static com.facebook.presto.sql.parser.StatementSplitter.isEmptyStatement;
@@ -327,6 +329,13 @@ private static void process(QueryRunner queryRunner, String sql, OutputFormat ou
                 session = withProperties(session, sessionProperties);
             }
 
+            // update session roles
+            if (!query.getSetRoles().isEmpty()) {
+                Map<String, SelectedRole> roles = new HashMap<>(session.getRoles());
+                roles.putAll(query.getSetRoles());
+                session = withRoles(session, roles);
+            }
+
             // update prepared statements if present
             if (!query.getAddedPreparedStatements().isEmpty() || !query.getDeallocatedPreparedStatements().isEmpty()) {
                 Map<String, String> preparedStatements = new HashMap<>(session.getPreparedStatements());

presto-cli/src/main/java/com/facebook/presto/cli/Query.java

@@ -19,6 +19,7 @@
 import com.facebook.presto.client.QueryError;
 import com.facebook.presto.client.QueryResults;
 import com.facebook.presto.client.StatementClient;
+import com.facebook.presto.spi.security.SelectedRole;
 import com.google.common.base.Splitter;
 import com.google.common.base.Strings;
 import com.google.common.base.Throwables;
@@ -72,6 +73,11 @@ public Set<String> getResetSessionProperties()
         return client.getResetSessionProperties();
     }
 
+    public Map<String, SelectedRole> getSetRoles()
+    {
+        return client.getSetRoles();
+    }
+
     public Map<String, String> getAddedPreparedStatements()
     {
         return client.getAddedPreparedStatements();

presto-client/src/main/java/com/facebook/presto/client/ClientSession.java

@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.client;
 
+import com.facebook.presto.spi.security.SelectedRole;
 import com.facebook.presto.spi.type.TimeZoneKey;
 import com.google.common.collect.ImmutableMap;
 import io.airlift.units.Duration;
@@ -41,6 +42,7 @@ public class ClientSession
     private final Locale locale;
     private final Map<String, String> properties;
     private final Map<String, String> preparedStatements;
+    private final Map<String, SelectedRole> roles;
     private final String transactionId;
     private final boolean debug;
     private final Duration clientRequestTimeout;
@@ -58,6 +60,7 @@ public static ClientSession withCatalogAndSchema(ClientSession session, String c
                 session.getLocale(),
                 session.getProperties(),
                 session.getPreparedStatements(),
+                session.getRoles(),
                 session.getTransactionId(),
                 session.isDebug(),
                 session.getClientRequestTimeout());
@@ -76,6 +79,26 @@ public static ClientSession withProperties(ClientSession session, Map<String, St
                 session.getLocale(),
                 properties,
                 session.getPreparedStatements(),
+                session.getRoles(),
+                session.getTransactionId(),
+                session.isDebug(),
+                session.getClientRequestTimeout());
+    }
+
+    public static ClientSession withRoles(ClientSession session, Map<String, SelectedRole> roles)
+    {
+        return new ClientSession(
+                session.getServer(),
+                session.getUser(),
+                session.getSource(),
+                session.getClientInfo(),
+                session.getCatalog(),
+                session.getSchema(),
+                session.getTimeZone().getId(),
+                session.getLocale(),
+                session.getProperties(),
+                session.getPreparedStatements(),
+                roles,
                 session.getTransactionId(),
                 session.isDebug(),
                 session.getClientRequestTimeout());
@@ -94,6 +117,7 @@ public static ClientSession withPreparedStatements(ClientSession session, Map<St
                 session.getLocale(),
                 session.getProperties(),
                 preparedStatements,
+                session.getRoles(),
                 session.getTransactionId(),
                 session.isDebug(),
                 session.getClientRequestTimeout());
@@ -112,6 +136,7 @@ public static ClientSession withTransactionId(ClientSession session, String tran
                 session.getLocale(),
                 session.getProperties(),
                 session.getPreparedStatements(),
+                session.getRoles(),
                 transactionId,
                 session.isDebug(),
                 session.getClientRequestTimeout());
@@ -130,6 +155,7 @@ public static ClientSession stripTransactionId(ClientSession session)
                 session.getLocale(),
                 session.getProperties(),
                 session.getPreparedStatements(),
+                session.getRoles(),
                 null,
                 session.isDebug(),
                 session.getClientRequestTimeout());
@@ -166,6 +192,25 @@ public ClientSession(
             String transactionId,
             boolean debug,
             Duration clientRequestTimeout)
+    {
+        this(server, user, source, clientInfo, catalog, schema, timeZoneId, locale, properties, preparedStatements, emptyMap(), transactionId, debug, clientRequestTimeout);
+    }
+
+    public ClientSession(
+            URI server,
+            String user,
+            String source,
+            String clientInfo,
+            String catalog,
+            String schema,
+            String timeZoneId,
+            Locale locale,
+            Map<String, String> properties,
+            Map<String, String> preparedStatements,
+            Map<String, SelectedRole> roles,
+            String transactionId,
+            boolean debug,
+            Duration clientRequestTimeout)
     {
         this.server = requireNonNull(server, "server is null");
         this.user = user;
@@ -179,6 +224,7 @@ public ClientSession(
         this.debug = debug;
         this.properties = ImmutableMap.copyOf(requireNonNull(properties, "properties is null"));
         this.preparedStatements = ImmutableMap.copyOf(requireNonNull(preparedStatements, "preparedStatements is null"));
+        this.roles = ImmutableMap.copyOf(requireNonNull(roles, "roles is null"));
         this.clientRequestTimeout = clientRequestTimeout;
 
         // verify the properties are valid
@@ -241,6 +287,14 @@ public Map<String, String> getPreparedStatements()
         return preparedStatements;
     }
 
+    /**
+     * Returns the map of catalog name -> selected role
+     */
+    public Map<String, SelectedRole> getRoles()
+    {
+        return roles;
+    }
+
     public String getTransactionId()
     {
         return transactionId;

presto-client/src/main/java/com/facebook/presto/client/PrestoHeaders.java

@@ -24,6 +24,8 @@ public final class PrestoHeaders
     public static final String PRESTO_SESSION = "X-Presto-Session";
     public static final String PRESTO_SET_SESSION = "X-Presto-Set-Session";
     public static final String PRESTO_CLEAR_SESSION = "X-Presto-Clear-Session";
+    public static final String PRESTO_SET_ROLE = "X-Presto-Set-Role";
+    public static final String PRESTO_ROLE = "X-Presto-Role";
     public static final String PRESTO_PREPARED_STATEMENT = "X-Presto-Prepared-Statement";
     public static final String PRESTO_ADDED_PREPARE = "X-Presto-Added-Prepare";
     public static final String PRESTO_DEALLOCATED_PREPARE = "X-Presto-Deallocated-Prepare";

presto-client/src/main/java/com/facebook/presto/client/StatementClient.java

@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.client;
 
+import com.facebook.presto.spi.security.SelectedRole;
 import com.facebook.presto.spi.type.TimeZoneKey;
 import com.google.common.base.Splitter;
 import com.google.common.base.Throwables;
@@ -48,6 +49,7 @@
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_CLEAR_SESSION;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_CLEAR_TRANSACTION_ID;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_DEALLOCATED_PREPARE;
+import static com.facebook.presto.client.PrestoHeaders.PRESTO_SET_ROLE;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_SET_SESSION;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_STARTED_TRANSACTION_ID;
 import static com.google.common.base.MoreObjects.firstNonNull;
@@ -86,6 +88,7 @@ public class StatementClient
     private final AtomicReference<QueryResults> currentResults = new AtomicReference<>();
     private final Map<String, String> setSessionProperties = new ConcurrentHashMap<>();
     private final Set<String> resetSessionProperties = Sets.newConcurrentHashSet();
+    private final Map<String, SelectedRole> setRoles = new ConcurrentHashMap<>();
     private final Map<String, String> addedPreparedStatements = new ConcurrentHashMap<>();
     private final Set<String> deallocatedPreparedStatements = Sets.newConcurrentHashSet();
     private final AtomicReference<String> startedtransactionId = new AtomicReference<>();
@@ -149,6 +152,11 @@ private Request buildQueryRequest(ClientSession session, String query)
             builder.addHeader(PrestoHeaders.PRESTO_SESSION, entry.getKey() + "=" + entry.getValue());
         }
 
+        Map<String, SelectedRole> roles = session.getRoles();
+        for (Entry<String, SelectedRole> entry : roles.entrySet()) {
+            builder.addHeader(PrestoHeaders.PRESTO_ROLE, entry.getKey() + '=' + urlEncode(entry.getValue().toString()));
+        }
+
         Map<String, String> statements = session.getPreparedStatements();
         for (Entry<String, String> entry : statements.entrySet()) {
             builder.addHeader(PrestoHeaders.PRESTO_PREPARED_STATEMENT, urlEncode(entry.getKey()) + "=" + urlEncode(entry.getValue()));
@@ -216,6 +224,11 @@ public Set<String> getResetSessionProperties()
         return ImmutableSet.copyOf(resetSessionProperties);
     }
 
+    public Map<String, SelectedRole> getSetRoles()
+    {
+        return ImmutableMap.copyOf(setRoles);
+    }
+
     public Map<String, String> getAddedPreparedStatements()
     {
         return ImmutableMap.copyOf(addedPreparedStatements);
@@ -319,6 +332,14 @@ private void processResponse(JsonResponse<QueryResults> response)
             resetSessionProperties.add(clearSession);
         }
 
+        for (String setRole : response.getHeaders(PRESTO_SET_ROLE)) {
+            List<String> keyValue = SESSION_HEADER_SPLITTER.splitToList(setRole);
+            if (keyValue.size() != 2) {
+                continue;
+            }
+            setRoles.put(keyValue.get(0), SelectedRole.valueOf(urlDecode(keyValue.get(1))));
+        }
+
         for (String entry : response.getHeaders(PRESTO_ADDED_PREPARE)) {
             List<String> keyValue = SESSION_HEADER_SPLITTER.splitToList(entry);
             if (keyValue.size() != 2) {

presto-docs/src/main/sphinx/connector/mysql.rst

@@ -67,3 +67,6 @@ The following SQL statements are not yet supported:
 * :doc:`/sql/create-table` (:doc:`/sql/create-table-as` is supported)
 * :doc:`/sql/grant`
 * :doc:`/sql/revoke`
+* :doc:`/sql/show-grants`
+* :doc:`/sql/show-roles`
+* :doc:`/sql/show-role-grants`

presto-docs/src/main/sphinx/connector/postgresql.rst

@@ -72,3 +72,6 @@ The following SQL statements are not yet supported:
 * :doc:`/sql/create-table` (:doc:`/sql/create-table-as` is supported)
 * :doc:`/sql/grant`
 * :doc:`/sql/revoke`
+* :doc:`/sql/show-grants`
+* :doc:`/sql/show-roles`
+* :doc:`/sql/show-role-grants`

presto-docs/src/main/sphinx/sql.rst

@@ -34,7 +34,10 @@ This chapter describes the SQL syntax used in Presto.
     sql/show-create-table
     sql/show-create-view
     sql/show-functions
+    sql/show-grants
     sql/show-partitions
+    sql/show-role-grants
+    sql/show-roles
     sql/show-schemas
     sql/show-session
     sql/show-tables

presto-docs/src/main/sphinx/sql/grant.rst

@@ -8,11 +8,9 @@ Synopsis
 .. code-block:: none
 
     GRANT ( privilege [, ...] | ( ALL PRIVILEGES ) )
-    ON [ TABLE ] table_name TO ( grantee | PUBLIC )
+    ON [ TABLE ] table_name TO ( user | USER user | ROLE role )
     [ WITH GRANT OPTION ]
 
-Usage of the term ``grantee`` denotes both users and roles.
-
 Description
 -----------
 
@@ -39,7 +37,7 @@ Grant ``SELECT`` privilege on the table ``nation`` to user ``alice``, additional
 
 Grant ``SELECT`` privilege on the table ``orders`` to everyone::
 
-    GRANT SELECT ON orders TO PUBLIC;
+    GRANT SELECT ON orders TO ROLE PUBLIC;
 
 Limitations
 -----------
@@ -51,3 +49,4 @@ See Also
 --------
 
 :doc:`revoke`
+:doc:`show-grants`

presto-docs/src/main/sphinx/sql/revoke.rst

@@ -9,9 +9,7 @@ Synopsis
 
     REVOKE [ GRANT OPTION FOR ]
     ( privilege [, ...] | ALL PRIVILEGES )
-    ON [ TABLE ] table_name FROM ( grantee | PUBLIC )
-
-Usage of the term ``grantee`` denotes both users and roles.
+    ON [ TABLE ] table_name FROM ( user | USER user | ROLE role )
 
 Description
 -----------
@@ -35,7 +33,7 @@ Revoke ``INSERT`` and ``SELECT`` privileges on the table ``orders`` from user ``
 
 Revoke ``SELECT`` privilege on the table ``nation`` from everyone, additionally revoking the privilege to grant ``SELECT`` privilege::
 
-    REVOKE GRANT OPTION FOR SELECT ON nation FROM PUBLIC;
+    REVOKE GRANT OPTION FOR SELECT ON nation FROM ROLE PUBLIC;
 
 Revoke all privileges on the table ``test`` from user ``alice``::
 
@@ -51,3 +49,4 @@ See Also
 --------
 
 :doc:`grant`
+:doc:`show-grants`

presto-docs/src/main/sphinx/sql/show-grants.rst

@@ -0,0 +1,46 @@
+===========
+SHOW GRANTS
+===========
+
+Synopsis
+--------
+
+.. code-block:: none
+
+    SHOW GRANTS ON ( [ TABLE ] table_name | ALL )
+
+Description
+-----------
+
+List the grants for the current user on the specified table in the current catalog.
+
+Specifying ``ALL`` lists the grants for the current user on all the tables in all the schemas in the current catalog.
+
+The command requires the current catalog to be set.
+
+.. note::
+
+    Ensure that authentication has been enabled before running any of the authorization commands.
+
+Examples
+--------
+
+List the grants for the current user on table ``orders``::
+
+    SHOW GRANTS ON TABLE ``orders``
+
+List the grants for the current user on all the tables in all the schemas in the current catalog::
+
+    SHOW GRANTS ON ALL;
+
+Limitations
+-----------
+
+Some connectors have no support for ``SHOW GRANTS``.
+See connector documentation for more details.
+
+See Also
+--------
+
+:doc:`grant`
+:doc:`revoke`