Merge pull request #1934 from ClickHouse/jdbc_fix_result_overflow_mode

Fixed overriding max_result_rows settings

Author: chernser

clickhouse-client/src/main/java/com/clickhouse/client/config/ClickHouseClientOption.java

@@ -218,6 +218,8 @@ public enum ClickHouseClientOption implements ClickHouseOption {
     MAX_RESULT_ROWS("max_result_rows", 0L,
             "Limit on the number of rows in the result. "
                     + "Also checked for subqueries, and on remote servers when running parts of a distributed query."),
+
+    RESULT_OVERFLOW_MODE("result_overflow_mode", "throw","What to do if the result is overflowed."),
     /**
      * Maximum size of thread pool for each client.
      */

clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpConnection.java

@@ -118,10 +118,25 @@ static String buildQueryParams(ClickHouseRequest<?> request, Map<String, Seriali
             appendQueryParameter(builder, settingKey, String.valueOf(config.getMaxExecutionTime()));
         }
         settingKey = ClickHouseClientOption.MAX_RESULT_ROWS.getKey();
-        if (config.getMaxResultRows() > 0L && !settings.containsKey(settingKey)) {
+        boolean hasRequestSetting = settings.containsKey(settingKey);
+        if (config.getMaxResultRows() > 0L && !hasRequestSetting) {
+            // set on client level
             appendQueryParameter(builder, settingKey, String.valueOf(config.getMaxResultRows()));
-            appendQueryParameter(builder, "result_overflow_mode", "break");
+        } else if (hasRequestSetting) {
+            // set on request level
+            Object value = settings.get(settingKey);
+            if (value instanceof Number && ((Number) value).longValue() > 0L) {
+                appendQueryParameter(builder, settingKey, String.valueOf(value));
+            } else if (value instanceof String && !(((String) value).isEmpty() || "0".equals(value))) {
+                appendQueryParameter(builder, settingKey, (String) value);
+            }
+        }
+
+        if (config.hasOption(ClickHouseClientOption.RESULT_OVERFLOW_MODE)) {
+            appendQueryParameter(builder, ClickHouseClientOption.RESULT_OVERFLOW_MODE.getKey(),
+                    config.getStrOption(ClickHouseClientOption.RESULT_OVERFLOW_MODE));
         }
+
         settingKey = "log_comment";
         if (!stmts.isEmpty() && config.getBoolOption(ClickHouseClientOption.LOG_LEADING_COMMENT)
                 && !settings.containsKey(settingKey)) {
@@ -162,11 +177,9 @@ static String buildQueryParams(ClickHouseRequest<?> request, Map<String, Seriali
         }
 
         for (Entry<String, Serializable> entry : settings.entrySet()) {
-            // Skip internal settings
-            if (entry.getKey().equalsIgnoreCase("_set_roles_stmt")) {
-                continue;
+            if (!processedSettings.contains(entry.getKey())) {
+                appendQueryParameter(builder, entry.getKey(), String.valueOf(entry.getValue()));
             }
-            appendQueryParameter(builder, entry.getKey(), String.valueOf(entry.getValue()));
         }
 
         if (builder.length() > 0) {
@@ -175,6 +188,12 @@ static String buildQueryParams(ClickHouseRequest<?> request, Map<String, Seriali
         return builder.toString();
     }
 
+    // Settings that are processed by the client and appropriate parameters are set
+    private static final HashSet<String> processedSettings = new HashSet<>(Arrays.asList(
+            "_set_roles_stmt",
+            ClickHouseClientOption.MAX_RESULT_ROWS.getKey()
+    ));
+
     static String buildUrl(String baseUrl, ClickHouseRequest<?> request, Map<String, Serializable> additionalParams) {
         StringBuilder builder = new StringBuilder().append(baseUrl);
         // TODO: Using default until we will remove

clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/internal/ClickHouseStatementImpl.java

@@ -610,11 +610,9 @@ public void setLargeMaxRows(long max) throws SQLException {
 
         if (this.maxRows != max) {
             if (max == 0L || !connection.allowCustomSetting()) {
-                request.removeSetting(ClickHouseClientOption.MAX_RESULT_ROWS.getKey());
-                request.removeSetting("result_overflow_mode");
+                request.set(ClickHouseClientOption.MAX_RESULT_ROWS.getKey(), 0);
             } else {
                 request.set(ClickHouseClientOption.MAX_RESULT_ROWS.getKey(), max);
-                request.set("result_overflow_mode", "break");
             }
             this.maxRows = max;
         }

clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHouseStatementTest.java

@@ -834,7 +834,9 @@ public void testWrapperObject() throws SQLException {
 
     @Test(groups = "integration")
     public void testQuerySystemLog() throws SQLException {
-        try (ClickHouseConnection conn = newConnection(new Properties())) {
+        Properties props = new Properties();
+        props.setProperty(ClickHouseClientOption.RESULT_OVERFLOW_MODE.getKey(), "break");
+        try (ClickHouseConnection conn = newConnection(props)) {
             ClickHouseStatement stmt = conn.createStatement();
             stmt.setMaxRows(10);
             stmt.setLargeMaxRows(11L);
@@ -1493,4 +1495,30 @@ public void testDescMetadata() {
             Assert.fail("Failed to create connection", e);
         }
     }
+
+    @Test(groups = "integration")
+    public void testMaxResultsRows() throws SQLException {
+        Properties props = new Properties();
+        int maxRows = 3;
+        props.setProperty(ClickHouseClientOption.MAX_RESULT_ROWS.getKey(), String.valueOf(maxRows));
+        props.setProperty(ClickHouseClientOption.RESULT_OVERFLOW_MODE.getKey(), "break");
+        try (ClickHouseConnection conn = newConnection(props);
+             ClickHouseStatement s = conn.createStatement()) {
+            ResultSet rs = s.executeQuery("SELECT number FROM system.numbers");
+            for (int i = 0; i < maxRows; i++) {
+                Assert.assertTrue(rs.next(), "Should have more rows, but have only " + i);
+            }
+        }
+
+        props.setProperty(ClickHouseClientOption.MAX_RESULT_ROWS.getKey(), "1");
+        props.remove(ClickHouseClientOption.RESULT_OVERFLOW_MODE.getKey());
+        try (ClickHouseConnection conn = newConnection(props);
+             ClickHouseStatement s = conn.createStatement()) {
+            s.executeQuery("SELECT number FROM system.numbers");
+            Assert.fail("Should throw exception");
+        } catch (SQLException e) {
+            Assert.assertTrue(e.getMessage().startsWith("Code: 396. DB::Exception: Limit for result exceeded, max rows"),
+                    "Unexpected exception: " + e.getMessage());
+        }
+    }
 }