improved code and added one more test

Author: chernser

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

@@ -2,15 +2,11 @@
 
 import com.clickhouse.client.api.Client;
 import com.clickhouse.client.api.ClientConfigProperties;
-import com.clickhouse.client.api.ClientException;
-import com.clickhouse.client.api.data_formats.ClickHouseBinaryFormatReader;
 import com.clickhouse.client.api.internal.ServerSettings;
 import com.clickhouse.client.api.metadata.TableSchema;
 import com.clickhouse.client.api.query.GenericRecord;
-import com.clickhouse.client.api.query.QueryResponse;
 import com.clickhouse.client.api.query.QuerySettings;
 import com.clickhouse.data.ClickHouseDataType;
-import com.clickhouse.data.ClickHouseFormat;
 import com.clickhouse.jdbc.internal.ExceptionUtils;
 import com.clickhouse.jdbc.internal.JdbcConfiguration;
 import com.clickhouse.jdbc.internal.JdbcUtils;
@@ -20,7 +16,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.Closeable;
 import java.sql.Array;
 import java.sql.Blob;
 import java.sql.CallableStatement;
@@ -49,7 +44,6 @@
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.Executor;
-import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 public class ConnectionImpl implements Connection, JdbcV2Wrapper {
@@ -643,9 +637,9 @@ public void setNetworkTimeout(Executor executor, int milliseconds) throws SQLExc
 
 
     // Should be called by child object to notify about timeout.
-    public void onNetworkTimeout() throws SQLException {
-        if (isClosed() || networkTimeoutExecutor == null) {
-            return; // we closed already so do nothing.
+    public synchronized void onNetworkTimeout() {
+        if (this.closed || networkTimeoutExecutor == null) {
+            return; // we closed already or have not set network timeout so do nothing.
         }
 
         networkTimeoutExecutor.execute(() -> {

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

@@ -183,11 +183,7 @@ protected ResultSetImpl executeQueryImpl(String sql, QuerySettings settings) thr
 
     protected void handleSocketTimeoutException(Exception e) {
         if (e.getCause() instanceof SocketTimeoutException || e instanceof SocketTimeoutException) {
-            try {
-                this.connection.onNetworkTimeout();
-            } catch (SQLException e1) {
-                LOG.warn("Failed to handle network timeout exception", e1);
-            }
+            this.connection.onNetworkTimeout();
         }
     }
 

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

@@ -30,6 +30,7 @@
 import java.util.UUID;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 
 import static org.testng.Assert.assertEquals;
@@ -425,10 +426,32 @@ public void testNetworkTimeout() throws Exception {
                 Assert.assertTrue(e.getMessage().contains("closed"));
             }
         }
-        try {
 
-        } catch (Exception e) {
+        Properties connConfig = new Properties();
+        connConfig.setProperty(ClientConfigProperties.SOCKET_OPERATION_TIMEOUT.getKey(), "10");
+        try (Connection conn = getJdbcConnection(connConfig)) {
+            Statement stmt = conn.createStatement();
+            try (ResultSet rs = stmt.executeQuery("SELECT sleepEachRow(1) FROM system.numbers LIMIT 2")) {
+                fail("Exception expected");
+            } catch (Exception e) {
+                Assert.assertFalse(conn.isClosed());
+                Assert.assertTrue(conn.isValid(1000));
+            }
+        }
 
+        try (Connection conn = getJdbcConnection(connConfig)) {
+            ExecutorService executorService = Executors.newSingleThreadExecutor();
+            conn.setNetworkTimeout(executorService, 10);
+            try (Statement stmt1 = conn.createStatement(); Statement stmt2 = conn.createStatement()) {
+                ScheduledExecutorService stmtExecutor = Executors.newScheduledThreadPool(2);
+                long t1 = System.currentTimeMillis();
+                stmtExecutor.schedule(() -> stmt1.executeQuery("SELECT sleepEachRow(1) FROM system.numbers LIMIT 2"), 100, TimeUnit.MILLISECONDS);
+                long t2 = System.currentTimeMillis() - t1;
+                stmtExecutor.schedule(() -> stmt2.executeQuery("SELECT sleepEachRow(1) FROM system.numbers LIMIT 1"), 100 - t2, TimeUnit.MILLISECONDS);
+
+                stmtExecutor.shutdown();
+                stmtExecutor.awaitTermination(10, TimeUnit.SECONDS);
+            }
         }
     }