Implement SET ROLE

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

Hide file tree

Showing 48 changed files with 805 additions and 15 deletions.
diff --git a/...-benchmark-driver/src/main/java/io/prestosql/benchmark/driver/BenchmarkDriverOptions.java b/...-benchmark-driver/src/main/java/io/prestosql/benchmark/driver/BenchmarkDriverOptions.java
@@ -102,6 +102,7 @@ public ClientSession getClientSession()
                 ImmutableMap.of(),
                 toProperties(this.sessionProperties),
                 ImmutableMap.of(),
+                ImmutableMap.of(),
                 null,
                 clientRequestTimeout);
     }

diff --git a/presto-cli/src/main/java/io/prestosql/cli/ClientOptions.java b/presto-cli/src/main/java/io/prestosql/cli/ClientOptions.java
@@ -165,6 +165,7 @@ public ClientSession toClientSession()
                 toResourceEstimates(resourceEstimates),
                 toProperties(sessionProperties),
                 emptyMap(),
+                emptyMap(),
                 null,
                 clientRequestTimeout);
     }

diff --git a/presto-cli/src/main/java/io/prestosql/cli/Console.java b/presto-cli/src/main/java/io/prestosql/cli/Console.java
@@ -23,6 +23,7 @@
 import io.airlift.log.LoggingConfiguration;
 import io.airlift.units.Duration;
 import io.prestosql.cli.ClientOptions.OutputFormat;
+import io.prestosql.client.ClientSelectedRole;
 import io.prestosql.client.ClientSession;
 import io.prestosql.sql.parser.StatementSplitter;
 import jline.console.history.FileHistory;
@@ -346,6 +347,13 @@ private static boolean process(QueryRunner queryRunner, String sql, OutputFormat
                 builder = builder.withProperties(sessionProperties);
             }
 
+            // update session roles
+            if (!query.getSetRoles().isEmpty()) {
+                Map<String, ClientSelectedRole> roles = new HashMap<>(session.getRoles());
+                roles.putAll(query.getSetRoles());
+                builder = builder.withRoles(roles);
+            }
+
             // update prepared statements if present
             if (!query.getAddedPreparedStatements().isEmpty() || !query.getDeallocatedPreparedStatements().isEmpty()) {
                 Map<String, String> preparedStatements = new HashMap<>(session.getPreparedStatements());

diff --git a/presto-cli/src/main/java/io/prestosql/cli/Query.java b/presto-cli/src/main/java/io/prestosql/cli/Query.java
@@ -18,6 +18,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import io.prestosql.cli.ClientOptions.OutputFormat;
+import io.prestosql.client.ClientSelectedRole;
 import io.prestosql.client.Column;
 import io.prestosql.client.ErrorLocation;
 import io.prestosql.client.QueryError;
@@ -89,6 +90,11 @@ public Set<String> getResetSessionProperties()
         return client.getResetSessionProperties();
     }
 
