fixed for HttpURLConnection. Fixed bugwith one role

chernser · chernser · commit 6ce419c5879a · 2024-04-23T20:15:24.000-07:00
diff --git a/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ApacheHttpConnectionImpl.java b/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ApacheHttpConnectionImpl.java
@@ -60,6 +60,7 @@
 import java.net.Socket;
 import java.net.StandardSocketOptions;
 import java.nio.charset.StandardCharsets;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.TimeZone;
@@ -76,7 +77,7 @@ public class ApacheHttpConnectionImpl extends ClickHouseHttpConnection {
 
     protected ApacheHttpConnectionImpl(ClickHouseNode server, ClickHouseRequest<?> request, ExecutorService executor)
             throws IOException {
-        super(server, request);
+        super(server, request, Collections.emptyMap());
 
         client = newConnection(config);
     }
diff --git a/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpClient.java b/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpClient.java
@@ -114,7 +114,8 @@ protected ClickHouseHttpConnection newConnection(ClickHouseHttpConnection connec
         }
 
         try {
-            return ClickHouseHttpConnectionFactory.createConnection(server, request, getExecutor());
+
+            return ClickHouseHttpConnectionFactory.createConnection(server, request, getExecutor(), buildAdditionalReqParams(request));
         } catch (IOException e) {
             throw new CompletionException(e);
         }
@@ -146,6 +147,18 @@ protected String buildQueryParams(Map<String, String> params) {
         return builder.toString();
     }
 
+    private Map<String, Serializable> buildAdditionalReqParams(ClickHouseRequest<?> sealedRequest) {
+        ClickHouseConfig config = sealedRequest.getConfig();
+        if (config.getBoolOption(ClickHouseHttpOption.REMEMBER_LAST_SET_ROLES)) {
+            if (sealedRequest.hasSetting("_set_roles_stmt")) {
+                return Collections.singletonMap("_roles", sealedRequest.getSettings().get("_set_roles_stmt"));
+            } else if (!roles.isEmpty()) {
+                return Collections.singletonMap("_roles", roles);
+            }
+        }
+        return Collections.emptyMap();
+    }
+
     @Override
     protected ClickHouseResponse send(ClickHouseRequest<?> sealedRequest) throws ClickHouseException, IOException {
         ClickHouseHttpConnection conn = getConnection(sealedRequest);
@@ -174,19 +187,10 @@ protected ClickHouseResponse send(ClickHouseRequest<?> sealedRequest) throws Cli
                     }
                 }
                 : null;
