mirror of
https://github.com/t8y2/dbx.git
synced 2026-10-21 03:21:53 +00:00
fix(cache): page native table metadata without JDBC catalog scans
Closes #8341
This commit is contained in:
@@ -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<TableInfo> listTables(String schema) {
|
||||
return listTables(schema, MetadataListConstraints.NONE);
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<TableInfo> listTables(String schema, MetadataListConstraints constraints) {
|
||||
MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints);
|
||||
List<String> 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<TableInfo> readTables(
|
||||
Connection connection,
|
||||
String schema,
|
||||
MetadataListConstraints constraints,
|
||||
List<String> 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<TableInfo> 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<ColumnInfo> getColumns(String schema, String table) {
|
||||
return cacheColumns(requireConnection(), getProfile(), configuredDatabase, schema, table);
|
||||
|
||||
+258
@@ -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<TableInfo> 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<Map<String, String>> 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<TableInfo> 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<TableInfo> 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<String, String> table(String name, String type, String comment) {
|
||||
return Map.of("TABLE_NAME", name, "TABLE_TYPE", type, "DESCRIPTION", comment, "REMARKS", comment);
|
||||
}
|
||||
|
||||
private static <Value> Value proxy(Class<Value> 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<String> 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<Map<String, String>> 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<Map<String, String>> 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());
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<usize>, returned_len: usize) -> bool {
|
||||
|
||||
Reference in New Issue
Block a user