xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(jdbc): handle taos query result sets

This commit is contained in:
t8y2 2026-06-16 16:09:55 +08:00
parent d59ce9f483
commit cef9ce8994
2 changed files with 309 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

101
plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java
View File

@ -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<JdbcDriverQuirkRule> 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.
*

213
plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java
View File

@ -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<String> 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<String> calls;
private BrokenResultSetDriver(String urlPrefix, boolean executeReturnsResultSet, int updateCount) {
this(urlPrefix, executeReturnsResultSet, updateCount, new ArrayList<>());
}
private BrokenResultSetDriver(String urlPrefix, boolean executeReturnsResultSet, int updateCount, List<String> 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<String> 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<String> 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);