-        Map<String, Serializable> additionalParams = null;
-
-        if (config.getBoolOption(ClickHouseHttpOption.REMEMBER_LAST_SET_ROLES)) {
-            if (sealedRequest.hasSetting("_set_roles_stmt")) {
-                additionalParams = Collections.singletonMap("_roles", sealedRequest.getSettings().get("_set_roles_stmt"));
-            } else if (!roles.isEmpty()) {
-                additionalParams = Collections.singletonMap("_roles", roles);
-            }
-        } else {
-            additionalParams = Collections.emptyMap();
-        }
 
         if (conn.isReusable()) {
+            Map<String, Serializable> additionalParams = buildAdditionalReqParams(sealedRequest);
+
             ClickHouseNode server = sealedRequest.getServer();
             httpResponse = conn.post(config, sql, sealedRequest.getInputStream().orElse(null),
                     sealedRequest.getExternalTables(), sealedRequest.getOutputStream().orElse(null),
diff --git a/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpConnection.java b/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpConnection.java
@@ -374,7 +374,8 @@ protected static void postData(ClickHouseConfig config, byte[] boundary, String
     protected final Map<String, String> defaultHeaders;
     protected final String url;
 
-    protected ClickHouseHttpConnection(ClickHouseNode server, ClickHouseRequest<?> request) {
+    protected ClickHouseHttpConnection(ClickHouseNode server, ClickHouseRequest<?> request,
+                                       Map<String, Serializable> additionalParams) {
         if (server == null || request == null) {
             throw new IllegalArgumentException("Non-null server and request are required");
         }
@@ -385,7 +386,7 @@ protected ClickHouseHttpConnection(ClickHouseNode server, ClickHouseRequest<?> r
         ClickHouseConfig c = request.getConfig();
         this.config = c;
         this.defaultHeaders = Collections.unmodifiableMap(createDefaultHeaders(c, server, getUserAgent(), ClickHouseHttpClient.getReferer(config)));
-        this.url = buildUrl(server.getBaseUri(), request, Collections.emptyMap());
+        this.url = buildUrl(server.getBaseUri(), request, additionalParams);
         log.debug("url [%s]", this.url);
     }
 
diff --git a/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpConnectionFactory.java b/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpConnectionFactory.java
@@ -1,6 +1,9 @@
 package com.clickhouse.client.http;
 
 import java.io.IOException;
+import java.io.Serializable;
+import java.util.Collections;
+import java.util.Map;
 import java.util.concurrent.ExecutorService;
 
 import com.clickhouse.client.ClickHouseNode;
@@ -14,7 +17,14 @@ public final class ClickHouseHttpConnectionFactory {
     private static final Logger log = LoggerFactory.getLogger(ClickHouseHttpConnectionFactory.class);
 
     public static ClickHouseHttpConnection createConnection(ClickHouseNode server, ClickHouseRequest<?> request,
-            ExecutorService executor) throws IOException {
+                                                            ExecutorService executor) throws IOException
+    {
+        return createConnection(server, request, executor, Collections.emptyMap());
+    }
+
+    public static ClickHouseHttpConnection createConnection(ClickHouseNode server, ClickHouseRequest<?> request,
+                                                            ExecutorService executor, Map<String,
+            Serializable> additionalRequestParams) throws IOException {
         HttpConnectionProvider provider = request.getConfig().getOption(ClickHouseHttpOption.CONNECTION_PROVIDER,
                 HttpConnectionProvider.class);
         if (provider == HttpConnectionProvider.APACHE_HTTP_CLIENT) {
@@ -27,7 +37,7 @@ public static ClickHouseHttpConnection createConnection(ClickHouseNode server, C
             log.warn("HTTP_CLIENT is only supported in JDK 11 or above, fall back to HTTP_URL_CONNECTION");
         }
 
-        return new HttpUrlConnectionImpl(server, request, executor);
+        return new HttpUrlConnectionImpl(server, request, executor, additionalRequestParams);
     }
 
     private ClickHouseHttpConnectionFactory() {
diff --git a/clickhouse-http-client/src/main/java/com/clickhouse/client/http/HttpUrlConnectionImpl.java b/clickhouse-http-client/src/main/java/com/clickhouse/client/http/HttpUrlConnectionImpl.java
@@ -17,15 +17,18 @@
 import com.clickhouse.logging.Logger;
 import com.clickhouse.logging.LoggerFactory;
 
+import javax.net.ssl.HostnameVerifier;
+import javax.net.ssl.HttpsURLConnection;
+import javax.net.ssl.SSLContext;
 import java.io.BufferedReader;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.io.OutputStream;
+import java.io.Serializable;
 import java.io.UncheckedIOException;
-
 import java.net.ConnectException;
 import java.net.HttpURLConnection;
 import java.net.Proxy;
@@ -35,15 +38,11 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Set;
 import java.util.TimeZone;
-import java.util.Map.Entry;
 import java.util.concurrent.ExecutorService;
 
-import javax.net.ssl.HostnameVerifier;
-import javax.net.ssl.HttpsURLConnection;
-import javax.net.ssl.SSLContext;
-
 public class HttpUrlConnectionImpl extends ClickHouseHttpConnection {
     private static final Logger log = LoggerFactory.getLogger(HttpUrlConnectionImpl.class);
 
@@ -205,9 +204,9 @@ private void checkResponse(HttpURLConnection conn) throws IOException {
         }
     }
 
-    protected HttpUrlConnectionImpl(ClickHouseNode server, ClickHouseRequest<?> request, ExecutorService executor)
-            throws IOException {
-        super(server, request);
+    protected HttpUrlConnectionImpl(ClickHouseNode server, ClickHouseRequest<?> request, ExecutorService executor,
+                                    Map<String, Serializable> additionalParams) throws IOException {
+        super(server, request, additionalParams);
 
         conn = newConnection(url, true);
     }
diff --git a/clickhouse-http-client/src/test/java/com/clickhouse/client/http/ClickHouseHttpConnectionTest.java b/clickhouse-http-client/src/test/java/com/clickhouse/client/http/ClickHouseHttpConnectionTest.java
@@ -1,16 +1,5 @@
 package com.clickhouse.client.http;
 
-import java.io.IOException;
-import java.io.Serializable;
-import java.net.URI;
-import java.net.URL;
-import java.net.URLDecoder;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-
 import com.clickhouse.client.ClickHouseClient;
 import com.clickhouse.client.ClickHouseConfig;
 import com.clickhouse.client.ClickHouseNode;
@@ -21,18 +10,23 @@
 import com.clickhouse.data.ClickHouseFormat;
 import com.clickhouse.data.ClickHouseInputStream;
 import com.clickhouse.data.ClickHouseOutputStream;
-
 import org.apache.hc.core5.net.URIBuilder;
 import org.testng.Assert;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 import org.testng.collections.Sets;
-import org.testng.internal.invokers.Arguments;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.net.URI;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
 
 public class ClickHouseHttpConnectionTest {
     static class SimpleHttpConnection extends ClickHouseHttpConnection {
         protected SimpleHttpConnection(ClickHouseNode server, ClickHouseRequest<?> request) {
-            super(server, request);
+            super(server, request, Collections.emptyMap());
         }
 
         @Override
diff --git a/clickhouse-jdbc/src/main/javacc/ClickHouseSqlParser.jj b/clickhouse-jdbc/src/main/javacc/ClickHouseSqlParser.jj
@@ -155,6 +155,9 @@ public class ClickHouseSqlParser {
                     roleBuff.append(ch);
                 }
             }
+            if (roleBuff.length() > 0) {
+                roles.add(roleBuff.toString());
+            }
         }
         anyArgsListStart = -1;
     }
diff --git a/clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/AccessManagementTest.java b/clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/AccessManagementTest.java
@@ -1,9 +1,9 @@
 package com.clickhouse.jdbc;
 
-import com.clickhouse.client.ClickHouseClient;
 import com.clickhouse.client.ClickHouseException;
 import com.clickhouse.client.ClickHouseProtocol;
 import com.clickhouse.client.http.config.ClickHouseHttpOption;
+import com.clickhouse.client.http.config.HttpConnectionProvider;
 import com.clickhouse.data.ClickHouseVersion;
 import org.testng.Assert;
 import org.testng.annotations.DataProvider;
@@ -19,12 +19,14 @@
 public class AccessManagementTest extends JdbcIntegrationTest {
 
     @Test(groups = "integration", dataProvider = "setRolesArgsForTestSetRole")
-    public void testSetRoleDifferentConnections(String[] roles, String setRoleExpr, String[] activeRoles) throws SQLException {
+    public void testSetRoleDifferentConnections(String[] roles, String setRoleExpr, String[] activeRoles,
+                                                String connectionProvider) throws SQLException {
 
         String httpEndpoint = "http://" + getServerAddress(ClickHouseProtocol.HTTP) + "/";
         String url = String.format("jdbc:ch:%s", httpEndpoint);
         Properties properties = new Properties();
         properties.setProperty(ClickHouseHttpOption.REMEMBER_LAST_SET_ROLES.getKey(), "true");
+        properties.setProperty(ClickHouseHttpOption.CONNECTION_PROVIDER.getKey(), connectionProvider);
         ClickHouseDataSource dataSource = new ClickHouseDataSource(url, properties);
 
         try (Connection connection = dataSource.getConnection("access_dba", "123")) {
@@ -59,13 +61,20 @@ public void testSetRoleDifferentConnections(String[] roles, String setRoleExpr,
     @DataProvider(name = "setRolesArgsForTestSetRole")
     private static Object[][] setRolesArgsForTestSetRole() {
         return new Object[][]{
-                {new String[]{"ROL1", "ROL2"}, "set role ROL2, ROL1", new String[]{"ROL2"}},
-                {new String[]{"ROL1", "ROL2"}, "set role ROL1,ROL2", new String[]{"ROL1"}},
-                {new String[]{"ROL1", "\"ROL2,☺\""}, "set role  \"ROL2,☺\", ROL1", new String[]{"ROL2,☺"}},
+                {new String[]{"ROL1", "ROL2"}, "set role ROL2", new String[]{"ROL2"},
+                        HttpConnectionProvider.HTTP_URL_CONNECTION.name()},
+                {new String[]{"ROL1", "ROL2"}, "set role ROL2", new String[]{"ROL2"},
+                        HttpConnectionProvider.APACHE_HTTP_CLIENT.name()},
+                {new String[]{"ROL1", "ROL2"}, "set role ROL2, ROL1", new String[]{"ROL1", "ROL2"},
+                        HttpConnectionProvider.APACHE_HTTP_CLIENT.name()},
+                {new String[]{"ROL1", "\"ROL2,☺\""}, "set role  \"ROL2,☺\", ROL1", new String[]{"ROL2,☺", "ROL1"},
+                        HttpConnectionProvider.APACHE_HTTP_CLIENT.name()},
+                {new String[]{"ROL1", "ROL2"}, "set role  ROL2 ,   ROL1,  ", new String[]{"ROL", "ROL1"},
+                        HttpConnectionProvider.APACHE_HTTP_CLIENT.name()},
         };
     }
 
-    private void assertRolesEquals(Connection connection, String ...expected) {
+    private void assertRolesEquals(Connection connection, String... expected) {
         try {
             Statement st = connection.createStatement();
             ResultSet resultSet = st.executeQuery("select currentRoles()");
@@ -76,7 +85,7 @@ private void assertRolesEquals(Connection connection, String ...expected) {
             Arrays.sort(expected);
             System.out.print("Roles: ");
             for (String role : roles) {
-                System.out.print("'" + role+ "', ");
+                System.out.print("'" + role + "', ");
             }
             System.out.println();
             Assert.assertEquals(roles, expected);

Original file line number	Diff line number	Diff line change
`@@ -155,6 +155,9 @@ public class ClickHouseSqlParser {`
`155`	`155`	`roleBuff.append(ch);`
`156`	`156`	`}`
`157`	`157`	`}`
	`158`	`+ if (roleBuff.length() > 0) {`
	`159`	`+ roles.add(roleBuff.toString());`
	`160`	`+ }`
`158`	`161`	`}`
`159`	`162`	`anyArgsListStart = -1;`
`160`	`163`	`}`