feat(tdengine): optimize metadata comment search
This commit is contained in:
二丫讲梵
GitHub
parent
2cbbaaef4c
commit
4a5eb3f04a
|
|
@ -10,7 +10,6 @@ import com.dbx.agent.IndexInfo;
|
|||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
import com.dbx.agent.ObjectSource;
|
||||
import com.dbx.agent.QueryPageOptions;
|
||||
|
|
@ -48,6 +47,7 @@ import java.util.regex.Matcher;
|
|||
import java.util.regex.Pattern;
|
||||
|
||||
public final class TDengineAgent extends BaseDatabaseAgent {
|
||||
private static final long TABLE_CACHE_TTL_MILLIS = 10_000L;
|
||||
private static final DateTimeFormatter TDENGINE_TIMESTAMP_FORMAT =
|
||||
new DateTimeFormatterBuilder()
|
||||
.appendPattern("yyyy-MM-dd HH:mm:ss")
|
||||
|
|
@ -63,6 +63,10 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
|||
Pattern.compile("(?i)\\bCOMPOSITE\\s+KEY\\b");
|
||||
|
||||
private Connection connection;
|
||||
private final Object tableCacheLock = new Object();
|
||||
private String tableCacheSchema = "";
|
||||
private long tableCacheTimeMillis;
|
||||
private List<TableInfo> tableCache = Collections.emptyList();
|
||||
|
||||
@Override
|
||||
public Connection getConnection() {
|
||||
|
|
@ -73,6 +77,7 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
|||
public void connect(ConnectParams params) {
|
||||
uncheckedVoid(() -> {
|
||||
connection = TDengineConnectionFactory.open(params);
|
||||
clearTableCache();
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -118,19 +123,21 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
|||
if (!normalized.tableTypeAllowed("TABLE")) {
|
||||
return Collections.emptyList();
|
||||
}
|
||||
try {
|
||||
return queryConstrainedTables(schema, normalized);
|
||||
} catch (RuntimeException ignored) {
|
||||
// TDengine 2.x does not expose information_schema.ins_stables/ins_tables.
|
||||
return normalized.filterTables(listTablesLegacy(schema));
|
||||
}
|
||||
return normalized.filterTables(listTablesFromShow(schema));
|
||||
}
|
||||
|
||||
private List<TableInfo> listTablesLegacy(String schema) {
|
||||
private List<TableInfo> listTablesFromShow(String schema) {
|
||||
List<TableInfo> cached = cachedTables(schema);
|
||||
if (cached != null) {
|
||||
return cached;
|
||||
}
|
||||
return unchecked(() -> {
|
||||
List<TableInfo> result = new ArrayList<>();
|
||||
result.addAll(queryTables("SHOW " + quoteQualifiedPrefix(schema) + "STABLES", "STABLE"));
|
||||
result.addAll(queryTables("SHOW " + quoteQualifiedPrefix(schema) + "TABLES", "TABLE"));
|
||||
// Connector/J 3.6.3 ignores Statement#setMaxRows. Read the SHOW
|
||||
// results once and page locally from a short-lived cache instead of
|
||||
// issuing the same full scan for every sidebar page.
|
||||
result.addAll(queryTables("SHOW " + quoteQualifiedPrefix(schema) + "STABLES", "STABLE", false));
|
||||
result.addAll(queryTables("SHOW " + quoteQualifiedPrefix(schema) + "TABLES", "TABLE", true));
|
||||
|
||||
Map<String, TableInfo> distinct = new LinkedHashMap<>();
|
||||
for (TableInfo table : result) {
|
||||
|
|
@ -138,7 +145,8 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
|||
}
|
||||
List<TableInfo> sorted = new ArrayList<>(distinct.values());
|
||||
sortTablesForHierarchy(sorted);
|
||||
return sorted;
|
||||
cacheTables(schema, sorted);
|
||||
return copyTables(sorted);
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -199,7 +207,7 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
|||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
return JdbcExecutor.current().execute(
|
||||
QueryResult result = JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
sql,
|
||||
prepareExecutionSchema(schema),
|
||||
|
|
@ -209,6 +217,10 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
|||
options.getTimeoutSecs(),
|
||||
this::tdengineResultValue
|
||||
);
|
||||
if (mayChangeMetadata(sql)) {
|
||||
clearTableCache();
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -242,12 +254,20 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
|||
|
||||
@Override
|
||||
public QueryResult executeTransaction(List<String> statements, String schema) {
|
||||
return super.executeTransaction(statements, prepareExecutionSchema(schema));
|
||||
QueryResult result = super.executeTransaction(statements, prepareExecutionSchema(schema));
|
||||
if (statements.stream().anyMatch(TDengineAgent::mayChangeMetadata)) {
|
||||
clearTableCache();
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
@Override
|
||||
public QueryResult executeBatch(List<String> statements, String schema) {
|
||||
return super.executeBatch(statements, prepareExecutionSchema(schema));
|
||||
QueryResult result = super.executeBatch(statements, prepareExecutionSchema(schema));
|
||||
if (statements.stream().anyMatch(TDengineAgent::mayChangeMetadata)) {
|
||||
clearTableCache();
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private String prepareExecutionSchema(String schema) {
|
||||
|
|
@ -273,98 +293,84 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
|||
connection.close();
|
||||
}
|
||||
connection = null;
|
||||
clearTableCache();
|
||||
});
|
||||
}
|
||||
|
||||
private List<TableInfo> queryTables(String sql, String tableType) throws Exception {
|
||||
private List<TableInfo> queryTables(String sql, String tableType, boolean includesStableName) throws Exception {
|
||||
List<TableInfo> result = new ArrayList<>();
|
||||
try (java.sql.Statement stmt = requireConnected().createStatement();
|
||||
ResultSet rs = stmt.executeQuery(sql)) {
|
||||
while (rs.next()) {
|
||||
result.add(new TableInfo(rs.getString(1), tableType, null));
|
||||
try (java.sql.Statement stmt = requireConnected().createStatement()) {
|
||||
try (ResultSet rs = stmt.executeQuery(sql)) {
|
||||
while (rs.next()) {
|
||||
// SHOW TABLES returns the owning STABLE as its fourth column. It
|
||||
// is absent for ordinary tables and older servers, where this
|
||||
// best-effort read simply leaves the table at the root level.
|
||||
String parentName = includesStableName ? optionalString(rs, 4) : null;
|
||||
if (parentName != null && parentName.trim().isEmpty()) {
|
||||
parentName = null;
|
||||
}
|
||||
result.add(new TableInfo(rs.getString(1), tableType, null, null, parentName));
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private List<TableInfo> queryConstrainedTables(String database, MetadataListConstraints constraints) {
|
||||
return unchecked(() -> {
|
||||
TableMetadataQuery query = buildTableMetadataQuery(database, constraints);
|
||||
List<TableInfo> result = new ArrayList<>();
|
||||
try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(query.sql())) {
|
||||
MetadataSqlSupport.bind(stmt, query.args());
|
||||
try (ResultSet rs = stmt.executeQuery()) {
|
||||
while (rs.next()) {
|
||||
result.add(new TableInfo(
|
||||
rs.getString(1),
|
||||
rs.getString(2),
|
||||
optionalString(rs, 3),
|
||||
null,
|
||||
optionalString(rs, 4)
|
||||
));
|
||||
}
|
||||
}
|
||||
private List<TableInfo> cachedTables(String schema) {
|
||||
String normalizedSchema = normalizedSchema(schema);
|
||||
synchronized (tableCacheLock) {
|
||||
if (cacheFresh(tableCacheTimeMillis) && tableCacheSchema.equals(normalizedSchema)) {
|
||||
return copyTables(tableCache);
|
||||
}
|
||||
MetadataListConstraints postFilter = constraints.hasLimit() ? constraints.withoutPaging() : constraints;
|
||||
return postFilter.filterTables(result);
|
||||
});
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
static TableMetadataQuery buildTableMetadataQuery(String database, MetadataListConstraints constraints) {
|
||||
MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints);
|
||||
List<Object> args = new ArrayList<>();
|
||||
StringBuilder sql = new StringBuilder("""
|
||||
SELECT table_name, table_type, table_comment, parent_name
|
||||
FROM (
|
||||
SELECT stable_name AS table_name,
|
||||
'STABLE' AS table_type,
|
||||
table_comment,
|
||||
CAST(NULL AS VARCHAR(192)) AS parent_name,
|
||||
stable_name AS hierarchy_name,
|
||||
0 AS hierarchy_rank
|
||||
FROM information_schema.ins_stables
|
||||
WHERE db_name = ?
|
||||
UNION ALL
|
||||
SELECT table_name,
|
||||
'TABLE' AS table_type,
|
||||
table_comment,
|
||||
stable_name AS parent_name,
|
||||
CASE WHEN stable_name IS NULL THEN table_name ELSE stable_name END AS hierarchy_name,
|
||||
1 AS hierarchy_rank
|
||||
FROM information_schema.ins_tables
|
||||
WHERE db_name = ?
|
||||
) metadata
|
||||
WHERE 1 = 1
|
||||
""".stripIndent().trim());
|
||||
args.add(database);
|
||||
args.add(database);
|
||||
if (normalized.hasFilter()) {
|
||||
String pattern = normalized.fuzzyLikePattern();
|
||||
sql.append(" AND (table_name LIKE ? OR table_comment LIKE ?)");
|
||||
args.add(pattern);
|
||||
args.add(pattern);
|
||||
private void cacheTables(String schema, List<TableInfo> tables) {
|
||||
synchronized (tableCacheLock) {
|
||||
tableCacheSchema = normalizedSchema(schema);
|
||||
tableCache = copyTables(tables);
|
||||
tableCacheTimeMillis = System.currentTimeMillis();
|
||||
}
|
||||
sql.append(" ORDER BY hierarchy_name, hierarchy_rank, table_name");
|
||||
MetadataSqlSupport.appendLiteralLimitOffset(sql, normalized);
|
||||
return new TableMetadataQuery(sql.toString(), args);
|
||||
}
|
||||
|
||||
static final class TableMetadataQuery {
|
||||
private final String sql;
|
||||
private final List<Object> args;
|
||||
|
||||
TableMetadataQuery(String sql, List<Object> args) {
|
||||
this.sql = sql;
|
||||
this.args = args;
|
||||
private void clearTableCache() {
|
||||
synchronized (tableCacheLock) {
|
||||
tableCacheSchema = "";
|
||||
tableCacheTimeMillis = 0L;
|
||||
tableCache = Collections.emptyList();
|
||||
}
|
||||
}
|
||||
|
||||
String sql() {
|
||||
return sql;
|
||||
}
|
||||
private static boolean cacheFresh(long cachedAtMillis) {
|
||||
return cachedAtMillis > 0L && System.currentTimeMillis() - cachedAtMillis <= TABLE_CACHE_TTL_MILLIS;
|
||||
}
|
||||
|
||||
List<Object> args() {
|
||||
return args;
|
||||
private static String normalizedSchema(String schema) {
|
||||
return schema == null ? "" : schema.trim();
|
||||
}
|
||||
|
||||
private static List<TableInfo> copyTables(List<TableInfo> tables) {
|
||||
List<TableInfo> copies = new ArrayList<>(tables.size());
|
||||
for (TableInfo table : tables) {
|
||||
copies.add(new TableInfo(
|
||||
table.getName(),
|
||||
table.getTable_type(),
|
||||
table.getComment(),
|
||||
table.getParent_schema(),
|
||||
table.getParent_name()
|
||||
));
|
||||
}
|
||||
return copies;
|
||||
}
|
||||
|
||||
private static boolean mayChangeMetadata(String sql) {
|
||||
String normalized = sql == null ? "" : sql.trim().toLowerCase(Locale.ROOT);
|
||||
return normalized.startsWith("create ")
|
||||
|| normalized.startsWith("drop ")
|
||||
|| normalized.startsWith("alter ")
|
||||
|| normalized.startsWith("rename ")
|
||||
|| normalized.startsWith("truncate ");
|
||||
}
|
||||
|
||||
static void sortTablesForHierarchy(List<TableInfo> tables) {
|
||||
|
|
|
|||
|
|
@ -16,17 +16,12 @@ import java.lang.reflect.InvocationHandler;
|
|||
import java.lang.reflect.Proxy;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.sql.Connection;
|
||||
import java.sql.PreparedStatement;
|
||||
import java.sql.ResultSet;
|
||||
import java.sql.Timestamp;
|
||||
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;
|
||||
import org.junit.jupiter.api.Assertions;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
|
|
@ -80,7 +75,7 @@ class TDengineAgentMetadataTest {
|
|||
}
|
||||
|
||||
@Test
|
||||
void usesTdengineMetadataStatements() {
|
||||
void usesTdengineShowStatementsForMetadata() {
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection());
|
||||
|
||||
|
|
@ -89,66 +84,119 @@ class TDengineAgentMetadataTest {
|
|||
agent.getColumns("power", "meters");
|
||||
|
||||
Assertions.assertEquals("SHOW DATABASES", JdbcMetadataSqlFake.statements.get(0));
|
||||
Assertions.assertTrue(JdbcMetadataSqlFake.statements.get(1).contains("FROM information_schema.ins_stables"));
|
||||
Assertions.assertTrue(JdbcMetadataSqlFake.statements.get(1).contains("FROM information_schema.ins_tables"));
|
||||
Assertions.assertEquals("param:1=power", JdbcMetadataSqlFake.statements.get(2));
|
||||
Assertions.assertEquals("param:2=power", JdbcMetadataSqlFake.statements.get(3));
|
||||
Assertions.assertEquals("DESCRIBE `power`.`meters`", JdbcMetadataSqlFake.statements.get(4));
|
||||
Assertions.assertEquals("SHOW `power`.STABLES", JdbcMetadataSqlFake.statements.get(1));
|
||||
Assertions.assertEquals("SHOW `power`.TABLES", JdbcMetadataSqlFake.statements.get(2));
|
||||
Assertions.assertFalse(JdbcMetadataSqlFake.statements.stream().anyMatch(sql -> sql.contains("information_schema")));
|
||||
Assertions.assertTrue(JdbcMetadataSqlFake.statements.contains("DESCRIBE `power`.`meters`"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void constrainedMetadataUsesRequestedDatabaseAndPushesFilterAndPaging() {
|
||||
void showMetadataCachesFullResultsForLocalPaging() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection());
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows));
|
||||
|
||||
List<TableInfo> firstPage = agent.listTables(
|
||||
"dbx_alpha",
|
||||
new MetadataListConstraints(null, 2, null, List.of("TABLE"))
|
||||
);
|
||||
List<TableInfo> secondPage = agent.listTables(
|
||||
"dbx_alpha",
|
||||
new MetadataListConstraints(null, 2, 2, List.of("TABLE"))
|
||||
);
|
||||
|
||||
Assertions.assertEquals(List.of("meters", "device_a"), firstPage.stream().map(TableInfo::getName).toList());
|
||||
Assertions.assertEquals(List.of("device_b", "standalone"), secondPage.stream().map(TableInfo::getName).toList());
|
||||
Assertions.assertEquals("meters", secondPage.get(0).getParent_name());
|
||||
Assertions.assertTrue(maxRows.isEmpty());
|
||||
Assertions.assertEquals(
|
||||
List.of(
|
||||
"SHOW `dbx_alpha`.STABLES",
|
||||
"SHOW `dbx_alpha`.TABLES"
|
||||
),
|
||||
statements
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void showMetadataFiltersLocallyBeforeApplyingRowLimit() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows));
|
||||
|
||||
List<TableInfo> tables = agent.listTables(
|
||||
"dbx_alpha",
|
||||
new MetadataListConstraints("stand", 1, null, List.of("TABLE"))
|
||||
);
|
||||
|
||||
Assertions.assertEquals(List.of("standalone"), tables.stream().map(TableInfo::getName).toList());
|
||||
Assertions.assertTrue(maxRows.isEmpty());
|
||||
Assertions.assertEquals(
|
||||
List.of(
|
||||
"SHOW `dbx_alpha`.STABLES",
|
||||
"SHOW `dbx_alpha`.TABLES"
|
||||
),
|
||||
statements
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void showMetadataFallsBackToLocalFilteringForLongLikePatterns() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows));
|
||||
|
||||
agent.listTables(
|
||||
"dbx_alpha",
|
||||
new MetadataListConstraints("ord", 25, 50, List.of("TABLE"))
|
||||
new MetadataListConstraints("a".repeat(50), 1, null, List.of("TABLE"))
|
||||
);
|
||||
|
||||
String sql = JdbcMetadataSqlFake.statements.get(0);
|
||||
Assertions.assertFalse(sql.contains("DATABASE()"), sql);
|
||||
Assertions.assertTrue(sql.contains("FROM information_schema.ins_stables"), sql);
|
||||
Assertions.assertTrue(sql.contains("FROM information_schema.ins_tables"), sql);
|
||||
Assertions.assertTrue(sql.contains("WHERE db_name = ?"), sql);
|
||||
Assertions.assertTrue(sql.contains("table_name LIKE ? OR table_comment LIKE ?"), sql);
|
||||
Assertions.assertTrue(sql.contains("ORDER BY hierarchy_name, hierarchy_rank, table_name"), sql);
|
||||
Assertions.assertTrue(sql.endsWith("LIMIT 25 OFFSET 50"), sql);
|
||||
Assertions.assertTrue(maxRows.isEmpty());
|
||||
Assertions.assertEquals(
|
||||
List.of(
|
||||
"param:1=dbx_alpha",
|
||||
"param:2=dbx_alpha",
|
||||
"param:3=%o%r%d%",
|
||||
"param:4=%o%r%d%"
|
||||
"SHOW `dbx_alpha`.STABLES",
|
||||
"SHOW `dbx_alpha`.TABLES"
|
||||
),
|
||||
JdbcMetadataSqlFake.statements.subList(1, 5)
|
||||
statements
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void constrainedMetadataPagesDoNotRepeatOrSkipTables() {
|
||||
List<TableInfo> metadata = List.of(
|
||||
new TableInfo("meters", "STABLE", null),
|
||||
new TableInfo("device_a", "TABLE", null, null, "meters"),
|
||||
new TableInfo("device_b", "TABLE", null, null, "meters"),
|
||||
new TableInfo("standalone", "TABLE", null),
|
||||
new TableInfo("weather", "STABLE", null)
|
||||
);
|
||||
void showMetadataDoesNotLoadTableComments() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, pagedMetadataConnection(metadata, statements));
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows));
|
||||
|
||||
List<TableInfo> combined = new ArrayList<>();
|
||||
combined.addAll(agent.listTables("dbx_alpha", new MetadataListConstraints(null, 2, null, List.of("TABLE"))));
|
||||
combined.addAll(agent.listTables("dbx_alpha", new MetadataListConstraints(null, 2, 2, List.of("TABLE"))));
|
||||
combined.addAll(agent.listTables("dbx_alpha", new MetadataListConstraints(null, 2, 4, List.of("TABLE"))));
|
||||
List<TableInfo> tables = agent.listTables("dbx_alpha");
|
||||
|
||||
Assertions.assertEquals(metadata.stream().map(TableInfo::getName).toList(), combined.stream().map(TableInfo::getName).toList());
|
||||
Set<String> distinctNames = new HashSet<>(combined.stream().map(TableInfo::getName).toList());
|
||||
Assertions.assertEquals(metadata.size(), distinctNames.size());
|
||||
Assertions.assertTrue(statements.get(0).endsWith("LIMIT 2"), statements.get(0));
|
||||
Assertions.assertTrue(statements.get(1).endsWith("LIMIT 2 OFFSET 2"), statements.get(1));
|
||||
Assertions.assertTrue(statements.get(2).endsWith("LIMIT 2 OFFSET 4"), statements.get(2));
|
||||
Assertions.assertNull(tables.get(0).getComment());
|
||||
Assertions.assertEquals(
|
||||
List.of(
|
||||
"SHOW `dbx_alpha`.STABLES",
|
||||
"SHOW `dbx_alpha`.TABLES"
|
||||
),
|
||||
statements
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void showMetadataFiltersByNameWithoutInformationSchema() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows));
|
||||
|
||||
List<TableInfo> tables = agent.listTables(
|
||||
"dbx_alpha",
|
||||
new MetadataListConstraints("device_b", 1, null, List.of("TABLE"))
|
||||
);
|
||||
|
||||
Assertions.assertEquals(List.of("device_b"), tables.stream().map(TableInfo::getName).toList());
|
||||
Assertions.assertTrue(maxRows.isEmpty());
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -198,13 +246,11 @@ class TDengineAgentMetadataTest {
|
|||
Assertions.assertTrue(JdbcMetadataSqlFake.statements.isEmpty());
|
||||
}
|
||||
|
||||
private static Connection pagedMetadataConnection(List<TableInfo> metadata, List<String> statements) {
|
||||
private static Connection showMetadataConnection(List<String> statements, List<Integer> maxRows) {
|
||||
return proxy(Connection.class, (proxy, method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("prepareStatement".equals(name)) {
|
||||
String sql = (String) args[0];
|
||||
statements.add(sql);
|
||||
return pagedMetadataStatement(sql, metadata);
|
||||
if ("createStatement".equals(name)) {
|
||||
return showMetadataStatement(statements, maxRows);
|
||||
}
|
||||
if ("isClosed".equals(name)) return false;
|
||||
if ("close".equals(name)) return null;
|
||||
|
|
@ -212,45 +258,55 @@ class TDengineAgentMetadataTest {
|
|||
});
|
||||
}
|
||||
|
||||
private static PreparedStatement pagedMetadataStatement(String sql, List<TableInfo> metadata) {
|
||||
return proxy(PreparedStatement.class, (proxy, method, args) -> {
|
||||
private static java.sql.Statement showMetadataStatement(List<String> statements, List<Integer> maxRows) {
|
||||
int[] activeMaxRows = {0};
|
||||
return proxy(java.sql.Statement.class, (proxy, method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("setMaxRows".equals(name)) {
|
||||
activeMaxRows[0] = (Integer) args[0];
|
||||
maxRows.add(activeMaxRows[0]);
|
||||
return null;
|
||||
}
|
||||
if ("executeQuery".equals(name)) {
|
||||
int limit = sqlClauseValue(sql, "LIMIT", metadata.size());
|
||||
int offset = sqlClauseValue(sql, "OFFSET", 0);
|
||||
int from = Math.min(offset, metadata.size());
|
||||
int to = Math.min(from + limit, metadata.size());
|
||||
return tableInfoResultSet(metadata.subList(from, to));
|
||||
}
|
||||
if ("setString".equals(name) || "setInt".equals(name) || "setObject".equals(name) || "close".equals(name)) return null;
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static ResultSet tableInfoResultSet(List<TableInfo> tables) {
|
||||
int[] index = {-1};
|
||||
return proxy(ResultSet.class, (proxy, method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("next".equals(name)) {
|
||||
index[0] += 1;
|
||||
return index[0] < tables.size();
|
||||
}
|
||||
if ("getString".equals(name)) {
|
||||
TableInfo table = tables.get(index[0]);
|
||||
int column = (Integer) args[0];
|
||||
if (column == 1) return table.getName();
|
||||
if (column == 2) return table.getTable_type();
|
||||
if (column == 3) return table.getComment();
|
||||
if (column == 4) return table.getParent_name();
|
||||
String sql = (String) args[0];
|
||||
statements.add(sql);
|
||||
if (sql.endsWith("STABLES")) {
|
||||
return showTableResultSet(List.of(new String[] {"meters"}, new String[] {"weather"}), activeMaxRows[0]);
|
||||
}
|
||||
if (sql.endsWith("TABLES")) {
|
||||
return showTableResultSet(
|
||||
List.of(
|
||||
new String[] {"device_a", "", "", "meters"},
|
||||
new String[] {"device_b", "", "", "meters"},
|
||||
new String[] {"standalone", "", "", ""}
|
||||
),
|
||||
activeMaxRows[0]
|
||||
);
|
||||
}
|
||||
}
|
||||
if ("close".equals(name)) return null;
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static int sqlClauseValue(String sql, String clause, int fallback) {
|
||||
Matcher matcher = Pattern.compile("\\b" + clause + "\\s+(\\d+)").matcher(sql);
|
||||
return matcher.find() ? Integer.parseInt(matcher.group(1)) : fallback;
|
||||
private static ResultSet showTableResultSet(List<String[]> rows, int maxRows) {
|
||||
List<String[]> limitedRows = maxRows > 0 ? rows.subList(0, Math.min(rows.size(), maxRows)) : rows;
|
||||
int[] index = {-1};
|
||||
return proxy(ResultSet.class, (proxy, method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("next".equals(name)) {
|
||||
index[0] += 1;
|
||||
return index[0] < limitedRows.size();
|
||||
}
|
||||
if ("getString".equals(name)) {
|
||||
int column = (Integer) args[0];
|
||||
String[] row = limitedRows.get(index[0]);
|
||||
return column <= row.length ? row[column - 1] : null;
|
||||
}
|
||||
if ("isClosed".equals(name)) return false;
|
||||
if ("close".equals(name)) return null;
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static <T> T proxy(Class<T> type, InvocationHandler handler) {
|
||||
|
|
|
|||
|
|
@ -15,6 +15,13 @@ use crate::models::connection::DatabaseConnectionInfo;
|
|||
|
||||
type PendingAgentResponse = tokio::sync::oneshot::Sender<Result<Value, String>>;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct CachedAgentQuery {
|
||||
key: String,
|
||||
expires_at: Instant,
|
||||
result: Result<Value, String>,
|
||||
}
|
||||
|
||||
pub struct AgentRuntimeClient {
|
||||
child: Arc<Mutex<Child>>,
|
||||
stdin: Arc<Mutex<BufWriter<ChildStdin>>>,
|
||||
|
|
@ -153,17 +160,27 @@ impl AgentRuntimeClient {
|
|||
let response = match (timeout_duration, cancel_token) {
|
||||
(Some(duration), Some(token)) => tokio::select! {
|
||||
_ = token.cancelled() => {
|
||||
self.cancel_session_request(¶ms).await;
|
||||
self.cancel_session_request_for_method(method, ¶ms).await;
|
||||
Err("Query canceled".to_string())
|
||||
},
|
||||
result = tokio::time::timeout(duration, receive) => result.map_err(|_| format!("Agent RPC call timed out ({}s)", duration.as_secs()))?,
|
||||
result = tokio::time::timeout(duration, receive) => match result {
|
||||
Ok(result) => result,
|
||||
Err(_) => {
|
||||
self.cancel_session_request_for_method(method, ¶ms).await;
|
||||
Err(format!("Agent RPC call timed out ({}s)", duration.as_secs()))
|
||||
}
|
||||
},
|
||||
},
|
||||
(Some(duration), None) => match tokio::time::timeout(duration, receive).await {
|
||||
Ok(result) => result,
|
||||
Err(_) => {
|
||||
self.cancel_session_request_for_method(method, ¶ms).await;
|
||||
Err(format!("Agent RPC call timed out ({}s)", duration.as_secs()))
|
||||
}
|
||||
},
|
||||
(Some(duration), None) => tokio::time::timeout(duration, receive)
|
||||
.await
|
||||
.map_err(|_| format!("Agent RPC call timed out ({}s)", duration.as_secs()))?,
|
||||
(None, Some(token)) => tokio::select! {
|
||||
_ = token.cancelled() => {
|
||||
self.cancel_session_request(¶ms).await;
|
||||
self.cancel_session_request_for_method(method, ¶ms).await;
|
||||
Err("Query canceled".to_string())
|
||||
},
|
||||
result = receive => result,
|
||||
|
|
@ -176,6 +193,12 @@ impl AgentRuntimeClient {
|
|||
decode_agent_response(response?)
|
||||
}
|
||||
|
||||
async fn cancel_session_request_for_method(&self, method: &str, params: &Value) {
|
||||
if method != AgentMethod::CancelSession.as_str() {
|
||||
self.cancel_session_request(params).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn cancel_session_request(&self, params: &Value) {
|
||||
let Some(agent_session_id) = params.get("agentSessionId").and_then(Value::as_str) else {
|
||||
return;
|
||||
|
|
@ -260,6 +283,11 @@ fn decode_agent_response<T: DeserializeOwned>(response: Value) -> Result<T, Stri
|
|||
serde_json::from_value(result.clone()).map_err(|e| format!("Failed to deserialize agent result: {e}"))
|
||||
}
|
||||
|
||||
fn deserialize_cached_agent_result<T: DeserializeOwned>(result: Result<Value, String>) -> Result<T, String> {
|
||||
result
|
||||
.and_then(|value| serde_json::from_value(value).map_err(|e| format!("Failed to deserialize agent result: {e}")))
|
||||
}
|
||||
|
||||
fn fail_pending_requests(pending: &Arc<Mutex<HashMap<u64, PendingAgentResponse>>>, error: String) {
|
||||
let requests = std::mem::take(&mut *pending.lock().expect("agent pending response lock poisoned"));
|
||||
for (_, sender) in requests {
|
||||
|
|
@ -287,6 +315,7 @@ pub struct AgentDriverClient {
|
|||
next_id: u64,
|
||||
shared_runtime: Option<Arc<AgentRuntimeClient>>,
|
||||
agent_session_id: Option<String>,
|
||||
cached_query: Option<CachedAgentQuery>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
|
|
@ -756,6 +785,7 @@ impl AgentDriverClient {
|
|||
next_id: 0,
|
||||
shared_runtime: None,
|
||||
agent_session_id: None,
|
||||
cached_query: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -769,6 +799,7 @@ impl AgentDriverClient {
|
|||
next_id: 0,
|
||||
shared_runtime: Some(runtime),
|
||||
agent_session_id: Some(agent_session_id),
|
||||
cached_query: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -989,6 +1020,7 @@ impl AgentDriverClient {
|
|||
}
|
||||
|
||||
pub async fn disconnect(&mut self) -> Result<Value, String> {
|
||||
self.invalidate_cached_query();
|
||||
if self.shared_runtime.is_some() {
|
||||
let session_id = self.agent_session_id.as_ref().ok_or("Shared Agent session id is missing")?.clone();
|
||||
let result =
|
||||
|
|
@ -1292,6 +1324,7 @@ impl AgentDriverClient {
|
|||
}
|
||||
|
||||
pub async fn execute_query<T: DeserializeOwned + Send + 'static>(&mut self, params: Value) -> Result<T, String> {
|
||||
self.invalidate_cached_query();
|
||||
self.call_method(AgentMethod::ExecuteQuery, params).await
|
||||
}
|
||||
|
||||
|
|
@ -1300,6 +1333,7 @@ impl AgentDriverClient {
|
|||
params: Value,
|
||||
timeout_duration: Option<Duration>,
|
||||
) -> Result<T, String> {
|
||||
self.invalidate_cached_query();
|
||||
self.call_method_with_timeout(AgentMethod::ExecuteQuery, params, timeout_duration).await
|
||||
}
|
||||
|
||||
|
|
@ -1309,14 +1343,37 @@ impl AgentDriverClient {
|
|||
timeout_duration: Option<Duration>,
|
||||
cancel_token: Option<CancellationToken>,
|
||||
) -> Result<T, String> {
|
||||
self.invalidate_cached_query();
|
||||
self.call_method_with_timeout_and_cancel(AgentMethod::ExecuteQuery, params, timeout_duration, cancel_token)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn execute_query_cached_with_timeout<T: DeserializeOwned + Send + 'static>(
|
||||
&mut self,
|
||||
cache_key: String,
|
||||
cache_ttl: Duration,
|
||||
params: Value,
|
||||
timeout_duration: Option<Duration>,
|
||||
) -> Result<T, String> {
|
||||
let now = Instant::now();
|
||||
if let Some(cached) = self.cached_query.as_ref() {
|
||||
if cached.key == cache_key && cached.expires_at > now {
|
||||
return deserialize_cached_agent_result(cached.result.clone());
|
||||
}
|
||||
}
|
||||
|
||||
self.cached_query = None;
|
||||
let result = self.call_method_with_timeout::<Value>(AgentMethod::ExecuteQuery, params, timeout_duration).await;
|
||||
self.cached_query =
|
||||
Some(CachedAgentQuery { key: cache_key, expires_at: Instant::now() + cache_ttl, result: result.clone() });
|
||||
deserialize_cached_agent_result(result)
|
||||
}
|
||||
|
||||
pub async fn execute_query_page<T: DeserializeOwned + Send + 'static>(
|
||||
&mut self,
|
||||
params: Value,
|
||||
) -> Result<T, String> {
|
||||
self.invalidate_cached_query();
|
||||
self.call_method(AgentMethod::ExecuteQueryPage, params).await
|
||||
}
|
||||
|
||||
|
|
@ -1325,6 +1382,7 @@ impl AgentDriverClient {
|
|||
params: Value,
|
||||
timeout_duration: Option<Duration>,
|
||||
) -> Result<T, String> {
|
||||
self.invalidate_cached_query();
|
||||
self.call_method_with_timeout(AgentMethod::ExecuteQueryPage, params, timeout_duration).await
|
||||
}
|
||||
|
||||
|
|
@ -1334,6 +1392,7 @@ impl AgentDriverClient {
|
|||
timeout_duration: Option<Duration>,
|
||||
cancel_token: Option<CancellationToken>,
|
||||
) -> Result<T, String> {
|
||||
self.invalidate_cached_query();
|
||||
self.call_method_with_timeout_and_cancel(AgentMethod::ExecuteQueryPage, params, timeout_duration, cancel_token)
|
||||
.await
|
||||
}
|
||||
|
|
@ -1375,6 +1434,7 @@ impl AgentDriverClient {
|
|||
&mut self,
|
||||
params: AgentTableReadStartParams,
|
||||
) -> Result<T, String> {
|
||||
self.invalidate_cached_query();
|
||||
self.call_method(AgentMethod::StartTableRead, serde_json::to_value(params).map_err(|e| e.to_string())?).await
|
||||
}
|
||||
|
||||
|
|
@ -1409,6 +1469,7 @@ impl AgentDriverClient {
|
|||
statements: &[String],
|
||||
schema: Option<&str>,
|
||||
) -> Result<T, String> {
|
||||
self.invalidate_cached_query();
|
||||
self.call_method(AgentMethod::ExecuteTransaction, agent_transaction_params(database, statements, schema)).await
|
||||
}
|
||||
|
||||
|
|
@ -1419,6 +1480,7 @@ impl AgentDriverClient {
|
|||
schema: Option<&str>,
|
||||
timeout_duration: Option<Duration>,
|
||||
) -> Result<T, String> {
|
||||
self.invalidate_cached_query();
|
||||
self.call_method_with_timeout(
|
||||
AgentMethod::ExecuteBatch,
|
||||
agent_transaction_params(database, statements, schema),
|
||||
|
|
@ -1427,6 +1489,10 @@ impl AgentDriverClient {
|
|||
.await
|
||||
}
|
||||
|
||||
fn invalidate_cached_query(&mut self) {
|
||||
self.cached_query = None;
|
||||
}
|
||||
|
||||
pub async fn call_mongo_method<T: DeserializeOwned + Send + 'static>(
|
||||
&mut self,
|
||||
method: MongoAgentMethod,
|
||||
|
|
@ -1971,6 +2037,7 @@ impl AgentDriverClient {
|
|||
next_id: 0,
|
||||
shared_runtime: None,
|
||||
agent_session_id: None,
|
||||
cached_query: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -2013,7 +2080,7 @@ mod tests {
|
|||
use std::io::Write;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
#[test]
|
||||
|
|
@ -2216,6 +2283,7 @@ mod tests {
|
|||
next_id: 0,
|
||||
shared_runtime: None,
|
||||
agent_session_id: None,
|
||||
cached_query: None,
|
||||
};
|
||||
|
||||
let started_at = std::time::Instant::now();
|
||||
|
|
@ -2234,6 +2302,293 @@ mod tests {
|
|||
assert!(!is_agent_rpc_response_error("Failed to read response from agent: end of stream"));
|
||||
}
|
||||
|
||||
async fn spawn_stateful_test_runtime(prefix: &str) -> (Arc<AgentRuntimeClient>, std::path::PathBuf) {
|
||||
let script_path = std::env::temp_dir().join(format!("dbx-agent-{prefix}-{}.py", uuid::Uuid::new_v4()));
|
||||
std::fs::write(
|
||||
&script_path,
|
||||
r#"import json, sys, threading
|
||||
print(json.dumps({'ready': True}), flush=True)
|
||||
state_lock = threading.Lock()
|
||||
output_lock = threading.Lock()
|
||||
session_locks = {}
|
||||
blocked = {}
|
||||
query_count = 0
|
||||
cancel_count = 0
|
||||
|
||||
def write_response(req, result=None, error=None):
|
||||
response = {'jsonrpc': '2.0', 'id': req['id']}
|
||||
if error is None:
|
||||
response['result'] = result
|
||||
else:
|
||||
response['error'] = {'code': -1, 'message': error}
|
||||
with output_lock:
|
||||
print(json.dumps(response), flush=True)
|
||||
|
||||
def session_lock(session_id):
|
||||
with state_lock:
|
||||
return session_locks.setdefault(session_id, threading.Lock())
|
||||
|
||||
def respond(req):
|
||||
global query_count, cancel_count
|
||||
method = req['method']
|
||||
params = req.get('params', {})
|
||||
if method == 'handshake':
|
||||
write_response(req, {'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']})
|
||||
return
|
||||
if method == 'query_count':
|
||||
with state_lock:
|
||||
value = query_count
|
||||
write_response(req, value)
|
||||
return
|
||||
if method == 'cancel_count':
|
||||
with state_lock:
|
||||
value = cancel_count
|
||||
write_response(req, value)
|
||||
return
|
||||
if method == 'cancel_session':
|
||||
session_id = params['agentSessionId']
|
||||
with state_lock:
|
||||
cancel_count += 1
|
||||
event = blocked.get(session_id)
|
||||
if event is not None:
|
||||
event.set()
|
||||
write_response(req, {'ok': True})
|
||||
return
|
||||
|
||||
session_id = params.get('agentSessionId', '__legacy__')
|
||||
with session_lock(session_id):
|
||||
if method == 'execute_query':
|
||||
sql = params.get('sql', '')
|
||||
with state_lock:
|
||||
query_count += 1
|
||||
current_count = query_count
|
||||
if sql.startswith('slow'):
|
||||
event = threading.Event()
|
||||
with state_lock:
|
||||
blocked[session_id] = event
|
||||
event.wait(2)
|
||||
with state_lock:
|
||||
blocked.pop(session_id, None)
|
||||
if sql == 'error':
|
||||
write_response(req, error='synthetic query error')
|
||||
return
|
||||
write_response(req, {'sql': sql, 'count': current_count})
|
||||
return
|
||||
write_response(req, {'ok': True})
|
||||
|
||||
for line in sys.stdin:
|
||||
threading.Thread(target=respond, args=(json.loads(line),), daemon=True).start()
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let runtime = AgentRuntimeClient::spawn(
|
||||
AgentLaunchSpec::new("python3").with_args([script_path.to_string_lossy().to_string()]),
|
||||
"test",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
(runtime, script_path)
|
||||
}
|
||||
|
||||
async fn runtime_counter(runtime: &Arc<AgentRuntimeClient>, method: &str) -> u64 {
|
||||
runtime.call(method, serde_json::json!({}), Some(Duration::from_secs(2)), None).await.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shared_runtime_timeout_cancels_session_with_and_without_token() {
|
||||
let (runtime, script_path) = spawn_stateful_test_runtime("timeout-cancel-test").await;
|
||||
let mut client = AgentDriverClient::shared_session(runtime.clone(), "timeout-session".to_string());
|
||||
|
||||
let error = client
|
||||
.execute_query_with_timeout::<serde_json::Value>(
|
||||
serde_json::json!({"sql": "slow-without-token"}),
|
||||
Some(Duration::from_millis(75)),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.contains("Agent RPC call timed out"));
|
||||
let started = Instant::now();
|
||||
client
|
||||
.call_with_timeout::<serde_json::Value>("probe", serde_json::json!({}), Some(Duration::from_millis(500)))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(started.elapsed() < Duration::from_millis(500));
|
||||
|
||||
let error = client
|
||||
.execute_query_with_timeout_and_cancel::<serde_json::Value>(
|
||||
serde_json::json!({"sql": "slow-with-token"}),
|
||||
Some(Duration::from_millis(75)),
|
||||
Some(CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.contains("Agent RPC call timed out"));
|
||||
let started = Instant::now();
|
||||
client
|
||||
.call_with_timeout::<serde_json::Value>("probe", serde_json::json!({}), Some(Duration::from_millis(500)))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(started.elapsed() < Duration::from_millis(500));
|
||||
assert_eq!(runtime_counter(&runtime, "cancel_count").await, 2);
|
||||
|
||||
runtime.kill();
|
||||
let _ = std::fs::remove_file(script_path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cached_agent_query_handles_hits_errors_ttl_and_invalidation() {
|
||||
let (runtime, script_path) = spawn_stateful_test_runtime("query-cache-test").await;
|
||||
let mut client = AgentDriverClient::shared_session(runtime.clone(), "cache-session".to_string());
|
||||
let timeout = Some(Duration::from_secs(2));
|
||||
let cache_ttl = Duration::from_secs(1);
|
||||
|
||||
let first: serde_json::Value = client
|
||||
.execute_query_cached_with_timeout(
|
||||
"same-key".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "same"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let second: serde_json::Value = client
|
||||
.execute_query_cached_with_timeout(
|
||||
"same-key".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "same"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(first, second);
|
||||
assert_eq!(runtime_counter(&runtime, "query_count").await, 1);
|
||||
|
||||
let first_error = client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"error-key".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "slow-cache-error"}),
|
||||
Some(Duration::from_millis(75)),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
let second_error = client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"error-key".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "slow-cache-error"}),
|
||||
Some(Duration::from_millis(75)),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(first_error, second_error);
|
||||
assert!(first_error.contains("Agent RPC call timed out"));
|
||||
assert_eq!(runtime_counter(&runtime, "query_count").await, 2);
|
||||
|
||||
client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"ttl-key".to_string(),
|
||||
Duration::from_millis(20),
|
||||
serde_json::json!({"sql": "ttl"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(40)).await;
|
||||
client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"ttl-key".to_string(),
|
||||
Duration::from_millis(20),
|
||||
serde_json::json!({"sql": "ttl"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(runtime_counter(&runtime, "query_count").await, 4);
|
||||
|
||||
client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"query-invalidation".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "cached-before-query"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
client.execute_query::<serde_json::Value>(serde_json::json!({"sql": "ordinary"})).await.unwrap();
|
||||
client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"query-invalidation".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "cached-before-query"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(runtime_counter(&runtime, "query_count").await, 7);
|
||||
|
||||
client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"batch-invalidation".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "cached-before-batch"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
client
|
||||
.execute_batch::<serde_json::Value>(None, &["UPDATE t SET a = 1".to_string()], None, timeout)
|
||||
.await
|
||||
.unwrap();
|
||||
client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"batch-invalidation".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "cached-before-batch"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(runtime_counter(&runtime, "query_count").await, 9);
|
||||
|
||||
client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"transaction-invalidation".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "cached-before-transaction"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
client.execute_transaction::<serde_json::Value>(None, &["UPDATE t SET a = 2".to_string()], None).await.unwrap();
|
||||
client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"transaction-invalidation".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "cached-before-transaction"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(runtime_counter(&runtime, "query_count").await, 11);
|
||||
|
||||
client
|
||||
.execute_query_cached_with_timeout::<serde_json::Value>(
|
||||
"disconnect-invalidation".to_string(),
|
||||
cache_ttl,
|
||||
serde_json::json!({"sql": "cached-before-disconnect"}),
|
||||
timeout,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(client.cached_query.is_some());
|
||||
client.disconnect().await.unwrap();
|
||||
assert!(client.cached_query.is_none());
|
||||
|
||||
runtime.kill();
|
||||
let _ = std::fs::remove_file(script_path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multiplexed_runtime_correlates_out_of_order_responses() {
|
||||
let script_path = std::env::temp_dir().join(format!("dbx-agent-runtime-test-{}.py", uuid::Uuid::new_v4()));
|
||||
|
|
|
|||
|
|
@ -40,6 +40,9 @@ macro_rules! try_sqlserver {
|
|||
}
|
||||
|
||||
const ORACLE_TABLE_COMMENT_BATCH_SIZE: usize = 500;
|
||||
const TDENGINE_COMMENT_SEARCH_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const TDENGINE_COMMENT_SEARCH_CACHE_TTL: Duration = Duration::from_secs(10);
|
||||
const TDENGINE_LIKE_PATTERN_MAX_BYTES: usize = 100;
|
||||
|
||||
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
|
|
@ -1473,6 +1476,25 @@ pub async fn get_table_comment_core(
|
|||
let mut client = client.lock().await;
|
||||
return client.get_table_comment::<Option<String>>(database, schema, table, timeout).await;
|
||||
}
|
||||
if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Tdengine) {
|
||||
let metadata_database = if schema.trim().is_empty() { database } else { schema };
|
||||
let sql = tdengine_table_comment_sql(metadata_database, table);
|
||||
let timeout = agent_metadata_timeout(db_config.as_ref());
|
||||
drop(connections);
|
||||
let mut client = client.lock().await;
|
||||
let result = client
|
||||
.execute_query_with_timeout::<db::QueryResult>(
|
||||
agent_execute_query_params(
|
||||
&sql,
|
||||
Some(database),
|
||||
(!schema.trim().is_empty()).then_some(schema),
|
||||
QueryExecutionOptions { max_rows: Some(2), ..Default::default() },
|
||||
),
|
||||
timeout,
|
||||
)
|
||||
.await?;
|
||||
return oracle_table_comment_from_query_result(result);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1504,6 +1526,55 @@ fn oracle_table_comment_sql(schema: &str, table: &str) -> String {
|
|||
)
|
||||
}
|
||||
|
||||
fn tdengine_table_comment_sql(database: &str, table: &str) -> String {
|
||||
format!(
|
||||
"SELECT table_comment FROM information_schema.ins_stables WHERE db_name = {} AND stable_name = {} \
|
||||
UNION ALL SELECT table_comment FROM information_schema.ins_tables WHERE db_name = {} AND table_name = {}",
|
||||
sql_string(database),
|
||||
sql_string(table),
|
||||
sql_string(database),
|
||||
sql_string(table),
|
||||
)
|
||||
}
|
||||
|
||||
fn tdengine_table_comments_sql(database: &str, filter: &str) -> String {
|
||||
let pattern = tdengine_table_comment_like_pattern(filter);
|
||||
format!(
|
||||
"SELECT stable_name, table_comment FROM information_schema.ins_stables \
|
||||
WHERE db_name = {database} AND table_comment IS NOT NULL AND LOWER(table_comment) LIKE {pattern} \
|
||||
UNION ALL SELECT table_name, table_comment FROM information_schema.ins_tables \
|
||||
WHERE db_name = {database} AND table_comment IS NOT NULL AND LOWER(table_comment) LIKE {pattern}",
|
||||
database = sql_string(database),
|
||||
pattern = sql_string(&pattern),
|
||||
)
|
||||
}
|
||||
|
||||
fn tdengine_table_comment_like_pattern(filter: &str) -> String {
|
||||
let normalized_filter = filter.trim().to_lowercase();
|
||||
if normalized_filter.is_empty() {
|
||||
return "%%".to_string();
|
||||
}
|
||||
|
||||
let mut pattern = String::with_capacity(TDENGINE_LIKE_PATTERN_MAX_BYTES);
|
||||
pattern.push('%');
|
||||
for ch in normalized_filter.chars() {
|
||||
let fragment = match ch {
|
||||
'\\' | '%' | '_' => format!("\\{ch}"),
|
||||
_ => ch.to_string(),
|
||||
};
|
||||
if pattern.len() + fragment.len() + 1 > TDENGINE_LIKE_PATTERN_MAX_BYTES {
|
||||
break;
|
||||
}
|
||||
pattern.push_str(&fragment);
|
||||
pattern.push('%');
|
||||
}
|
||||
pattern
|
||||
}
|
||||
|
||||
fn tdengine_table_comment_cache_key(database: &str, schema: &str, filter: &str) -> String {
|
||||
serde_json::json!([database, schema, filter.trim().to_lowercase()]).to_string()
|
||||
}
|
||||
|
||||
fn oracle_table_comment_from_query_result(result: db::QueryResult) -> Result<Option<String>, String> {
|
||||
Ok(result
|
||||
.rows
|
||||
|
|
@ -1524,7 +1595,7 @@ fn oracle_table_comments_sql(schema: &str, table_names: &[String]) -> Option<Str
|
|||
))
|
||||
}
|
||||
|
||||
fn oracle_table_comments_from_query_result(result: db::QueryResult) -> HashMap<String, String> {
|
||||
fn table_comments_from_query_result(result: db::QueryResult) -> HashMap<String, String> {
|
||||
result
|
||||
.rows
|
||||
.into_iter()
|
||||
|
|
@ -1953,7 +2024,7 @@ fn unique_oracle_comment_names<'a>(names: impl Iterator<Item = &'a str>) -> Vec<
|
|||
unique
|
||||
}
|
||||
|
||||
fn apply_oracle_table_comments(tables: &mut [db::TableInfo], comments: &HashMap<String, String>) {
|
||||
fn apply_table_comments(tables: &mut [db::TableInfo], comments: &HashMap<String, String>) {
|
||||
for table in tables {
|
||||
if !comment_is_blank(&table.comment) {
|
||||
continue;
|
||||
|
|
@ -2004,7 +2075,7 @@ async fn oracle_table_comments_for_names(
|
|||
timeout_duration,
|
||||
)
|
||||
.await?;
|
||||
comments.extend(oracle_table_comments_from_query_result(result));
|
||||
comments.extend(table_comments_from_query_result(result));
|
||||
}
|
||||
Ok(comments)
|
||||
}
|
||||
|
|
@ -2021,7 +2092,35 @@ async fn load_oracle_table_comments_for_tables(
|
|||
return Ok(());
|
||||
}
|
||||
let comments = oracle_table_comments_for_names(client, database, schema, &table_names, timeout_duration).await?;
|
||||
apply_oracle_table_comments(tables, &comments);
|
||||
apply_table_comments(tables, &comments);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn load_tdengine_table_comments_for_filter(
|
||||
client: &mut db::agent_driver::AgentDriverClient,
|
||||
database: &str,
|
||||
schema: &str,
|
||||
filter: &str,
|
||||
tables: &mut [db::TableInfo],
|
||||
) -> Result<(), String> {
|
||||
let metadata_database = if schema.trim().is_empty() { database } else { schema };
|
||||
let sql = tdengine_table_comments_sql(metadata_database, filter);
|
||||
let cache_key = tdengine_table_comment_cache_key(database, schema, filter);
|
||||
let result = client
|
||||
.execute_query_cached_with_timeout::<db::QueryResult>(
|
||||
cache_key,
|
||||
TDENGINE_COMMENT_SEARCH_CACHE_TTL,
|
||||
agent_execute_query_params(
|
||||
&sql,
|
||||
if database.is_empty() { None } else { Some(database) },
|
||||
if schema.is_empty() { None } else { Some(schema) },
|
||||
QueryExecutionOptions::default(),
|
||||
),
|
||||
Some(TDENGINE_COMMENT_SEARCH_TIMEOUT),
|
||||
)
|
||||
.await?;
|
||||
let comments = table_comments_from_query_result(result);
|
||||
apply_table_comments(tables, &comments);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
|
@ -2275,23 +2374,28 @@ async fn list_tables_once(
|
|||
try_sqlserver!(connections, &pool_key, list_tables, schema, filter, limit, offset);
|
||||
if let Some(client) = extract_pool!(&connections, &pool_key, Agent) {
|
||||
let is_oracle = db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Oracle);
|
||||
let is_tdengine = db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Tdengine);
|
||||
let use_agent_table_paging = db_config.as_ref().is_some_and(supports_agent_table_paging);
|
||||
let filter_locally_after_oracle_comments =
|
||||
is_oracle && filter.is_some_and(|filter| !filter.trim().is_empty());
|
||||
let filter_locally_after_tdengine_comments =
|
||||
is_tdengine && filter.is_some_and(|filter| !filter.trim().is_empty());
|
||||
let filter_locally_after_comments =
|
||||
filter_locally_after_oracle_comments || filter_locally_after_tdengine_comments;
|
||||
let timeout_duration = agent_metadata_timeout(db_config.as_ref());
|
||||
let fallback_config = db_config.clone();
|
||||
drop(connections);
|
||||
let mut client = client.lock().await;
|
||||
let agent_filter = if filter_locally_after_oracle_comments { None } else { filter };
|
||||
let agent_filter = if filter_locally_after_comments { None } else { filter };
|
||||
let force_local_table_name_filter = table_name_filter.is_some_and(|filter| !filter.is_empty());
|
||||
let agent_limit = if filter_locally_after_oracle_comments || force_local_table_name_filter {
|
||||
let agent_limit = if filter_locally_after_comments || force_local_table_name_filter {
|
||||
None
|
||||
} else if use_agent_table_paging {
|
||||
limit
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let agent_offset = if filter_locally_after_oracle_comments || force_local_table_name_filter {
|
||||
let agent_offset = if filter_locally_after_comments || force_local_table_name_filter {
|
||||
None
|
||||
} else if use_agent_table_paging {
|
||||
offset
|
||||
|
|
@ -2321,7 +2425,29 @@ async fn list_tables_once(
|
|||
)
|
||||
.await?;
|
||||
}
|
||||
let final_offset = if filter_locally_after_oracle_comments || force_local_table_name_filter {
|
||||
if filter_locally_after_tdengine_comments {
|
||||
if let Err(error) = load_tdengine_table_comments_for_filter(
|
||||
&mut client,
|
||||
database,
|
||||
schema,
|
||||
filter.expect("TDengine comment filtering requires a non-empty filter"),
|
||||
&mut tables,
|
||||
)
|
||||
.await
|
||||
{
|
||||
// TDengine 2.x can lack the information_schema views. SHOW
|
||||
// metadata remains usable, so preserve name filtering when
|
||||
// the optional comment lookup is unavailable or times out.
|
||||
log::warn!(
|
||||
"[schema][tdengine:list_tables:comment-search-failed] connection_id={} database={} schema={} error={}",
|
||||
connection_id,
|
||||
database,
|
||||
schema,
|
||||
error
|
||||
);
|
||||
}
|
||||
}
|
||||
let final_offset = if filter_locally_after_comments || force_local_table_name_filter {
|
||||
offset
|
||||
} else if agent_paging_likely_applied(use_agent_table_paging, limit, tables.len()) {
|
||||
Some(0)
|
||||
|
|
@ -2855,15 +2981,16 @@ mod tests {
|
|||
dameng_object_statistics_rows_only_sql, dameng_object_statistics_user_segments_sql, deduplicate_column_infos,
|
||||
filter_mysql_system_databases_for_config, filter_object_infos, filter_table_infos, filter_visible_schema_names,
|
||||
gbase8a_object_statistics_sql, is_agent_postgres_metadata_fallback_config, is_retryable_metadata_error,
|
||||
mysql_object_source_ddl_column_index, mysql_object_source_sql, mysql_table_metadata_catalog,
|
||||
normalize_information_schema_table_type, oracle_columns_from_query_result, oracle_columns_sql,
|
||||
oracle_object_statistics_dba_segments_sql, oracle_object_statistics_from_query_result,
|
||||
metadata_name_or_comment_matches, mysql_object_source_ddl_column_index, mysql_object_source_sql,
|
||||
mysql_table_metadata_catalog, normalize_information_schema_table_type, oracle_columns_from_query_result,
|
||||
oracle_columns_sql, oracle_object_statistics_dba_segments_sql, oracle_object_statistics_from_query_result,
|
||||
oracle_object_statistics_rows_only_sql, oracle_object_statistics_sql,
|
||||
oracle_object_statistics_user_segments_sql, oracle_table_comment_from_query_result, oracle_table_comment_sql,
|
||||
oracle_table_comments_from_query_result, oracle_table_comments_sql, presto_like_columns_from_query_result,
|
||||
presto_like_information_schema_columns_sql, presto_like_information_schema_tables_sql,
|
||||
presto_like_tables_from_query_result, should_query_oracle_columns_via_sql_first, table_name_filter_matches,
|
||||
visible_schema_filter, TableNameFilter,
|
||||
oracle_table_comments_sql, presto_like_columns_from_query_result, presto_like_information_schema_columns_sql,
|
||||
presto_like_information_schema_tables_sql, presto_like_tables_from_query_result,
|
||||
should_query_oracle_columns_via_sql_first, table_comments_from_query_result, table_name_filter_matches,
|
||||
tdengine_table_comment_like_pattern, tdengine_table_comment_sql, tdengine_table_comments_sql,
|
||||
visible_schema_filter, TableNameFilter, TDENGINE_COMMENT_SEARCH_TIMEOUT, TDENGINE_LIKE_PATTERN_MAX_BYTES,
|
||||
};
|
||||
#[cfg(feature = "duckdb-bundled")]
|
||||
use super::{
|
||||
|
|
@ -3724,6 +3851,69 @@ mod tests {
|
|||
assert!(!sql.contains("ALL_OBJECTS"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tdengine_table_comment_sql_targets_one_name_and_escapes_literals() {
|
||||
let sql = tdengine_table_comment_sql("dbx's", "meter's");
|
||||
|
||||
assert!(sql.contains("information_schema.ins_stables"));
|
||||
assert!(sql.contains("information_schema.ins_tables"));
|
||||
assert!(sql.contains("db_name = 'dbx''s'"));
|
||||
assert!(sql.contains("stable_name = 'meter''s'"));
|
||||
assert!(sql.contains("table_name = 'meter''s'"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tdengine_table_comments_sql_only_queries_comments_matching_the_filter() {
|
||||
let sql = tdengine_table_comments_sql("dbx's", "s_%\\q");
|
||||
|
||||
assert!(sql.contains("information_schema.ins_stables"));
|
||||
assert!(sql.contains("information_schema.ins_tables"));
|
||||
assert!(sql.contains("db_name = 'dbx''s'"));
|
||||
assert!(sql.contains("table_comment IS NOT NULL"));
|
||||
assert!(sql.contains("LOWER(table_comment) LIKE '%s%\\_%\\%%\\\\%q%'"));
|
||||
assert!(!sql.contains("LIMIT"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tdengine_table_comment_pattern_respects_ascii_boundary() {
|
||||
let filter_49 = "a".repeat(49);
|
||||
let filter_50 = format!("{filter_49}b");
|
||||
let pattern_49 = tdengine_table_comment_like_pattern(&filter_49);
|
||||
let pattern_50 = tdengine_table_comment_like_pattern(&filter_50);
|
||||
|
||||
assert_eq!(pattern_49.len(), 99);
|
||||
assert_eq!(pattern_50, pattern_49);
|
||||
assert!(pattern_50.len() <= TDENGINE_LIKE_PATTERN_MAX_BYTES);
|
||||
assert!(!metadata_name_or_comment_matches("table", Some(&filter_49), &filter_50));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tdengine_table_comment_pattern_keeps_escaped_fragments_within_limit() {
|
||||
assert_eq!(tdengine_table_comment_like_pattern("%_\\"), r"%\%%\_%\\%");
|
||||
|
||||
let pattern_33 = tdengine_table_comment_like_pattern(&"%".repeat(33));
|
||||
let pattern_34 = tdengine_table_comment_like_pattern(&"%".repeat(34));
|
||||
assert_eq!(pattern_33.len(), TDENGINE_LIKE_PATTERN_MAX_BYTES);
|
||||
assert_eq!(pattern_34, pattern_33);
|
||||
assert!(pattern_34.ends_with('%'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tdengine_table_comment_pattern_truncates_only_at_utf8_boundaries() {
|
||||
let pattern_24 = tdengine_table_comment_like_pattern(&"你".repeat(24));
|
||||
let pattern_25 = tdengine_table_comment_like_pattern(&"你".repeat(25));
|
||||
|
||||
assert_eq!(pattern_24.len(), 97);
|
||||
assert_eq!(pattern_25, pattern_24);
|
||||
assert!(pattern_25.is_char_boundary(pattern_25.len()));
|
||||
assert!(pattern_25.len() <= TDENGINE_LIKE_PATTERN_MAX_BYTES);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tdengine_comment_search_uses_a_short_outer_deadline() {
|
||||
assert_eq!(TDENGINE_COMMENT_SEARCH_TIMEOUT, std::time::Duration::from_secs(5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oracle_table_comment_from_query_result_returns_optional_non_blank_comment() {
|
||||
let result = db::QueryResult {
|
||||
|
|
@ -3770,7 +3960,7 @@ mod tests {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn oracle_table_comments_from_query_result_maps_non_blank_comments() {
|
||||
fn table_comments_from_query_result_maps_non_blank_comments() {
|
||||
let result = db::QueryResult {
|
||||
columns: vec!["TABLE_NAME".to_string(), "COMMENTS".to_string()],
|
||||
column_types: Vec::new(),
|
||||
|
|
@ -3787,7 +3977,7 @@ mod tests {
|
|||
elasticsearch_raw_body: None,
|
||||
};
|
||||
|
||||
let comments = oracle_table_comments_from_query_result(result);
|
||||
let comments = table_comments_from_query_result(result);
|
||||
assert_eq!(comments.get("ORDERS").map(String::as_str), Some("Orders table"));
|
||||
assert!(!comments.contains_key("PRODUCTS"));
|
||||
}
|
||||
|
|
@ -4021,7 +4211,7 @@ mod tests {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn apply_oracle_table_comments_only_fills_missing_table_comments() {
|
||||
fn apply_table_comments_only_fills_missing_table_comments() {
|
||||
let mut tables = vec![
|
||||
super::db::TableInfo {
|
||||
name: "ORDERS".to_string(),
|
||||
|
|
@ -4043,7 +4233,7 @@ mod tests {
|
|||
("PRODUCTS".to_string(), "Products table".to_string()),
|
||||
]);
|
||||
|
||||
super::apply_oracle_table_comments(&mut tables, &comments);
|
||||
super::apply_table_comments(&mut tables, &comments);
|
||||
|
||||
assert_eq!(tables[0].comment.as_deref(), Some("Orders table"));
|
||||
assert_eq!(tables[1].comment.as_deref(), Some("Existing"));
|
||||
|
|
|
|||
Loading…
Reference in New Issue