From 7c5e9dcda006e473664ee0c824978cdc4e3c5d3d Mon Sep 17 00:00:00 2001 From: t8y2 Date: Fri, 9 Oct 2026 00:16:48 +0800 Subject: [PATCH] fix(cache): page native table metadata without JDBC catalog scans Closes #8341 --- .../java/com/dbx/agent/cache/CacheAgent.java | 84 ++++++ .../dbx/agent/cache/CacheAgentTablesTest.java | 258 ++++++++++++++++++ crates/dbx-core/src/schema/mod.rs | 12 +- 3 files changed, 352 insertions(+), 2 deletions(-) create mode 100644 agents/drivers/cache/src/test/java/com/dbx/agent/cache/CacheAgentTablesTest.java diff --git a/agents/drivers/cache/src/main/java/com/dbx/agent/cache/CacheAgent.java b/agents/drivers/cache/src/main/java/com/dbx/agent/cache/CacheAgent.java index 0ba44afd64..19900f54a1 100644 --- a/agents/drivers/cache/src/main/java/com/dbx/agent/cache/CacheAgent.java +++ b/agents/drivers/cache/src/main/java/com/dbx/agent/cache/CacheAgent.java @@ -5,15 +5,19 @@ import com.dbx.agent.ColumnInfo; import com.dbx.agent.ConnectParams; import com.dbx.agent.JdbcAgentProfile; import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.MultiSessionJsonRpcServer; import com.dbx.agent.StandardJdbcMetadata; +import com.dbx.agent.TableInfo; import java.sql.Connection; import java.sql.DatabaseMetaData; import java.sql.PreparedStatement; import java.sql.ResultSet; +import java.sql.SQLException; import java.sql.Types; import java.util.ArrayList; +import java.util.Collections; import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.List; @@ -37,6 +41,7 @@ public final class CacheAgent extends ConfiguredJdbcAgent { @Override protected void afterConnect(ConnectParams params, Connection connection) { + super.afterConnect(params, connection); configuredDatabase = params.getDatabase(); } @@ -74,6 +79,85 @@ public final class CacheAgent extends ConfiguredJdbcAgent { return dedupeCaseInsensitiveSchemas(StandardJdbcMetadata.INSTANCE.listSchemas(requireConnection(), CACHE_PROFILE)); } + @Override + public List listTables(String schema) { + return listTables(schema, MetadataListConstraints.NONE); + } + + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + List tableTypes = new ArrayList<>(); + for (String tableType : CACHE_PROFILE.getTableTypes()) { + if (normalized.tableTypeAllowed(tableType)) { + tableTypes.add(tableType); + } + } + if (tableTypes.isEmpty()) { + return Collections.emptyList(); + } + return unchecked(() -> { + Connection connection = requireConnection(); + try { + return readTables(connection, schema, normalized, tableTypes); + } catch (SQLException error) { + if (error.getErrorCode() != -30 && !"42S02".equals(error.getSQLState())) { + throw error; + } + return StandardJdbcMetadata.INSTANCE.listTables(connection, CACHE_PROFILE, configuredDatabase, schema, normalized); + } + }); + } + + private static List readTables( + Connection connection, + String schema, + MetadataListConstraints constraints, + List tableTypes + ) throws SQLException { + int offset = constraints.hasOffset() ? constraints.getOffset() : 0; + int limit = constraints.hasLimit() ? constraints.getLimit() : Integer.MAX_VALUE; + long fetchLimit = (long) offset + limit; + StringBuilder sql = new StringBuilder("SELECT "); + if (constraints.hasLimit() && !constraints.hasFilter() && fetchLimit <= Integer.MAX_VALUE) { + sql.append("TOP ").append(fetchLimit).append(' '); + } + sql.append("TABLE_NAME, TABLE_TYPE, DESCRIPTION FROM INFORMATION_SCHEMA.TABLES WHERE "); + if (!isBlank(schema)) { + sql.append("TABLE_SCHEMA = ? AND "); + } + sql.append("TABLE_TYPE IN (").append(String.join(", ", Collections.nCopies(tableTypes.size(), "?"))) + .append(") ORDER BY TABLE_NAME, TABLE_SCHEMA"); + + List tables = new ArrayList<>(); + try (PreparedStatement statement = connection.prepareStatement(sql.toString())) { + int parameter = 1; + if (!isBlank(schema)) { + statement.setString(parameter++, schema.trim()); + } + for (String tableType : tableTypes) { + statement.setString(parameter++, tableType); + } + try (ResultSet result = statement.executeQuery()) { + int skipped = 0; + while (tables.size() < limit && result.next()) { + String name = result.getString("TABLE_NAME"); + String comment = result.getString("DESCRIPTION"); + if (!constraints.nameOrCommentMatches(name, comment)) { + continue; + } + if (skipped < offset) { + skipped++; + continue; + } + String tableType = result.getString("TABLE_TYPE"); + tables.add(new TableInfo(name, "BASE TABLE".equalsIgnoreCase(tableType) ? "TABLE" : tableType, comment)); + } + } + } + return tables; + } + @Override public List getColumns(String schema, String table) { return cacheColumns(requireConnection(), getProfile(), configuredDatabase, schema, table); diff --git a/agents/drivers/cache/src/test/java/com/dbx/agent/cache/CacheAgentTablesTest.java b/agents/drivers/cache/src/test/java/com/dbx/agent/cache/CacheAgentTablesTest.java new file mode 100644 index 0000000000..816c3c2eec --- /dev/null +++ b/agents/drivers/cache/src/test/java/com/dbx/agent/cache/CacheAgentTablesTest.java @@ -0,0 +1,258 @@ +package com.dbx.agent.cache; + +import com.dbx.agent.AbstractJdbcAgent; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.TableInfo; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Field; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class CacheAgentTablesTest { + @Test + void bindsSpecialSchemaNamesLiterallyWithoutJdbcMetadata() throws Exception { + for (String schema : List.of("_system", "%Atelier", "SQLUser", "a'b")) { + Fixture fixture = new Fixture(List.of(table("Example", "BASE TABLE", "description"))); + List tables = fixture.agent.listTables(schema); + + assertEquals("Example", tables.get(0).getName()); + assertEquals("TABLE", tables.get(0).getTable_type()); + assertEquals("description", tables.get(0).getComment()); + assertTrue(fixture.sql.contains("FROM INFORMATION_SCHEMA.TABLES")); + assertTrue(fixture.sql.contains("TABLE_SCHEMA = ?")); + assertEquals(schema, fixture.bindings.get(0)); + assertFalse(fixture.sql.contains(schema)); + assertTrue(fixture.statementClosed); + assertTrue(fixture.resultClosed); + } + } + + @Test + void stopsReadingMetadataAtTheRequestedPage() throws Exception { + List> rows = new ArrayList<>(); + for (int index = 0; index < 5000; index++) { + rows.add(table(String.format("Table%05d", index), "TABLE", "")); + } + Fixture fixture = new Fixture(rows); + + List tables = fixture.agent.listTables("_system", constraints(null, 2, 3, "TABLE")); + + assertEquals(2, tables.size()); + assertEquals("Table00003", tables.get(0).getName()); + assertEquals("Table00004", tables.get(1).getName()); + assertEquals(5, fixture.rowsRead); + assertTrue(fixture.sql.startsWith("SELECT TOP 5 ")); + assertTrue(fixture.sql.endsWith("ORDER BY TABLE_NAME, TABLE_SCHEMA")); + assertEquals(List.of("_system", "TABLE", "BASE TABLE", "SYSTEM TABLE"), fixture.bindings); + } + + @Test + void filtersViewsBeforeApplyingThePageLimit() throws Exception { + Fixture fixture = new Fixture(List.of(table("Report", "VIEW", ""))); + + assertEquals("VIEW", fixture.agent.listTables("SQLUser", constraints(null, 1, 0, "VIEW")).get(0).getTable_type()); + assertEquals(List.of("SQLUser", "VIEW", "SYSTEM VIEW"), fixture.bindings); + } + + @Test + void routesObjectTypeOverloadThroughTheNativeQuery() throws Exception { + Fixture fixture = new Fixture(List.of(table("Report", "SYSTEM VIEW", ""))); + + assertEquals("SYSTEM VIEW", fixture.agent.listTables("SQLUser", List.of("VIEW")).get(0).getTable_type()); + assertEquals(List.of("SQLUser", "VIEW", "SYSTEM VIEW"), fixture.bindings); + } + + @Test + void handlesOffsetBeyondTheLastTable() throws Exception { + Fixture fixture = new Fixture(List.of(table("Example", "TABLE", ""))); + + assertTrue(fixture.agent.listTables("SQLUser", constraints(null, 2, 5, "TABLE")).isEmpty()); + assertEquals(1, fixture.rowsRead); + } + + @Test + void avoidsOverflowInTheServerFetchLimit() throws Exception { + Fixture fixture = new Fixture(List.of()); + + assertTrue(fixture.agent.listTables("SQLUser", constraints(null, 2, Integer.MAX_VALUE, "TABLE")).isEmpty()); + assertFalse(fixture.sql.contains("TOP")); + } + + @Test + void preservesFuzzyNameAndCommentSearchBeforeOffset() throws Exception { + Fixture fixture = new Fixture(List.of( + table("Alpha", "TABLE", ""), + table("Budget", "TABLE", "monthly report"), + table("Other", "TABLE", ""), + table("MonthlyReport", "TABLE", ""), + table("Unused", "TABLE", "") + )); + + List tables = fixture.agent.listTables("SQLUser", constraints("mrt", 1, 1, "TABLE")); + + assertEquals("MonthlyReport", tables.get(0).getName()); + assertEquals(4, fixture.rowsRead); + assertFalse(fixture.sql.contains("TOP")); + } + + @Test + void supportsUnscopedAndUnpagedLists() throws Exception { + Fixture fixture = new Fixture(List.of(table("Alpha", "TABLE", ""), table("Beta", "SYSTEM VIEW", ""))); + + assertEquals(2, fixture.agent.listTables(null).size()); + assertFalse(fixture.sql.contains("TABLE_SCHEMA = ?")); + assertFalse(fixture.sql.contains("TOP")); + } + + @Test + void unsupportedObjectTypesDoNotContactTheDatabase() throws Exception { + Fixture fixture = new Fixture(List.of()); + + assertTrue(fixture.agent.listTables("SQLUser", constraints(null, 1, 0, "PROCEDURE")).isEmpty()); + assertEquals("", fixture.sql); + } + + @Test + void doesNotFallbackForAnEmptySchema() throws Exception { + Fixture fixture = new Fixture(List.of()); + + assertTrue(fixture.agent.listTables("Empty").isEmpty()); + assertTrue(fixture.statementClosed); + assertTrue(fixture.resultClosed); + } + + @Test + void fallsBackToJdbcForServersWithoutInformationSchema() throws Exception { + Fixture fixture = new Fixture(List.of(table("Legacy", "TABLE", ""))); + fixture.queryFailure = new SQLException("Table INFORMATION_SCHEMA.TABLES not found", "42S02", -30); + fixture.allowJdbc = true; + + assertEquals("Legacy", fixture.agent.listTables("SQLUser", constraints(null, 1, 0, "TABLE")).get(0).getName()); + assertTrue(fixture.jdbcUsed); + assertTrue(fixture.statementClosed); + } + + @Test + void preservesConfiguredNamespaceForInheritedMetadata() throws Exception { + Fixture fixture = new Fixture(List.of()); + fixture.allowJdbc = true; + + assertEquals("USER", fixture.agent.listDatabases().get(0).getName()); + } + + @Test + void closesResourcesWhenReadingFails() throws Exception { + Fixture fixture = new Fixture(List.of()); + fixture.readFailure = new SQLException("connection lost", "08006"); + + assertThrows(RuntimeException.class, () -> fixture.agent.listTables("SQLUser")); + assertTrue(fixture.statementClosed); + assertTrue(fixture.resultClosed); + assertFalse(fixture.jdbcUsed); + } + + @Test + void doesNotRetryPermissionOrConnectionErrorsThroughLegacyMetadata() throws Exception { + for (String state : List.of("42000", "08006")) { + Fixture fixture = new Fixture(List.of()); + fixture.queryFailure = new SQLException("metadata query failed", state); + + assertThrows(RuntimeException.class, () -> fixture.agent.listTables("SQLUser")); + assertTrue(fixture.statementClosed); + assertFalse(fixture.jdbcUsed); + } + } + + private static MetadataListConstraints constraints(String filter, int limit, int offset, String type) { + return new MetadataListConstraints(filter, limit, offset, List.of(type)); + } + + private static Map table(String name, String type, String comment) { + return Map.of("TABLE_NAME", name, "TABLE_TYPE", type, "DESCRIPTION", comment, "REMARKS", comment); + } + + private static Value proxy(Class type, InvocationHandler handler) { + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, handler)); + } + + private static final class Fixture { + private final CacheAgent agent = new CacheAgent(); + private final List bindings = new ArrayList<>(); + private String sql = ""; + private int rowsRead; + private boolean statementClosed; + private boolean resultClosed; + private boolean allowJdbc; + private boolean jdbcUsed; + private SQLException queryFailure; + private SQLException readFailure; + + private Fixture(List> values) throws Exception { + Connection connection = proxy(Connection.class, (instance, method, arguments) -> { + switch (method.getName()) { + case "prepareStatement": + sql = (String) arguments[0]; + return proxy(PreparedStatement.class, (statement, operation, parameters) -> { + switch (operation.getName()) { + case "setString": bindings.add((String) parameters[1]); return null; + case "executeQuery": + if (queryFailure != null) throw queryFailure; + return rows(values); + case "close": statementClosed = true; return null; + default: throw new AssertionError(operation.getName()); + } + }); + case "getMetaData": + if (!allowJdbc) throw new AssertionError("unexpected legacy JDBC metadata query"); + jdbcUsed = true; + return proxy(DatabaseMetaData.class, (metadata, operation, parameters) -> { + switch (operation.getName()) { + case "getTables": return rows(values); + case "getTableTypes": return rows(List.of(Map.of("TABLE_TYPE", "TABLE"))); + case "getCatalogs": return rows(List.of()); + case "getSearchStringEscape": return "\\"; + default: throw new AssertionError(operation.getName()); + } + }); + case "getCatalog": return null; + default: throw new AssertionError(method.getName()); + } + }); + Field connectionField = AbstractJdbcAgent.class.getDeclaredField("connection"); + connectionField.setAccessible(true); + connectionField.set(agent, connection); + agent.afterConnect(new ConnectParams("db.local", 1972, "USER", "", "", "", "", false), connection); + } + + private ResultSet rows(List> values) { + int[] position = {-1}; + return proxy(ResultSet.class, (instance, method, arguments) -> { + switch (method.getName()) { + case "next": + if (readFailure != null) throw readFailure; + boolean hasNext = ++position[0] < values.size(); + if (hasNext) rowsRead++; + return hasNext; + case "getString": return values.get(position[0]).get((String) arguments[0]); + case "close": resultClosed = true; return null; + default: throw new AssertionError(method.getName()); + } + }); + } + } +} diff --git a/crates/dbx-core/src/schema/mod.rs b/crates/dbx-core/src/schema/mod.rs index 98a1593cb3..78d3431e01 100644 --- a/crates/dbx-core/src/schema/mod.rs +++ b/crates/dbx-core/src/schema/mod.rs @@ -4845,7 +4845,13 @@ for line in sys.stdin: } #[test] - fn agent_table_paging_supports_tdengine_and_default_oracle_only() { + fn agent_table_paging_supports_cache_tdengine_and_default_oracle() { + let mut cache = test_connection_config(DatabaseType::Iris); + assert!(!super::supports_agent_table_paging(&cache)); + cache.driver_profile = Some("cache".to_string()); + assert!(super::supports_agent_table_paging(&cache)); + cache.driver_profile = Some("CACHE".to_string()); + assert!(super::supports_agent_table_paging(&cache)); assert!(super::supports_agent_table_paging(&test_connection_config(DatabaseType::Tdengine))); assert!(super::supports_agent_table_paging(&test_connection_config(DatabaseType::Oracle))); assert!(!super::supports_agent_table_paging(&test_connection_config(DatabaseType::Dameng))); @@ -9743,7 +9749,9 @@ fn uses_oracle_metadata_object_source(config: Option<&ConnectionConfig>, object_ fn supports_agent_table_paging(config: &ConnectionConfig) -> bool { // Keep paging opt-in until each legacy agent is known to apply metadata constraints server-side. - matches!(config.db_type, DatabaseType::Tdengine) || is_default_oracle_agent_config(config) + matches!(config.db_type, DatabaseType::Tdengine) + || crate::agent_catalog::agent_key(&config.db_type, config.driver_profile.as_deref()) == Some("cache") + || is_default_oracle_agent_config(config) } fn agent_paging_likely_applied(enabled: bool, limit: Option, returned_len: usize) -> bool {