Support session property and extra credentials with comma or equality…

… sign
trinodb · Mar 13, 2019 · db1fa3a · db1fa3a
1 parent 63ec248
commit db1fa3a
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 6 deletions.
diff --git a/presto-client/src/main/java/io/prestosql/client/StatementClientV1.java b/presto-client/src/main/java/io/prestosql/client/StatementClientV1.java
@@ -177,7 +177,7 @@ private Request buildQueryRequest(ClientSession session, String query)
 
         Map<String, String> property = session.getProperties();
         for (Entry<String, String> entry : property.entrySet()) {
-            builder.addHeader(PRESTO_SESSION, entry.getKey() + "=" + entry.getValue());
+            builder.addHeader(PRESTO_SESSION, entry.getKey() + "=" + urlEncode(entry.getValue()));
         }
 
         Map<String, String> resourceEstimates = session.getResourceEstimates();
@@ -192,7 +192,7 @@ private Request buildQueryRequest(ClientSession session, String query)
 
         Map<String, String> extraCredentials = session.getExtraCredentials();
         for (Entry<String, String> entry : extraCredentials.entrySet()) {
-            builder.addHeader(PRESTO_EXTRA_CREDENTIAL, entry.getKey() + "=" + entry.getValue());
+            builder.addHeader(PRESTO_EXTRA_CREDENTIAL, entry.getKey() + "=" + urlEncode(entry.getValue()));
         }
 
         Map<String, String> statements = session.getPreparedStatements();
@@ -415,7 +415,7 @@ private void processResponse(Headers headers, QueryResults results)
             if (keyValue.size() != 2) {
                 continue;
             }
-            setSessionProperties.put(keyValue.get(0), keyValue.get(1));
+            setSessionProperties.put(keyValue.get(0), urlDecode(keyValue.get(1)));
         }
         resetSessionProperties.addAll(headers.values(PRESTO_CLEAR_SESSION));
 

diff --git a/presto-jdbc/src/test/java/io/prestosql/jdbc/TestJdbcConnection.java b/presto-jdbc/src/test/java/io/prestosql/jdbc/TestJdbcConnection.java
@@ -53,6 +53,7 @@
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
 
 public class TestJdbcConnection
 {
@@ -186,6 +187,28 @@ public void testSession()
             assertThat(listSession(connection))
                     .contains("join_distribution_type|BROADCAST|PARTITIONED")
                     .contains("exchange_compression|true|false");
+
+            try (Statement statement = connection.createStatement()) {
+                // setting Hive session properties requires the admin role
+                statement.execute("SET ROLE admin");
+            }
+
+            for (String part : ImmutableList.of(",", "=", ":", "|", "/", "\\", "'", "\\'", "''", "\"", "\\\"", "[", "]")) {
+                String value = format("/tmp/presto-%s-${USER}", part);
+                try {
+                    try (Statement statement = connection.createStatement()) {
+                        statement.execute(format("SET SESSION hive.temporary_staging_directory_path = '%s'", value.replace("'", "''")));
+                    }
+
+                    assertThat(listSession(connection))
+                            .contains("join_distribution_type|BROADCAST|PARTITIONED")
+                            .contains("exchange_compression|true|false")
+                            .contains(format("hive.temporary_staging_directory_path|%s|/tmp/presto-${USER}", value));
+                }
+                catch (Exception e) {
+                    fail(format("Failed to set session property value to [%s]", value), e);
+                }
+            }
         }
     }
 

diff --git a/presto-main/src/main/java/io/prestosql/server/HttpRequestSessionContext.java b/presto-main/src/main/java/io/prestosql/server/HttpRequestSessionContext.java
@@ -317,7 +317,12 @@ private static Map<String, String> parseProperty(HttpServletRequest servletReque
         for (String header : splitSessionHeader(servletRequest.getHeaders(headerName))) {
             List<String> nameValue = Splitter.on('=').trimResults().splitToList(header);
             assertRequest(nameValue.size() == 2, "Invalid %s header", headerName);
-            properties.put(nameValue.get(0), nameValue.get(1));
+            try {
+                properties.put(nameValue.get(0), urlDecode(nameValue.get(1)));
+            }
+            catch (IllegalArgumentException e) {
+                throw badRequest(format("Invalid %s header: %s", headerName, e));
+            }
         }
         return properties;
     }

diff --git a/presto-main/src/main/java/io/prestosql/server/protocol/StatementResource.java b/presto-main/src/main/java/io/prestosql/server/protocol/StatementResource.java
@@ -220,7 +220,7 @@ private static Response toResponse(Query query, QueryResults queryResults)
 
         // add set session properties
         query.getSetSessionProperties().entrySet()
-                .forEach(entry -> response.header(PRESTO_SET_SESSION, entry.getKey() + '=' + entry.getValue()));
+                .forEach(entry -> response.header(PRESTO_SET_SESSION, entry.getKey() + '=' + urlEncode(entry.getValue())));
 
         // add clear session properties
         query.getResetSessionProperties()

diff --git a/presto-main/src/test/java/io/prestosql/server/TestHttpRequestSessionContext.java b/presto-main/src/test/java/io/prestosql/server/TestHttpRequestSessionContext.java
@@ -58,6 +58,7 @@ public void testSessionContext()
                         .put(PRESTO_CLIENT_INFO, "client-info")
                         .put(PRESTO_SESSION, QUERY_MAX_MEMORY + "=1GB")
                         .put(PRESTO_SESSION, JOIN_DISTRIBUTION_TYPE + "=partitioned," + HASH_PARTITION_COUNT + " = 43")
+                        .put(PRESTO_SESSION, "some_session_property=some value with %2C comma")
                         .put(PRESTO_PREPARED_STATEMENT, "query1=select * from foo,query2=select * from bar")
                         .put(PRESTO_ROLE, "foo_connector=ALL")
                         .put(PRESTO_ROLE, "bar_connector=NONE")
@@ -76,7 +77,11 @@ public void testSessionContext()
         assertEquals(context.getClientInfo(), "client-info");
         assertEquals(context.getLanguage(), "zh-TW");
         assertEquals(context.getTimeZoneId(), "Asia/Taipei");
-        assertEquals(context.getSystemProperties(), ImmutableMap.of(QUERY_MAX_MEMORY, "1GB", JOIN_DISTRIBUTION_TYPE, "partitioned", HASH_PARTITION_COUNT, "43"));
+        assertEquals(context.getSystemProperties(), ImmutableMap.of(
+                QUERY_MAX_MEMORY, "1GB",
+                JOIN_DISTRIBUTION_TYPE, "partitioned",
+                HASH_PARTITION_COUNT, "43",
+                "some_session_property", "some value with , comma"));
         assertEquals(context.getPreparedStatements(), ImmutableMap.of("query1", "select * from foo", "query2", "select * from bar"));
         assertEquals(context.getIdentity().getRoles(), ImmutableMap.of(
                 "foo_connector", new SelectedRole(SelectedRole.Type.ALL, Optional.empty()),