diff --git a/agents/drivers/tdengine/src/main/java/com/dbx/agent/tdengine/TDengineAgent.java b/agents/drivers/tdengine/src/main/java/com/dbx/agent/tdengine/TDengineAgent.java index 102c6fa08..efb917f2d 100644 --- a/agents/drivers/tdengine/src/main/java/com/dbx/agent/tdengine/TDengineAgent.java +++ b/agents/drivers/tdengine/src/main/java/com/dbx/agent/tdengine/TDengineAgent.java @@ -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 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 listTablesLegacy(String schema) { + private List listTablesFromShow(String schema) { + List cached = cachedTables(schema); + if (cached != null) { + return cached; + } return unchecked(() -> { List 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 distinct = new LinkedHashMap<>(); for (TableInfo table : result) { @@ -138,7 +145,8 @@ public final class TDengineAgent extends BaseDatabaseAgent { } List 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 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 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 queryTables(String sql, String tableType) throws Exception { + private List queryTables(String sql, String tableType, boolean includesStableName) throws Exception { List 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 queryConstrainedTables(String database, MetadataListConstraints constraints) { - return unchecked(() -> { - TableMetadataQuery query = buildTableMetadataQuery(database, constraints); - List 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 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 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 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 args; - - TableMetadataQuery(String sql, List 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 args() { - return args; + private static String normalizedSchema(String schema) { + return schema == null ? "" : schema.trim(); + } + + private static List copyTables(List tables) { + List 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 tables) { diff --git a/agents/drivers/tdengine/src/test/java/com/dbx/agent/tdengine/TDengineAgentTest.java b/agents/drivers/tdengine/src/test/java/com/dbx/agent/tdengine/TDengineAgentTest.java index df770f666..6e2072d81 100644 --- a/agents/drivers/tdengine/src/test/java/com/dbx/agent/tdengine/TDengineAgentTest.java +++ b/agents/drivers/tdengine/src/test/java/com/dbx/agent/tdengine/TDengineAgentTest.java @@ -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 statements = new ArrayList<>(); + List maxRows = new ArrayList<>(); TDengineAgent agent = new TDengineAgent(); - TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows)); + + List firstPage = agent.listTables( + "dbx_alpha", + new MetadataListConstraints(null, 2, null, List.of("TABLE")) + ); + List 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 statements = new ArrayList<>(); + List maxRows = new ArrayList<>(); + TDengineAgent agent = new TDengineAgent(); + TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows)); + + List 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 statements = new ArrayList<>(); + List 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 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 statements = new ArrayList<>(); + List maxRows = new ArrayList<>(); TDengineAgent agent = new TDengineAgent(); - TestSupport.setPrivateConnection(agent, pagedMetadataConnection(metadata, statements)); + TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows)); - List 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 tables = agent.listTables("dbx_alpha"); - Assertions.assertEquals(metadata.stream().map(TableInfo::getName).toList(), combined.stream().map(TableInfo::getName).toList()); - Set 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 statements = new ArrayList<>(); + List maxRows = new ArrayList<>(); + TDengineAgent agent = new TDengineAgent(); + TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows)); + + List 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 metadata, List statements) { + private static Connection showMetadataConnection(List statements, List 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 metadata) { - return proxy(PreparedStatement.class, (proxy, method, args) -> { + private static java.sql.Statement showMetadataStatement(List statements, List 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 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 rows, int maxRows) { + List 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 proxy(Class type, InvocationHandler handler) { diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index ff43065d3..7673c8a18 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -15,6 +15,13 @@ use crate::models::connection::DatabaseConnectionInfo; type PendingAgentResponse = tokio::sync::oneshot::Sender>; +#[derive(Clone)] +struct CachedAgentQuery { + key: String, + expires_at: Instant, + result: Result, +} + pub struct AgentRuntimeClient { child: Arc>, stdin: Arc>>, @@ -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(response: Value) -> Result(result: Result) -> Result { + 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>>, 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>, agent_session_id: Option, + cached_query: Option, } #[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 { + 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(&mut self, params: Value) -> Result { + self.invalidate_cached_query(); self.call_method(AgentMethod::ExecuteQuery, params).await } @@ -1300,6 +1333,7 @@ impl AgentDriverClient { params: Value, timeout_duration: Option, ) -> Result { + self.invalidate_cached_query(); self.call_method_with_timeout(AgentMethod::ExecuteQuery, params, timeout_duration).await } @@ -1309,14 +1343,37 @@ impl AgentDriverClient { timeout_duration: Option, cancel_token: Option, ) -> Result { + 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( + &mut self, + cache_key: String, + cache_ttl: Duration, + params: Value, + timeout_duration: Option, + ) -> Result { + 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::(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( &mut self, params: Value, ) -> Result { + self.invalidate_cached_query(); self.call_method(AgentMethod::ExecuteQueryPage, params).await } @@ -1325,6 +1382,7 @@ impl AgentDriverClient { params: Value, timeout_duration: Option, ) -> Result { + self.invalidate_cached_query(); self.call_method_with_timeout(AgentMethod::ExecuteQueryPage, params, timeout_duration).await } @@ -1334,6 +1392,7 @@ impl AgentDriverClient { timeout_duration: Option, cancel_token: Option, ) -> Result { + 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 { + 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 { + 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, ) -> Result { + 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( &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, 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, 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::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::("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::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::("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::( + "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::( + "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::( + "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::( + "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::( + "query-invalidation".to_string(), + cache_ttl, + serde_json::json!({"sql": "cached-before-query"}), + timeout, + ) + .await + .unwrap(); + client.execute_query::(serde_json::json!({"sql": "ordinary"})).await.unwrap(); + client + .execute_query_cached_with_timeout::( + "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::( + "batch-invalidation".to_string(), + cache_ttl, + serde_json::json!({"sql": "cached-before-batch"}), + timeout, + ) + .await + .unwrap(); + client + .execute_batch::(None, &["UPDATE t SET a = 1".to_string()], None, timeout) + .await + .unwrap(); + client + .execute_query_cached_with_timeout::( + "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::( + "transaction-invalidation".to_string(), + cache_ttl, + serde_json::json!({"sql": "cached-before-transaction"}), + timeout, + ) + .await + .unwrap(); + client.execute_transaction::(None, &["UPDATE t SET a = 2".to_string()], None).await.unwrap(); + client + .execute_query_cached_with_timeout::( + "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::( + "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())); diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index eb8b4b9ba..a597b2605 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -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::>(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::( + 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, String> { Ok(result .rows @@ -1524,7 +1595,7 @@ fn oracle_table_comments_sql(schema: &str, table_names: &[String]) -> Option HashMap { +fn table_comments_from_query_result(result: db::QueryResult) -> HashMap { result .rows .into_iter() @@ -1953,7 +2024,7 @@ fn unique_oracle_comment_names<'a>(names: impl Iterator) -> Vec< unique } -fn apply_oracle_table_comments(tables: &mut [db::TableInfo], comments: &HashMap) { +fn apply_table_comments(tables: &mut [db::TableInfo], comments: &HashMap) { 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::( + 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"));