Merge pull request #2041 from ClickHouse/jdbc2_fix_connections

[jdbc-v2] Fixed closing prev result set

Author: Paultagoras

jdbc-v2/src/main/java/com/clickhouse/jdbc/StatementImpl.java

@@ -152,6 +152,18 @@ public ResultSetImpl executeQuery(String sql, QuerySettings settings) throws SQL
                 response = connection.client.query(lastSql, mergedSettings).get(queryTimeout, TimeUnit.SECONDS);
             }
             ClickHouseBinaryFormatReader reader = connection.client.newBinaryFormatReader(response);
+
+            if (currentResultSet != null) {
+                LOG.debug("Previous result set is open [resultSet = " + currentResultSet + "]");
+                // Closing request blindly assuming that user do not care about it anymore (DDL request for example)
+                try {
+                    currentResultSet.close();
+                } catch (Exception e) {
+                    LOG.error("Failed to close previous result set", e);
+                } finally {
+                    currentResultSet = null;
+                }
+            }
             currentResultSet = new ResultSetImpl(this, response, reader);
             metrics = response.getMetrics();
             lastQueryId = response.getQueryId();
@@ -289,6 +301,7 @@ private boolean execute(String sql, QuerySettings settings) throws SQLException
             executeQuery(sql, settings); // keep open to allow getResultSet()
             return true;
         } else if(type == StatementType.SET) {
+            executeUpdate(sql, settings);
             //SET ROLE
             List<String> tokens = JdbcUtils.tokenizeSQL(sql);
             if (JdbcUtils.containsIgnoresCase(tokens, "ROLE")) {

jdbc-v2/src/test/java/com/clickhouse/jdbc/StatementTest.java

@@ -1,5 +1,7 @@
 package com.clickhouse.jdbc;
 
+import com.clickhouse.client.api.Client;
+import com.clickhouse.client.api.ClientConfigProperties;
 import com.clickhouse.client.api.query.GenericRecord;
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -16,6 +18,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Properties;
+import java.util.concurrent.TimeUnit;
 
 import static org.testng.Assert.*;
 import static org.testng.Assert.assertEquals;
@@ -361,28 +364,30 @@ private void testSettingRole() throws SQLException {
 
         try (ConnectionImpl conn = new ConnectionImpl(getEndpointString(), info)) {
             GenericRecord record = conn.client.queryAll("SELECT currentRoles()").get(0);
-            assertEquals(record.getList(1).size(), 2);
+            assertEquals(record.getList(1).size(), 0);
 
             try (Statement stmt = conn.createStatement()) {
                 stmt.execute("SET ROLE role1");
             }
 
             record = conn.client.queryAll("SELECT currentRoles()").get(0);
             assertEquals(record.getList(1).size(), 1);
+            assertEquals(record.getList(1).get(0), "role1");
 
             try (Statement stmt = conn.createStatement()) {
                 stmt.execute("SET ROLE role2");
             }
 
             record = conn.client.queryAll("SELECT currentRoles()").get(0);
             assertEquals(record.getList(1).size(), 1);
+            assertEquals(record.getList(1).get(0), "role2");
 
             try (Statement stmt = conn.createStatement()) {
                 stmt.execute("SET ROLE NONE");
             }
 
             record = conn.client.queryAll("SELECT currentRoles()").get(0);
-            assertEquals(record.getList(1).size(), 2);
+            assertEquals(record.getList(1).size(), 0);
         }
     }
 
@@ -449,4 +454,21 @@ public void testWithIPs() throws Exception {
             }
         }
     }
+
+    @Test
+    public void testConnectionExhaustion() throws Exception {
+
+        int maxNumConnections = 3;
+        Properties properties = new Properties();
+        properties.put(ClientConfigProperties.HTTP_MAX_OPEN_CONNECTIONS.getKey(), "" + maxNumConnections);
+        properties.put(ClientConfigProperties.CONNECTION_REQUEST_TIMEOUT.getKey(), "" + 1000); // 1 sec connection req timeout
+
+        try (Connection conn = getJdbcConnection(properties)) {
+            try (Statement stmt = conn.createStatement()) {
+                for (int i = 0; i< maxNumConnections * 2; i++) {
+                    stmt.executeQuery("SELECT number FROM system.numbers LIMIT 100");
+                }
+            }
+        }
+    }
 }