Merge pull request #387 from databendlabs/fix/regex-parse-err

fix: extractColumnTypes parse sql err

Author: hantmac

databend-jdbc/src/main/java/com/databend/jdbc/StatementUtil.java

@@ -68,6 +68,9 @@ public Optional<Pair<String, String>> extractParamFromSetStatement(@NonNull Stri
      */
     public static Map<Integer, String> extractColumnTypes(String sql) {
         Map<Integer, String> columnTypes = new LinkedHashMap<>();
+        if (isQuery(sql)) {
+            return columnTypes;
+        }
         Pattern pattern = Pattern.compile("\\((.*?)\\)");
         Matcher matcher = pattern.matcher(sql);
         if (matcher.find()) {

databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java

@@ -620,6 +620,48 @@ public void testBatchAndNoBatch() throws SQLException {
     }
 
 
+    @Test(groups = {"IT"})
+    public void testSelectCountStarWithPreparedStatement() throws SQLException {
+        if (Compatibility.skipDriverBugLowerThen("0.4.5")) {
+            return;
+        }
+        try (Connection conn = getConn();
+             Statement s = conn.createStatement()) {
+            s.execute("create or replace table t1(a int, b string)");
+            s.execute("insert into t1 values(1, 'a')");
+            s.execute("insert into t1 values(2, 'b')");
+
+            // COUNT(*) should not be parsed as a parameter
+            String countSql = "SELECT COUNT(*) FROM t1";
+            try (PreparedStatement ps = conn.prepareStatement(countSql)) {
+                Assert.assertEquals(ps.getParameterMetaData().getParameterCount(), 0,
+                        "COUNT(*) should have 0 parameters");
+                ResultSet rs = ps.executeQuery();
+                Assert.assertTrue(rs.next());
+                Assert.assertEquals(rs.getInt(1), 2);
+            }
+
+            // SUM with column should not be parsed as a parameter
+            String sumSql = "SELECT SUM(a) FROM t1";
+            try (PreparedStatement ps = conn.prepareStatement(sumSql)) {
+                Assert.assertEquals(ps.getParameterMetaData().getParameterCount(), 0,
+                        "SUM(a) should have 0 parameters");
+                ResultSet rs = ps.executeQuery();
+                Assert.assertTrue(rs.next());
+                Assert.assertEquals(rs.getInt(1), 3);
+            }
+
+            // SELECT with actual parameter marker should still work
+            String selectWithParam = "SELECT COUNT(*) FROM t1 WHERE a = ?";
+            try (PreparedStatement ps = conn.prepareStatement(selectWithParam)) {
+                ps.setInt(1, 1);
+                ResultSet rs = ps.executeQuery();
+                Assert.assertTrue(rs.next());
+                Assert.assertEquals(rs.getInt(1), 1);
+            }
+        }
+    }
+
     @Test(groups = {"IT"})
     public void testSelectWithPreparedStatement()
             throws SQLException {

databend-jdbc/src/test/java/com/databend/jdbc/TestStatementUtil.java

@@ -19,4 +19,25 @@ public void testExtractColumnTypes() {
         assertEquals("Int8", columnTypes.get(1));
         assertEquals("VARIANT", columnTypes.get(2));
     }
+
+    @Test(groups = {"UNIT"})
+    public void testExtractColumnTypesSelectCountStar() {
+        String sql = "SELECT COUNT(*) FROM 'szps_dwd'.'DWD_WSC_WS_GYYX_LJSLLJ_1H'";
+        Map<Integer, String> columnTypes = StatementUtil.extractColumnTypes(sql);
+        assertEquals(0, columnTypes.size());
+    }
+
+    @Test(groups = {"UNIT"})
+    public void testExtractColumnTypesSelectWithFunctions() {
+        String sql = "SELECT SUM(a), AVG(b) FROM test_table WHERE id = ?";
+        Map<Integer, String> columnTypes = StatementUtil.extractColumnTypes(sql);
+        assertEquals(0, columnTypes.size());
+    }
+
+    @Test(groups = {"UNIT"})
+    public void testExtractColumnTypesShowTables() {
+        String sql = "SHOW TABLES";
+        Map<Integer, String> columnTypes = StatementUtil.extractColumnTypes(sql);
+        assertEquals(0, columnTypes.size());
+    }
 }