From 9a1ead8ca4cd9375fe8a6cf404260b8a3008c412 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 19 Aug 2026 09:41:41 +0000 Subject: [PATCH] fix(jdbc): preserve TDengine database context Closes #6682 --- .../main/java/app/dbx/jdbc/DbxJdbcPlugin.java | 42 +++- .../java/app/dbx/jdbc/DbxJdbcPluginTest.java | 184 +++++++++++++++++- 2 files changed, 213 insertions(+), 13 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 dbf00e6388..4f90afcc4f 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -30,6 +30,7 @@ import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; +import java.sql.SQLClientInfoException; import java.sql.SQLFeatureNotSupportedException; import java.lang.reflect.Method; import java.sql.Statement; @@ -74,12 +75,15 @@ public final class DbxJdbcPlugin { false, false, false, + null, StatementMaxRowsMode.READ_LOOP_ONLY ); private static final JdbcDriverQuirks USE_CATALOG_QUIRKS = DEFAULT_QUIRKS.withUseCatalogFallbackSql(true); private static final JdbcDriverQuirks HIVE_QUIRKS = USE_CATALOG_QUIRKS.withSchemasAsDatabasesFallback(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 TAOS_QUIRKS = DEFAULT_QUIRKS + .withPreferExecuteQueryForResultSetSql(true) + .withDatabaseClientInfoProperty("dbname"); private static final JdbcDriverQuirks YASHAN_QUIRKS = new JdbcDriverQuirks( true, true, @@ -88,6 +92,7 @@ public final class DbxJdbcPlugin { false, false, false, + null, StatementMaxRowsMode.APPLY_STATEMENT_MAX_ROWS ); private static final JdbcDriverQuirks IRIS_QUIRKS = new JdbcDriverQuirks( @@ -98,6 +103,7 @@ public final class DbxJdbcPlugin { false, false, false, + null, StatementMaxRowsMode.READ_LOOP_ONLY ); private static final JdbcDriverQuirks ORACLE_QUIRKS = new JdbcDriverQuirks( @@ -108,6 +114,7 @@ public final class DbxJdbcPlugin { false, false, false, + null, StatementMaxRowsMode.APPLY_STATEMENT_MAX_ROWS ); private static final List DRIVER_QUIRK_RULES = List.of( @@ -122,7 +129,8 @@ public final class DbxJdbcPlugin { new JdbcDriverQuirkRule("jdbc:oracle:", ORACLE_QUIRKS), new JdbcDriverQuirkRule("jdbc:dm:", ORACLE_QUIRKS), new JdbcDriverQuirkRule("jdbc:taos:", TAOS_QUIRKS), - new JdbcDriverQuirkRule("jdbc:taos-ws:", TAOS_QUIRKS) + new JdbcDriverQuirkRule("jdbc:taos-ws:", TAOS_QUIRKS), + new JdbcDriverQuirkRule("jdbc:taos-rs:", TAOS_QUIRKS) ); private static String registeredDriverKey = ""; private static String sharedConnectionKey = ""; @@ -137,6 +145,7 @@ public final class DbxJdbcPlugin { boolean ignoreCatalogForSchemaMetadata, boolean preferExecuteQueryForResultSetSql, boolean schemasAsDatabasesFallback, + String databaseClientInfoProperty, StatementMaxRowsMode statementMaxRowsMode ) { JdbcDriverQuirks withUseCatalogFallbackSql(boolean value) { @@ -148,6 +157,7 @@ public final class DbxJdbcPlugin { ignoreCatalogForSchemaMetadata, preferExecuteQueryForResultSetSql, schemasAsDatabasesFallback, + databaseClientInfoProperty, statementMaxRowsMode ); } @@ -161,6 +171,7 @@ public final class DbxJdbcPlugin { value, preferExecuteQueryForResultSetSql, schemasAsDatabasesFallback, + databaseClientInfoProperty, statementMaxRowsMode ); } @@ -174,6 +185,7 @@ public final class DbxJdbcPlugin { ignoreCatalogForSchemaMetadata, value, schemasAsDatabasesFallback, + databaseClientInfoProperty, statementMaxRowsMode ); } @@ -187,6 +199,21 @@ public final class DbxJdbcPlugin { ignoreCatalogForSchemaMetadata, preferExecuteQueryForResultSetSql, value, + databaseClientInfoProperty, + statementMaxRowsMode + ); + } + + JdbcDriverQuirks withDatabaseClientInfoProperty(String value) { + return new JdbcDriverQuirks( + skipExecutionContext, + useOracleMetadata, + caseInsensitiveSchemaMetadata, + useCatalogFallbackSql, + ignoreCatalogForSchemaMetadata, + preferExecuteQueryForResultSetSql, + schemasAsDatabasesFallback, + value, statementMaxRowsMode ); } @@ -1505,7 +1532,8 @@ public final class DbxJdbcPlugin { } private static void applyExecutionContext(JsonNode connection, Connection conn, String database, String schema) throws SQLException { - if (driverQuirks(connection).skipExecutionContext()) { + JdbcDriverQuirks quirks = driverQuirks(connection); + if (quirks.skipExecutionContext()) { return; } String catalog = emptyToNull(database); @@ -1514,7 +1542,13 @@ public final class DbxJdbcPlugin { conn.setCatalog(catalog); } catch (SQLFeatureNotSupportedException | AbstractMethodError | UnsupportedOperationException ignored) { } - if (driverQuirks(connection).useCatalogFallbackSql()) { + if (quirks.databaseClientInfoProperty() != null) { + try { + conn.setClientInfo(quirks.databaseClientInfoProperty(), catalog); + } catch (SQLClientInfoException | AbstractMethodError | UnsupportedOperationException ignored) { + } + } + if (quirks.useCatalogFallbackSql()) { applyUseCatalogFallback(conn, catalog); } } 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 f2d5bcc087..4ebfb43599 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -21,6 +21,7 @@ import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; +import java.sql.SQLClientInfoException; import java.sql.SQLFeatureNotSupportedException; import java.sql.Statement; import java.sql.Time; @@ -623,12 +624,144 @@ final class DbxJdbcPluginTest { assertFalse(response.has("error"), response.toString()); assertEquals("row-value", response.path("result").path("rows").path(0).path(0).asText()); - assertEquals(List.of("executeQuery"), calls); + assertEquals(List.of("createStatement", "executeQuery"), calls); } finally { DriverManager.deregisterDriver(driver); } } + @Test + void tdengineUrlsApplyDatabaseClientInfoBeforeCreatingStatement() throws Exception { + for (String url : List.of( + "jdbc:taos://dbx-fake:6030", + "jdbc:taos-ws://dbx-fake:6041", + "jdbc:taos-rs://dbx-fake:6041" + )) { + List calls = new ArrayList<>(); + String prefix = url.substring(0, url.indexOf("//")); + Driver driver = new BrokenResultSetDriver(prefix, true, -1, calls); + DriverManager.registerDriver(driver); + String connection = """ + { + "connection_string": "%s", + "connect_timeout_secs": 30 + } + """.formatted(url); + try { + JsonNode response = request("executeQuery", """ + { + "connection": %s, + "database": "bopu_light", + "sql": "SELECT v FROM meters" + } + """.formatted(connection)); + + assertFalse(response.has("error"), response.toString()); + assertEquals( + List.of( + "setCatalog:bopu_light", + "setClientInfo:dbname:bopu_light", + "createStatement", + "executeQuery" + ), + calls, + url + ); + } finally { + closeAndDeregister(connection, driver); + } + } + } + + @Test + void ordinaryJdbcDatabaseContextDoesNotSetTdengineClientInfo() throws Exception { + List calls = new ArrayList<>(); + Driver driver = new BrokenResultSetDriver("jdbc:dbx-context:", true, -1, calls); + DriverManager.registerDriver(driver); + String connection = """ + { + "connection_string": "jdbc:dbx-context:demo", + "connect_timeout_secs": 30 + } + """; + try { + JsonNode response = request("executeQuery", """ + { + "connection": %s, + "database": "app", + "sql": "SELECT v FROM meters" + } + """.formatted(connection)); + + assertFalse(response.has("error"), response.toString()); + assertEquals(List.of("setCatalog:app", "createStatement", "execute", "executeQuery"), calls); + } finally { + closeAndDeregister(connection, driver); + } + } + + @Test + void tdengineBlankDatabaseSkipsDatabaseContext() throws Exception { + List calls = new ArrayList<>(); + Driver driver = new BrokenResultSetDriver("jdbc:taos-rs:", true, -1, calls); + DriverManager.registerDriver(driver); + String connection = """ + { + "connection_string": "jdbc:taos-rs://dbx-fake:6041", + "connect_timeout_secs": 30 + } + """; + try { + JsonNode response = request("executeQuery", """ + { + "connection": %s, + "database": " ", + "sql": "SELECT v FROM meters" + } + """.formatted(connection)); + + assertFalse(response.has("error"), response.toString()); + assertEquals(List.of("createStatement", "executeQuery"), calls); + } finally { + closeAndDeregister(connection, driver); + } + } + + @Test + void tdengineClientInfoCompatibilityFailuresDoNotBlockQueries() throws Exception { + for (Throwable failure : List.of( + new SQLClientInfoException(), + new UnsupportedOperationException("unsupported"), + new AbstractMethodError("unsupported") + )) { + List calls = new ArrayList<>(); + Driver driver = new BrokenResultSetDriver("jdbc:taos-rs:", true, -1, calls, failure); + DriverManager.registerDriver(driver); + String connection = """ + { + "connection_string": "jdbc:taos-rs://dbx-fake:6041", + "connect_timeout_secs": 30 + } + """; + try { + JsonNode response = request("executeQuery", """ + { + "connection": %s, + "database": "bopu_light", + "sql": "SELECT v FROM meters" + } + """.formatted(connection)); + + assertFalse(response.has("error"), failure.getClass().getSimpleName() + ": " + response); + assertEquals("setCatalog:bopu_light", calls.get(0)); + assertEquals("setClientInfo:dbname:bopu_light", calls.get(1)); + assertEquals("executeQuery", calls.get(calls.size() - 1)); + } finally { + closeAndDeregister(connection, driver); + } + } + } + @Test void connectionUsernameWithMultipleAtSignsIsPassedToDriverProperties() throws Exception { RecordingConnectDriver driver = new RecordingConnectDriver("jdbc:dbx-proxysql-form:"); @@ -3418,16 +3551,28 @@ final class DbxJdbcPluginTest { private final boolean executeReturnsResultSet; private final int updateCount; private final List calls; + private final Throwable clientInfoFailure; private BrokenResultSetDriver(String urlPrefix, boolean executeReturnsResultSet, int updateCount) { - this(urlPrefix, executeReturnsResultSet, updateCount, new ArrayList<>()); + this(urlPrefix, executeReturnsResultSet, updateCount, new ArrayList<>(), null); } private BrokenResultSetDriver(String urlPrefix, boolean executeReturnsResultSet, int updateCount, List calls) { + this(urlPrefix, executeReturnsResultSet, updateCount, calls, null); + } + + private BrokenResultSetDriver( + String urlPrefix, + boolean executeReturnsResultSet, + int updateCount, + List calls, + Throwable clientInfoFailure + ) { this.urlPrefix = urlPrefix; this.executeReturnsResultSet = executeReturnsResultSet; this.updateCount = updateCount; this.calls = calls; + this.clientInfoFailure = clientInfoFailure; } @Override @@ -3435,7 +3580,7 @@ final class DbxJdbcPluginTest { if (!acceptsURL(url)) { return null; } - return brokenResultSetConnection(executeReturnsResultSet, updateCount, calls); + return brokenResultSetConnection(executeReturnsResultSet, updateCount, calls, clientInfoFailure); } @Override @@ -3469,15 +3614,36 @@ final class DbxJdbcPluginTest { } } - private static Connection brokenResultSetConnection(boolean executeReturnsResultSet, int updateCount, List calls) { + private static Connection brokenResultSetConnection( + boolean executeReturnsResultSet, + int updateCount, + List calls, + Throwable clientInfoFailure + ) { 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()); + (proxy, method, args) -> { + if ("setClientInfo".equals(method.getName())) { + calls.add("setClientInfo:" + args[0] + ":" + args[1]); + if (clientInfoFailure != null) { + throw clientInfoFailure; + } + return null; + } + return switch (method.getName()) { + case "createStatement" -> { + calls.add("createStatement"); + yield brokenResultSetStatement(executeReturnsResultSet, updateCount, calls); + } + case "setCatalog" -> { + calls.add("setCatalog:" + args[0]); + yield null; + } + case "isClosed" -> false; + case "close" -> null; + default -> defaultValue(method.getReturnType()); + }; } ); }