+    public Map<String, ClientSelectedRole> getSetRoles()
+    {
+        return client.getSetRoles();
+    }
+
     public Map<String, String> getAddedPreparedStatements()
     {
         return client.getAddedPreparedStatements();

diff --git a/presto-cli/src/test/java/io/prestosql/cli/TestQueryRunner.java b/presto-cli/src/test/java/io/prestosql/cli/TestQueryRunner.java
@@ -99,6 +99,7 @@ public void testCookie()
                         ImmutableMap.of(),
                         ImmutableMap.of(),
                         ImmutableMap.of(),
+                        ImmutableMap.of(),
                         null,
                         new Duration(2, MINUTES)));
         try (Query query = queryRunner.startQuery("first query will introduce a cookie")) {

diff --git a/presto-client/src/main/java/io/prestosql/client/ClientSelectedRole.java b/presto-client/src/main/java/io/prestosql/client/ClientSelectedRole.java
@@ -0,0 +1,103 @@
+/*
+ * 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.client;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+import java.util.Optional;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * TODO: This class is a duplicate of SelectedRole class from SPI. Think about better way
+ * of implementing it without duplicating code.
+ */
+public class ClientSelectedRole
+{
+    public enum Type
+    {
+        ROLE, ALL, NONE
+    }
+
+    private static final Pattern PATTERN = Pattern.compile("(ROLE|ALL|NONE)(\\{(.+?)\\})?");
+
+    private final Type type;
+    private final Optional<String> role;
+
+    @JsonCreator
+    public ClientSelectedRole(@JsonProperty("type") Type type, @JsonProperty("role") Optional<String> role)
+    {
+        this.type = requireNonNull(type, "type is null");
+        this.role = requireNonNull(role, "role is null");
+        if (type == Type.ROLE && !role.isPresent()) {
+            throw new IllegalArgumentException("Role must be present for the selected role type: " + type);
+        }
+    }
+
+    @JsonProperty
+    public Type getType()
+    {
+        return type;
+    }
+
+    @JsonProperty
+    public Optional<String> getRole()
+    {
+        return role;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        ClientSelectedRole that = (ClientSelectedRole) o;
+        return type == that.type &&
+                Objects.equals(role, that.role);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(type, role);
+    }
+
+    @Override
+    public String toString()
+    {
+        StringBuilder result = new StringBuilder();
+        result.append(type);
+        role.ifPresent(s -> result.append("{").append(s).append("}"));
+        return result.toString();
+    }
+
+    public static ClientSelectedRole valueOf(String value)
+    {
+        Matcher m = PATTERN.matcher(value);
+        if (m.matches()) {
+            Type type = Type.valueOf(m.group(1));
+            Optional<String> role = Optional.ofNullable(m.group(3));
+            return new ClientSelectedRole(type, role);
+        }
+        throw new IllegalArgumentException("Could not parse selected role: " + value);
+    }
+}
diff --git a/presto-client/src/main/java/io/prestosql/client/ClientSession.java b/presto-client/src/main/java/io/prestosql/client/ClientSession.java
@@ -47,6 +47,7 @@ public class ClientSession
     private final Map<String, String> resourceEstimates;
     private final Map<String, String> properties;
     private final Map<String, String> preparedStatements;
+    private final Map<String, ClientSelectedRole> roles;
     private final String transactionId;
     private final Duration clientRequestTimeout;
 
@@ -77,6 +78,7 @@ public ClientSession(
             Map<String, String> resourceEstimates,
             Map<String, String> properties,
             Map<String, String> preparedStatements,
+            Map<String, ClientSelectedRole> roles,
             String transactionId,
             Duration clientRequestTimeout)
     {
@@ -95,6 +97,7 @@ public ClientSession(
         this.resourceEstimates = ImmutableMap.copyOf(requireNonNull(resourceEstimates, "resourceEstimates is null"));
         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;
 
         for (String clientTag : clientTags) {
@@ -188,6 +191,14 @@ public Map<String, String> getPreparedStatements()
         return preparedStatements;
     }
 
+    /**
+     * Returns the map of catalog name -> selected role
+     */
+    public Map<String, ClientSelectedRole> getRoles()
+    {
+        return roles;
+    }
+
     public String getTransactionId()
     {
         return transactionId;
@@ -239,6 +250,7 @@ public static final class Builder
         private Map<String, String> resourceEstimates;
         private Map<String, String> properties;
         private Map<String, String> preparedStatements;
+        private Map<String, ClientSelectedRole> roles;
         private String transactionId;
         private Duration clientRequestTimeout;
 
@@ -259,6 +271,7 @@ private Builder(ClientSession clientSession)
             resourceEstimates = clientSession.getResourceEstimates();
             properties = clientSession.getProperties();
             preparedStatements = clientSession.getPreparedStatements();
+            roles = clientSession.getRoles();
             transactionId = clientSession.getTransactionId();
             clientRequestTimeout = clientSession.getClientRequestTimeout();
         }
@@ -287,6 +300,12 @@ public Builder withProperties(Map<String, String> properties)
             return this;
         }
 
+        public Builder withRoles(Map<String, ClientSelectedRole> roles)
+        {
+            this.roles = roles;
+            return this;
+        }
+
         public Builder withPreparedStatements(Map<String, String> preparedStatements)
         {
             this.preparedStatements = requireNonNull(preparedStatements, "preparedStatements is null");
@@ -322,6 +341,7 @@ public ClientSession build()
                     resourceEstimates,
                     properties,
                     preparedStatements,
+                    roles,
                     transactionId,
                     clientRequestTimeout);
         }

diff --git a/presto-client/src/main/java/io/prestosql/client/PrestoHeaders.java b/presto-client/src/main/java/io/prestosql/client/PrestoHeaders.java
@@ -29,6 +29,8 @@ public final class PrestoHeaders
     public static final String PRESTO_SET_PATH = "X-Presto-Set-Path";
     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";

diff --git a/presto-client/src/main/java/io/prestosql/client/StatementClient.java b/presto-client/src/main/java/io/prestosql/client/StatementClient.java
@@ -54,6 +54,8 @@ public interface StatementClient
 
     Set<String> getResetSessionProperties();
 
+    Map<String, ClientSelectedRole> getSetRoles();
+
     Map<String, String> getAddedPreparedStatements();
 
     Set<String> getDeallocatedPreparedStatements();

diff --git a/presto-client/src/main/java/io/prestosql/client/StatementClientV1.java b/presto-client/src/main/java/io/prestosql/client/StatementClientV1.java
@@ -65,6 +65,7 @@
 import static io.prestosql.client.PrestoHeaders.PRESTO_SESSION;
 import static io.prestosql.client.PrestoHeaders.PRESTO_SET_CATALOG;
 import static io.prestosql.client.PrestoHeaders.PRESTO_SET_PATH;
+import static io.prestosql.client.PrestoHeaders.PRESTO_SET_ROLE;
 import static io.prestosql.client.PrestoHeaders.PRESTO_SET_SCHEMA;
 import static io.prestosql.client.PrestoHeaders.PRESTO_SET_SESSION;
 import static io.prestosql.client.PrestoHeaders.PRESTO_SOURCE;
@@ -100,6 +101,7 @@ class StatementClientV1
     private final AtomicReference<String> setPath = new AtomicReference<>();
     private final Map<String, String> setSessionProperties = new ConcurrentHashMap<>();
     private final Set<String> resetSessionProperties = Sets.newConcurrentHashSet();
+    private final Map<String, ClientSelectedRole> 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<>();
@@ -182,6 +184,11 @@ private Request buildQueryRequest(ClientSession session, String query)
             builder.addHeader(PRESTO_RESOURCE_ESTIMATE, entry.getKey() + "=" + entry.getValue());
         }
 
