Fix inconsistent session value from both server and client side

Previously, server only encode session value as http header content in response, but do not finish decoding and encoding at client, and decoding at server when receive request. Complete corresponding decoding and encoding from both server and client side.
prestodb · Jul 14, 2020 · 829b667 · 829b667
1 parent 20a1726
commit 829b667
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 4 deletions.
diff --git a/presto-client/src/main/java/com/facebook/presto/client/StatementClientV1.java b/presto-client/src/main/java/com/facebook/presto/client/StatementClientV1.java
@@ -169,7 +169,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();
@@ -398,7 +398,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/pom.xml b/presto-jdbc/pom.xml
@@ -32,6 +32,12 @@
             <artifactId>presto-common</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>com.facebook.presto</groupId>
+            <artifactId>presto-tests</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>com.squareup.okhttp3</groupId>
             <artifactId>okhttp</artifactId>

diff --git a/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestPrestoDriver.java b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestPrestoDriver.java
@@ -31,6 +31,8 @@
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.VarbinaryType;
 import com.facebook.presto.execution.QueryState;
+import com.facebook.presto.metadata.Catalog;
+import com.facebook.presto.metadata.SessionPropertyManager;
 import com.facebook.presto.plugin.blackhole.BlackHolePlugin;
 import com.facebook.presto.server.testing.TestingPrestoServer;
 import com.facebook.presto.spi.security.SelectedRole;
@@ -84,6 +86,9 @@
 import static com.facebook.presto.common.type.VarcharType.createVarcharType;
 import static com.facebook.presto.execution.QueryState.FAILED;
 import static com.facebook.presto.execution.QueryState.RUNNING;
+import static com.facebook.presto.testing.TestingSession.TESTING_CATALOG;
+import static com.facebook.presto.testing.TestingSession.createBogusTestingCatalog;
+import static com.facebook.presto.tests.AbstractTestQueries.TEST_CATALOG_PROPERTIES;
 import static io.airlift.units.Duration.nanosSince;
 import static java.lang.Float.POSITIVE_INFINITY;
 import static java.lang.String.format;
@@ -121,6 +126,10 @@ public void setup()
         server.createCatalog(TEST_CATALOG, "tpch");
         server.installPlugin(new BlackHolePlugin());
         server.createCatalog("blackhole", "blackhole");
+        Catalog bogusTestingCatalog = createBogusTestingCatalog(TESTING_CATALOG);
+        server.getCatalogManager().registerCatalog(bogusTestingCatalog);
+        SessionPropertyManager sessionPropertyManager = server.getMetadata().getSessionPropertyManager();
+        sessionPropertyManager.addConnectorSessionProperties(bogusTestingCatalog.getConnectorId(), TEST_CATALOG_PROPERTIES);
         waitForNodeRefresh(server);
         setupTestTables();
         executorService = newCachedThreadPool(daemonThreadsNamed("test-%s"));
@@ -379,7 +388,7 @@ public void testGetCatalogs()
     {
         try (Connection connection = createConnection()) {
             try (ResultSet rs = connection.getMetaData().getCatalogs()) {
-                assertEquals(readRows(rs), list(list("blackhole"), list("system"), list(TEST_CATALOG)));
+                assertEquals(readRows(rs), list(list("blackhole"), list("system"), list(TEST_CATALOG), list(TESTING_CATALOG)));
 
                 ResultSetMetaData metadata = rs.getMetaData();
                 assertEquals(metadata.getColumnCount(), 1);
@@ -1643,6 +1652,32 @@ public void testUpdatePartialCancel()
         }
     }
 
+    @Test
+    public void testEncodeDecodeSessionValue()
+            throws Exception
+    {
+        boolean isValidSessionValue = false;
+        String targetValue = "a+1==3";
+        try (Connection connection = createConnection(TESTING_CATALOG)) {
+            try (Statement statement = connection.createStatement()) {
+                assertFalse(statement.execute(format("set session %s.connector_string='%s'", TESTING_CATALOG, targetValue)));
+                assertNull(statement.getResultSet());
+
+                assertTrue(statement.execute("show session"));
+                ResultSet rs = statement.getResultSet();
+                while (rs.next()) {
+                    String sessionName = rs.getString("Name");
+                    if (sessionName.equals(format("%s.connector_string", TESTING_CATALOG))) {
+                        assertEquals(rs.getString("Value"), targetValue);
+                        isValidSessionValue = true;
+                        break;
+                    }
+                }
+            }
+        }
+        assertTrue(isValidSessionValue);
+    }
+
     private QueryState getQueryState(String queryId)
             throws SQLException
     {

diff --git a/presto-main/src/main/java/com/facebook/presto/server/HttpRequestSessionContext.java b/presto-main/src/main/java/com/facebook/presto/server/HttpRequestSessionContext.java
@@ -200,7 +200,7 @@ 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));
+            properties.put(nameValue.get(0), urlDecode(nameValue.get(1)));
         }
         return properties;
     }