From cef9ce8994aaa0d7f1f297b6fbd84dca6c13c948 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 16 Jun 2026 16:09:55 +0800 Subject: [PATCH] fix(jdbc): handle taos query result sets --- .../main/java/app/dbx/jdbc/DbxJdbcPlugin.java | 101 ++++++++- .../java/app/dbx/jdbc/DbxJdbcPluginTest.java | 213 ++++++++++++++++++ 2 files changed, 309 insertions(+), 5 deletions(-) diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index 08ffbff14..7d810cde7 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -59,16 +59,19 @@ public final class DbxJdbcPlugin { false, false, false, + false, StatementMaxRowsMode.APPLY_STATEMENT_MAX_ROWS ); private static final JdbcDriverQuirks USE_CATALOG_QUIRKS = DEFAULT_QUIRKS.withUseCatalogFallbackSql(true); private static final JdbcDriverQuirks KINGBASE_QUIRKS = DEFAULT_QUIRKS.withIgnoreCatalogForSchemaMetadata(true); + private static final JdbcDriverQuirks TAOS_QUIRKS = DEFAULT_QUIRKS.withPreferExecuteQueryForResultSetSql(true); private static final JdbcDriverQuirks YASHAN_QUIRKS = new JdbcDriverQuirks( true, true, false, false, false, + false, StatementMaxRowsMode.APPLY_STATEMENT_MAX_ROWS ); private static final JdbcDriverQuirks IRIS_QUIRKS = new JdbcDriverQuirks( @@ -77,6 +80,7 @@ public final class DbxJdbcPlugin { true, false, false, + false, StatementMaxRowsMode.READ_LOOP_ONLY ); private static final JdbcDriverQuirks ORACLE_QUIRKS = new JdbcDriverQuirks( @@ -85,6 +89,7 @@ public final class DbxJdbcPlugin { false, false, false, + false, StatementMaxRowsMode.APPLY_STATEMENT_MAX_ROWS ); private static final List DRIVER_QUIRK_RULES = List.of( @@ -97,7 +102,9 @@ public final class DbxJdbcPlugin { new JdbcDriverQuirkRule("jdbc:yasdb:", YASHAN_QUIRKS), new JdbcDriverQuirkRule("jdbc:iris:", IRIS_QUIRKS), new JdbcDriverQuirkRule("jdbc:oracle:", ORACLE_QUIRKS), - new JdbcDriverQuirkRule("jdbc:dm:", ORACLE_QUIRKS) + new JdbcDriverQuirkRule("jdbc:dm:", ORACLE_QUIRKS), + new JdbcDriverQuirkRule("jdbc:taos:", TAOS_QUIRKS), + new JdbcDriverQuirkRule("jdbc:taos-ws:", TAOS_QUIRKS) ); private static String registeredDriverKey = ""; private static String sharedConnectionKey = ""; @@ -109,6 +116,7 @@ public final class DbxJdbcPlugin { boolean caseInsensitiveSchemaMetadata, boolean useCatalogFallbackSql, boolean ignoreCatalogForSchemaMetadata, + boolean preferExecuteQueryForResultSetSql, StatementMaxRowsMode statementMaxRowsMode ) { JdbcDriverQuirks withUseCatalogFallbackSql(boolean value) { @@ -118,6 +126,7 @@ public final class DbxJdbcPlugin { caseInsensitiveSchemaMetadata, value, ignoreCatalogForSchemaMetadata, + preferExecuteQueryForResultSetSql, statementMaxRowsMode ); } @@ -129,6 +138,19 @@ public final class DbxJdbcPlugin { caseInsensitiveSchemaMetadata, useCatalogFallbackSql, value, + preferExecuteQueryForResultSetSql, + statementMaxRowsMode + ); + } + + JdbcDriverQuirks withPreferExecuteQueryForResultSetSql(boolean value) { + return new JdbcDriverQuirks( + skipExecutionContext, + useOracleMetadata, + caseInsensitiveSchemaMetadata, + useCatalogFallbackSql, + ignoreCatalogForSchemaMetadata, + value, statementMaxRowsMode ); } @@ -348,14 +370,15 @@ public final class DbxJdbcPlugin { JdbcDriverQuirks quirks = driverQuirks(connection); try (Statement statement = conn.createStatement()) { applyStatementOptions(statement, maxRows, fetchSize, timeoutSecs, quirks); - boolean hasResultSet = statement.execute(trimStatementSql(sql)); + String trimmedSql = trimStatementSql(sql); + ExecutedStatement executed = executeStatementForResult(statement, trimmedSql, quirks); ObjectNode result = MAPPER.createObjectNode(); ArrayNode columns = MAPPER.createArrayNode(); ArrayNode rows = MAPPER.createArrayNode(); boolean truncated = false; - if (hasResultSet) { - try (ResultSet rs = statement.getResultSet()) { + try (ResultSet rs = executed.resultSet()) { + if (rs != null) { ResultSetMetaData meta = rs.getMetaData(); int columnCount = meta.getColumnCount(); for (int i = 1; i <= columnCount; i++) { @@ -378,13 +401,81 @@ public final class DbxJdbcPlugin { result.set("columns", columns); result.set("rows", rows); - result.put("affected_rows", hasResultSet ? 0 : Math.max(statement.getUpdateCount(), 0)); + result.put("affected_rows", columns.isEmpty() ? Math.max(executed.updateCount(), 0) : 0); result.put("execution_time_ms", (System.nanoTime() - start) / 1_000_000); result.put("truncated", truncated); return result; } } + private record ExecutedStatement(ResultSet resultSet, int updateCount) { + } + + private static ExecutedStatement executeStatementForResult( + Statement statement, + String sql, + JdbcDriverQuirks quirks + ) throws SQLException { + if (quirks.preferExecuteQueryForResultSetSql() && looksLikeResultSetSql(sql)) { + return new ExecutedStatement(statement.executeQuery(sql), -1); + } + boolean hasResultSet = statement.execute(sql); + int updateCount = hasResultSet ? -1 : statement.getUpdateCount(); + ResultSet rs = hasResultSet ? statement.getResultSet() : null; + if (rs == null && shouldRetryWithExecuteQuery(sql, hasResultSet, updateCount)) { + rs = statement.executeQuery(sql); + } + return new ExecutedStatement(rs, updateCount); + } + + private static boolean shouldRetryWithExecuteQuery(String sql, boolean hasResultSet, int updateCount) { + if (hasResultSet) { + return true; + } + return updateCount < 0 && looksLikeResultSetSql(sql); + } + + static boolean looksLikeResultSetSql(String sql) { + String keyword = firstSqlKeyword(sql); + return switch (keyword) { + case "SELECT", "WITH", "SHOW", "DESCRIBE", "DESC", "EXPLAIN", "VALUES", "TABLE", "PRAGMA" -> true; + default -> false; + }; + } + + private static String firstSqlKeyword(String sql) { + String text = stripLeadingSqlComments(sql).trim(); + int end = 0; + while (end < text.length() && Character.isLetter(text.charAt(end))) { + end++; + } + return text.substring(0, end).toUpperCase(Locale.ROOT); + } + + private static String stripLeadingSqlComments(String sql) { + String text = sql.trim(); + boolean changed; + do { + changed = false; + if (text.startsWith("--")) { + int lineEnd = text.indexOf('\n'); + if (lineEnd < 0) { + return ""; + } + text = text.substring(lineEnd + 1).trim(); + changed = true; + } else if (text.startsWith("/*")) { + int commentEnd = text.indexOf("*/", 2); + if (commentEnd < 0) { + return ""; + } + text = text.substring(commentEnd + 2).trim(); + changed = true; + } + } while (changed); + return text; + } + /** * Get DM execution plan using DmdbConnection.getExplainInfo() via reflection. * diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java index e26163d7d..ea31195c7 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -10,9 +10,12 @@ import java.lang.reflect.Proxy; import java.sql.Connection; import java.sql.DatabaseMetaData; import java.sql.Date; +import java.sql.Driver; import java.sql.DriverManager; +import java.sql.DriverPropertyInfo; import java.sql.ResultSet; import java.sql.ResultSetMetaData; +import java.sql.SQLException; import java.sql.Statement; import java.sql.Timestamp; import java.sql.Types; @@ -136,6 +139,75 @@ final class DbxJdbcPluginTest { assertEquals(true, response.path("result").path("truncated").asBoolean()); } + @Test + void executeQueryFallsBackWhenExecutedStatementReturnsNullResultSet() throws Exception { + Driver driver = new BrokenResultSetDriver("jdbc:dbx-null-execute-rs:", true, -1); + DriverManager.registerDriver(driver); + try { + JsonNode response = request("executeQuery", """ + { + "connection": { + "connection_string": "jdbc:dbx-null-execute-rs:demo", + "connect_timeout_secs": 30 + }, + "sql": "SELECT v FROM meters" + } + """); + + assertFalse(response.has("error"), response.toString()); + assertEquals("VALUE", response.path("result").path("columns").path(0).asText()); + assertEquals("row-value", response.path("result").path("rows").path(0).path(0).asText()); + } finally { + DriverManager.deregisterDriver(driver); + } + } + + @Test + void executeQueryFallsBackForQuerySqlWithoutUpdateCount() throws Exception { + Driver driver = new BrokenResultSetDriver("jdbc:dbx-no-result-flag:", false, -1); + DriverManager.registerDriver(driver); + try { + JsonNode response = request("executeQuery", """ + { + "connection": { + "connection_string": "jdbc:dbx-no-result-flag:demo", + "connect_timeout_secs": 30 + }, + "sql": "-- generated preview\\nSHOW TABLES" + } + """); + + assertFalse(response.has("error"), response.toString()); + assertEquals("row-value", response.path("result").path("rows").path(0).path(0).asText()); + } finally { + DriverManager.deregisterDriver(driver); + } + } + + @Test + void taosQuerySqlUsesExecuteQueryDirectly() throws Exception { + List calls = new ArrayList<>(); + Driver driver = new BrokenResultSetDriver("jdbc:taos:", true, -1, calls); + DriverManager.registerDriver(driver); + try { + JsonNode response = request("executeQuery", """ + { + "connection": { + "connection_string": "jdbc:taos://dbx-fake:6030/power", + "connect_timeout_secs": 30 + }, + "sql": "SELECT v FROM meters" + } + """); + + assertFalse(response.has("error"), response.toString()); + assertEquals("row-value", response.path("result").path("rows").path(0).path(0).asText()); + assertEquals(List.of("executeQuery"), calls); + } finally { + DriverManager.deregisterDriver(driver); + } + } + @Test void connectTimeoutIsMappedToDriverProperties() throws Exception { Method method = DbxJdbcPlugin.class.getDeclaredMethod("applyConnectTimeout", JsonNode.class, Properties.class); @@ -269,6 +341,11 @@ final class DbxJdbcPluginTest { "jdbc_driver_class": "org.apache.kyuubi.jdbc.KyuubiHiveDriver" } """); + JsonNode taos = MAPPER.readTree(""" + { + "connection_string": "jdbc:TAOS://127.0.0.1:6030/power" + } + """); assertEquals(true, DbxJdbcPlugin.driverQuirks(yashan).skipExecutionContext()); assertEquals(true, DbxJdbcPlugin.driverQuirks(yashan).useOracleMetadata()); @@ -290,6 +367,7 @@ final class DbxJdbcPluginTest { assertEquals(true, DbxJdbcPlugin.driverQuirks(mysql).useCatalogFallbackSql()); assertEquals(true, DbxJdbcPlugin.driverQuirks(kingbase).ignoreCatalogForSchemaMetadata()); assertEquals(true, DbxJdbcPlugin.driverQuirks(kyuubi).useCatalogFallbackSql()); + assertEquals(true, DbxJdbcPlugin.driverQuirks(taos).preferExecuteQueryForResultSetSql()); } @Test @@ -666,6 +744,141 @@ final class DbxJdbcPluginTest { ); } + private static final class BrokenResultSetDriver implements Driver { + private final String urlPrefix; + private final boolean executeReturnsResultSet; + private final int updateCount; + private final List calls; + + private BrokenResultSetDriver(String urlPrefix, boolean executeReturnsResultSet, int updateCount) { + this(urlPrefix, executeReturnsResultSet, updateCount, new ArrayList<>()); + } + + private BrokenResultSetDriver(String urlPrefix, boolean executeReturnsResultSet, int updateCount, List calls) { + this.urlPrefix = urlPrefix; + this.executeReturnsResultSet = executeReturnsResultSet; + this.updateCount = updateCount; + this.calls = calls; + } + + @Override + public Connection connect(String url, Properties info) throws SQLException { + if (!acceptsURL(url)) { + return null; + } + return brokenResultSetConnection(executeReturnsResultSet, updateCount, calls); + } + + @Override + public boolean acceptsURL(String url) { + return url != null && url.startsWith(urlPrefix); + } + + @Override + public DriverPropertyInfo[] getPropertyInfo(String url, Properties info) { + return new DriverPropertyInfo[0]; + } + + @Override + public int getMajorVersion() { + return 1; + } + + @Override + public int getMinorVersion() { + return 0; + } + + @Override + public boolean jdbcCompliant() { + return false; + } + + @Override + public java.util.logging.Logger getParentLogger() { + return java.util.logging.Logger.getGlobal(); + } + } + + private static Connection brokenResultSetConnection(boolean executeReturnsResultSet, int updateCount, List calls) { + return (Connection) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Connection.class }, + (proxy, method, args) -> switch (method.getName()) { + case "createStatement" -> brokenResultSetStatement(executeReturnsResultSet, updateCount, calls); + case "isClosed" -> false; + case "close" -> null; + default -> defaultValue(method.getReturnType()); + } + ); + } + + private static Statement brokenResultSetStatement(boolean executeReturnsResultSet, int updateCount, List calls) { + return (Statement) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Statement.class }, + (proxy, method, args) -> { + if ("execute".equals(method.getName()) || "executeQuery".equals(method.getName())) { + calls.add(method.getName()); + } + return switch (method.getName()) { + case "execute" -> executeReturnsResultSet; + case "getResultSet" -> null; + case "getUpdateCount" -> updateCount; + case "executeQuery" -> singleRowResultSet(); + case "setMaxRows", "setFetchSize", "setQueryTimeout", "close" -> null; + default -> defaultValue(method.getReturnType()); + }; + } + ); + } + + private static ResultSet singleRowResultSet() { + return (ResultSet) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { ResultSet.class }, + new java.lang.reflect.InvocationHandler() { + private int index = -1; + + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return switch (method.getName()) { + case "next" -> ++index == 0; + case "getMetaData" -> singleColumnMeta(); + case "getObject", "getString" -> "row-value"; + case "close" -> null; + default -> defaultValue(method.getReturnType()); + }; + } + } + ); + } + + private static ResultSetMetaData singleColumnMeta() { + return (ResultSetMetaData) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { ResultSetMetaData.class }, + (proxy, method, args) -> switch (method.getName()) { + case "getColumnCount" -> 1; + case "getColumnLabel", "getColumnName" -> "VALUE"; + case "getColumnType" -> Types.VARCHAR; + default -> defaultValue(method.getReturnType()); + } + ); + } + + private static Object defaultValue(Class returnType) { + if (returnType == boolean.class) return false; + if (returnType == byte.class) return (byte) 0; + if (returnType == short.class) return (short) 0; + if (returnType == int.class) return 0; + if (returnType == long.class) return 0L; + if (returnType == float.class) return 0f; + if (returnType == double.class) return 0d; + if (returnType == char.class) return '\0'; + return null; + } + private static JsonNode request(String method, String params) throws Exception { Method handleLine = DbxJdbcPlugin.class.getDeclaredMethod("handleLine", String.class); handleLine.setAccessible(true);