github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-21 10:28:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(jdbc): preserve TDengine database context

Browse Source 
Closes #6682
This commit is contained in:
t8y2
2026-08-19 09:43:12 +00:00
parent bfbbc2bff2
commit 9a1ead8ca4
2 changed files with 213 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java
+38 -4
View File
@@ -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<JdbcDriverQuirkRule> 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);
}
}
plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java
+175 -9
View File
@@ -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<String> 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<String> 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<String> 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<String> 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<String> 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<String> calls) {
this(urlPrefix, executeReturnsResultSet, updateCount, calls, null);
}
private BrokenResultSetDriver(
String urlPrefix,
boolean executeReturnsResultSet,
int updateCount,
List<String> 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<String> calls) {
private static Connection brokenResultSetConnection(
boolean executeReturnsResultSet,
int updateCount,
List<String> 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());
};
}
);
}