fix(tdengine): switch REST databases without USE SQL

This commit is contained in:
t8y2 2026-07-21 15:58:06 +08:00
parent e3d5dbda2a
commit 8f98e28305
2 changed files with 73 additions and 3 deletions

View File

@ -18,10 +18,13 @@ import com.dbx.agent.QueryPageResult;
import com.dbx.agent.QueryResult;
import com.dbx.agent.TableInfo;
import com.dbx.agent.TriggerInfo;
import com.taosdata.jdbc.TSDBDriver;
import com.taosdata.jdbc.rs.RestfulConnection;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.sql.Types;
import java.time.format.DateTimeFormatter;
@ -194,7 +197,7 @@ public final class TDengineAgent extends BaseDatabaseAgent {
return JdbcExecutor.current().execute(
requireConnected(),
sql,
schema,
prepareExecutionSchema(schema),
this::setSchemaSQL,
options.getMaxRows(),
options.getFetchSize(),
@ -208,7 +211,7 @@ public final class TDengineAgent extends BaseDatabaseAgent {
return JdbcExecutor.current().executePage(
requireConnected(),
sql,
schema,
prepareExecutionSchema(schema),
this::setSchemaSQL,
options,
this::tdengineResultValue
@ -220,7 +223,7 @@ public final class TDengineAgent extends BaseDatabaseAgent {
return JdbcExecutor.current().startTableRead(
requireConnected(),
sql,
schema,
prepareExecutionSchema(schema),
this::setSchemaSQL,
options,
this::tdengineResultValue
@ -232,6 +235,32 @@ public final class TDengineAgent extends BaseDatabaseAgent {
return "USE " + quoteIdentifier(schema);
}
@Override
public QueryResult executeTransaction(List<String> statements, String schema) {
return super.executeTransaction(statements, prepareExecutionSchema(schema));
}
@Override
public QueryResult executeBatch(List<String> statements, String schema) {
return super.executeBatch(statements, prepareExecutionSchema(schema));
}
private String prepareExecutionSchema(String schema) {
return unchecked(() -> prepareExecutionSchema(requireConnected(), schema));
}
static String prepareExecutionSchema(Connection connection, String schema) throws SQLException {
if (!(connection instanceof RestfulConnection) || schema == null || schema.trim().isEmpty()) {
return schema;
}
String database = schema.trim();
// Connector/J 3.6.3 misparses quoted USE statements and puts the whole SQL in /rest/sql/<db>.
connection.setCatalog(database);
connection.setClientInfo(TSDBDriver.PROPERTY_KEY_DBNAME, database);
return null;
}
@Override
public void disconnect() {
uncheckedVoid(() -> {

View File

@ -10,6 +10,8 @@ import com.dbx.agent.test.JdbcAgentFake;
import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest;
import com.dbx.agent.test.JdbcMetadataSqlFake;
import com.dbx.agent.test.TestSupport;
import com.taosdata.jdbc.TSDBDriver;
import com.taosdata.jdbc.rs.RestfulConnection;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Proxy;
import java.nio.charset.StandardCharsets;
@ -21,6 +23,7 @@ import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Properties;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
@ -281,6 +284,30 @@ class TDengineAgentMetadataTest {
);
}
@Test
void restExecutionSwitchesCatalogWithoutGeneratingUseSql() throws Exception {
try (RestfulConnection connection = restfulConnection()) {
String executionSchema = TDengineAgent.prepareExecutionSchema(connection, " power-data ");
Assertions.assertNull(executionSchema);
Assertions.assertEquals("power-data", connection.getCatalog());
Assertions.assertEquals(
"power-data",
connection.getClientInfo(TSDBDriver.PROPERTY_KEY_DBNAME)
);
}
}
@Test
void websocketExecutionKeepsSchemaForUseSqlSwitching() throws Exception {
Connection connection = JdbcAgentFake.connection();
Assertions.assertEquals(
"power",
TDengineAgent.prepareExecutionSchema(connection, "power")
);
}
@Test
void decodesTdengineByteArrayTextValues() {
Assertions.assertEquals(
@ -344,4 +371,18 @@ class TDengineAgentMetadataTest {
}
);
}
private static RestfulConnection restfulConnection() {
return new RestfulConnection(
"127.0.0.1",
"6041",
new Properties(),
"",
"jdbc:TAOS-RS://127.0.0.1:6041/",
null,
false,
null,
null
);
}
}