+        Map<String, ClientSelectedRole> roles = session.getRoles();
+        for (Entry<String, ClientSelectedRole> 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(PRESTO_PREPARED_STATEMENT, urlEncode(entry.getKey()) + "=" + urlEncode(entry.getValue()));
@@ -283,6 +290,12 @@ public Set<String> getResetSessionProperties()
         return ImmutableSet.copyOf(resetSessionProperties);
     }
 
+    @Override
+    public Map<String, ClientSelectedRole> getSetRoles()
+    {
+        return ImmutableMap.copyOf(setRoles);
+    }
+
     @Override
     public Map<String, String> getAddedPreparedStatements()
     {
@@ -400,6 +413,14 @@ private void processResponse(Headers headers, QueryResults results)
         }
         resetSessionProperties.addAll(headers.values(PRESTO_CLEAR_SESSION));
 
+        for (String setRole : headers.values(PRESTO_SET_ROLE)) {
+            List<String> keyValue = SESSION_HEADER_SPLITTER.splitToList(setRole);
+            if (keyValue.size() != 2) {
+                continue;
+            }
+            setRoles.put(keyValue.get(0), ClientSelectedRole.valueOf(urlDecode(keyValue.get(1))));
+        }
+
         for (String entry : headers.values(PRESTO_ADDED_PREPARE)) {
             List<String> keyValue = SESSION_HEADER_SPLITTER.splitToList(entry);
             if (keyValue.size() != 2) {

diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/file/FileHiveMetastore.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/file/FileHiveMetastore.java
@@ -144,7 +144,7 @@ public FileHiveMetastore(HdfsEnvironment hdfsEnvironment, String catalogDirector
     {
         this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null");
         this.catalogDirectory = new Path(requireNonNull(catalogDirectory, "baseDirectory is null"));
-        this.hdfsContext = new HdfsContext(new ConnectorIdentity(metastoreUser, Optional.empty()));
+        this.hdfsContext = new HdfsContext(new ConnectorIdentity(metastoreUser, Optional.empty(), Optional.empty()));
         try {
             metadataFileSystem = hdfsEnvironment.getFileSystem(hdfsContext, this.catalogDirectory);
         }

diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastore.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastore.java
@@ -143,7 +143,7 @@ public GlueHiveMetastore(HdfsEnvironment hdfsEnvironment, GlueHiveMetastoreConfi
     public GlueHiveMetastore(HdfsEnvironment hdfsEnvironment, GlueHiveMetastoreConfig glueConfig, AWSGlueAsync glueClient)
     {
         this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null");
-        this.hdfsContext = new HdfsContext(new ConnectorIdentity(DEFAULT_METASTORE_USER, Optional.empty()));
+        this.hdfsContext = new HdfsContext(new ConnectorIdentity(DEFAULT_METASTORE_USER, Optional.empty(), Optional.empty()));
         this.glueClient = requireNonNull(glueClient, "glueClient is null");
         this.defaultDir = glueConfig.getDefaultWarehouseDir();
     }

diff --git a/...to-hive/src/main/java/io/prestosql/plugin/hive/security/PartitionsAwareAccessControl.java b/...to-hive/src/main/java/io/prestosql/plugin/hive/security/PartitionsAwareAccessControl.java
@@ -211,4 +211,10 @@ public void checkCanRevokeRoles(ConnectorTransactionHandle transactionHandle, Co
     {
         delegate.checkCanRevokeRoles(transactionHandle, identity, roles, grantees, adminOptionFor, grantor, catalogName);
     }
+
+    @Override
+    public void checkCanSetRole(ConnectorTransactionHandle transactionHandle, ConnectorIdentity identity, String role, String catalogName)
+    {
+        delegate.checkCanSetRole(transactionHandle, identity, role, catalogName);
+    }
 }
diff --git a/presto-hive/src/test/java/io/prestosql/plugin/hive/AbstractTestHiveFileSystem.java b/presto-hive/src/test/java/io/prestosql/plugin/hive/AbstractTestHiveFileSystem.java
@@ -105,7 +105,7 @@
 
 public abstract class AbstractTestHiveFileSystem
 {
-    private static final HdfsContext TESTING_CONTEXT = new HdfsContext(new ConnectorIdentity("test", Optional.empty()));
+    private static final HdfsContext TESTING_CONTEXT = new HdfsContext(new ConnectorIdentity("test", Optional.empty(), Optional.empty()));
 
     protected String database;
     protected SchemaTableName table;