From e74782f6511cf9f199be7a0463ee4d56c97e827c Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Sun, 5 Jul 2026 12:09:24 +0800 Subject: [PATCH] feat(metadata): optimize agent metadata loading --- .../com/dbx/agent/ConfiguredJdbcAgent.java | 24 + .../java/com/dbx/agent/DatabaseAgent.java | 10 + .../java/com/dbx/agent/JsonRpcServer.java | 13 +- .../dbx/agent/MetadataListConstraints.java | 211 +++ .../com/dbx/agent/MetadataSqlSupport.java | 95 ++ .../com/dbx/agent/StandardJdbcMetadata.java | 55 +- .../agent/CommonJavaCompatibilityTest.java | 114 ++ .../agent/MetadataConstraintCoverageTest.java | 114 ++ .../com/dbx/agent/bigquery/BigQueryAgent.java | 91 +- .../bigquery/BigQueryAgentMetadataTest.java | 16 + .../com/dbx/agent/dameng/DamengAgent.java | 229 +++- .../com/dbx/agent/dameng/DamengAgentTest.java | 45 + .../com/dbx/agent/databend/DatabendAgent.java | 66 + .../dbx/agent/databend/DatabendAgentTest.java | 142 ++ .../main/java/com/dbx/agent/db2/Db2Agent.java | 130 ++ .../java/com/dbx/agent/db2/Db2AgentTest.java | 32 + .../com/dbx/agent/gbase8a/Gbase8aAgent.java | 216 ++- .../dbx/agent/gbase8a/Gbase8aAgentTest.java | 51 +- .../com/dbx/agent/gbase8s/Gbase8sAgent.java | 87 ++ .../dbx/agent/gbase8s/Gbase8sAgentTest.java | 116 ++ .../com/dbx/agent/goldendb/GoldendbAgent.java | 146 +++ .../goldendb/GoldendbAgentMetadataTest.java | 29 + .../main/java/com/dbx/agent/h2/H2Agent.java | 148 +++ .../java/com/dbx/agent/h2/H2AgentTest.java | 19 + .../com/dbx/agent/informix/InformixAgent.java | 127 ++ .../dbx/agent/informix/InformixAgentTest.java | 33 + .../com/dbx/agent/kingbase/KingbaseAgent.java | 231 ++++ .../dbx/agent/kingbase/KingbaseAgentTest.java | 53 + .../oceanbaseoracle/OceanBaseOracleAgent.java | 157 ++- .../OceanBaseOracleAgentTest.java | 140 ++ agents/drivers/oracle-go/main.go | 175 ++- agents/drivers/oracle-go/main_test.go | 62 + .../dbx/agent/snowflake/SnowflakeAgent.java | 140 +- .../snowflake/SnowflakeAgentMetadataTest.java | 30 + .../java/com/dbx/agent/sundb/SundbAgent.java | 146 +++ .../agent/sundb/SundbAgentMetadataTest.java | 30 + .../java/com/dbx/agent/trino/TrinoAgent.java | 73 ++ .../com/dbx/agent/trino/TrinoAgentTest.java | 21 + agents/drivers/xugu/main.go | 215 ++- agents/drivers/xugu/main_test.go | 67 + agents/metadata-constraint-coverage.tsv | 23 + .../dbx/agent/test/JdbcMetadataSqlFake.java | 2 +- .../generate/DataGenerateDialog.vue | 3 + .../components/import/TableImportDialog.vue | 1 + .../src/components/sidebar/TreeItem.vue | 59 +- apps/desktop/src/lib/backend/http.ts | 5 +- apps/desktop/src/lib/backend/tauri.ts | 4 +- .../lib/metadata/metadataLoadCoordinator.ts | 145 +++ .../src/lib/metadata/metadataLoadScope.ts | 86 ++ .../src/lib/metadata/metadataResultCache.ts | 99 ++ .../src/lib/metadata/tableMetadataCache.ts | 126 ++ .../lib/sidebar/sidebarTableSearchControl.ts | 11 +- .../connectionStore.metadataLoading.spec.ts | 99 ++ apps/desktop/src/stores/connectionStore.ts | 1159 +++++++++++------ apps/desktop/src/stores/queryStore.ts | 27 +- crates/dbx-core/src/database_export.rs | 13 +- crates/dbx-core/src/db/agent_driver.rs | 51 +- crates/dbx-core/src/schema.rs | 233 +++- .../app-tests/metadataLoadCoordinator.test.ts | 169 +++ .../app-tests/metadataResultCache.test.ts | 54 + .../sidebarTableSearchControl.test.ts | 65 +- packages/app-tests/tableMetadataCache.test.ts | 101 ++ .../main/java/app/dbx/jdbc/DbxJdbcPlugin.java | 239 +++- .../java/app/dbx/jdbc/DbxJdbcPluginTest.java | 74 ++ src-tauri/src/commands/schema.rs | 16 +- 65 files changed, 6192 insertions(+), 571 deletions(-) create mode 100644 agents/common/src/main/java/com/dbx/agent/MetadataListConstraints.java create mode 100644 agents/common/src/main/java/com/dbx/agent/MetadataSqlSupport.java create mode 100644 agents/common/src/test/java/com/dbx/agent/MetadataConstraintCoverageTest.java create mode 100644 agents/metadata-constraint-coverage.tsv create mode 100644 apps/desktop/src/lib/metadata/metadataLoadCoordinator.ts create mode 100644 apps/desktop/src/lib/metadata/metadataLoadScope.ts create mode 100644 apps/desktop/src/lib/metadata/metadataResultCache.ts create mode 100644 apps/desktop/src/lib/metadata/tableMetadataCache.ts create mode 100644 apps/desktop/src/stores/__tests__/connectionStore.metadataLoading.spec.ts create mode 100644 packages/app-tests/metadataLoadCoordinator.test.ts create mode 100644 packages/app-tests/metadataResultCache.test.ts create mode 100644 packages/app-tests/tableMetadataCache.test.ts diff --git a/agents/common/src/main/java/com/dbx/agent/ConfiguredJdbcAgent.java b/agents/common/src/main/java/com/dbx/agent/ConfiguredJdbcAgent.java index 7bc7d55de2..a92d1ce1bb 100644 --- a/agents/common/src/main/java/com/dbx/agent/ConfiguredJdbcAgent.java +++ b/agents/common/src/main/java/com/dbx/agent/ConfiguredJdbcAgent.java @@ -55,11 +55,35 @@ public abstract class ConfiguredJdbcAgent extends AbstractJdbcAgent { return StandardJdbcMetadata.INSTANCE.listTables(requireConnection(), profile, configuredDatabase, schema); } + @Override + public List listTables(String schema, List objectTypes) { + return listTables(schema, new MetadataListConstraints(null, null, null, objectTypes)); + } + + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + return StandardJdbcMetadata.INSTANCE.listTables( + requireConnection(), + profile, + configuredDatabase, + schema, + constraints + ); + } + @Override public List listObjects(String schema) { return StandardJdbcMetadata.INSTANCE.listObjects(listTables(schema), schema); } + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + MetadataListConstraints tableConstraints = + new MetadataListConstraints(normalized.getFilter(), null, null, normalized.getObjectTypes()); + return StandardJdbcMetadata.INSTANCE.listObjects(listTables(schema, tableConstraints), schema, normalized); + } + @Override public List listDataTypes() { return StandardJdbcMetadata.INSTANCE.listDataTypes(requireConnection()); diff --git a/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java b/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java index f7eb84a149..5e1b9a852a 100644 --- a/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java +++ b/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java @@ -31,6 +31,12 @@ public interface DatabaseAgent { return listTables(schema); } + default List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + // Keep old driver overrides for object type filtering, then apply the new optional constraints locally. + return normalized.filterTables(listTables(schema, normalized.getObjectTypes())); + } + default List listObjects(String schema) { List result = new ArrayList<>(); for (TableInfo table : listTables(schema)) { @@ -39,6 +45,10 @@ public interface DatabaseAgent { return result; } + default List listObjects(String schema, MetadataListConstraints constraints) { + return MetadataListConstraints.orNone(constraints).filterObjects(listObjects(schema)); + } + default List listDataTypes() { return Collections.emptyList(); } diff --git a/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java b/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java index d7c348a219..b978a90d33 100644 --- a/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java +++ b/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java @@ -124,11 +124,11 @@ public final class JsonRpcServer { } if (AgentProtocol.METHOD_LIST_TABLES.equals(method)) { switchCatalog(params); - return agent.listTables(params.get("schema").getAsString(), stringListOrNull(params, "object_types")); + return agent.listTables(params.get("schema").getAsString(), metadataListConstraints(params)); } if (AgentProtocol.METHOD_LIST_OBJECTS.equals(method)) { switchCatalog(params); - return agent.listObjects(params.get("schema").getAsString()); + return agent.listObjects(params.get("schema").getAsString(), metadataListConstraints(params)); } if (AgentProtocol.METHOD_LIST_DATA_TYPES.equals(method)) { switchCatalog(params); @@ -356,6 +356,15 @@ public final class JsonRpcServer { return element.getAsInt(); } + private MetadataListConstraints metadataListConstraints(JsonObject params) { + return new MetadataListConstraints( + stringOrNull(params, "filter"), + intOrNull(params, "limit"), + intOrNull(params, "offset"), + stringListOrNull(params, "object_types") + ); + } + private static int intOrDefault(JsonObject object, String key, int defaultValue) { Integer value = intOrNull(object, key); return value == null ? defaultValue : value; diff --git a/agents/common/src/main/java/com/dbx/agent/MetadataListConstraints.java b/agents/common/src/main/java/com/dbx/agent/MetadataListConstraints.java new file mode 100644 index 0000000000..a21fb106b9 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/MetadataListConstraints.java @@ -0,0 +1,211 @@ +package com.dbx.agent; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Locale; + +public final class MetadataListConstraints { + public static final MetadataListConstraints NONE = new MetadataListConstraints(null, null, null, null); + + private final String filter; + private final Integer limit; + private final Integer offset; + private final List objectTypes; + + public MetadataListConstraints(String filter, Integer limit, Integer offset, List objectTypes) { + this.filter = normalizeFilter(filter); + this.limit = limit == null || limit <= 0 ? null : limit; + this.offset = offset == null || offset <= 0 ? null : offset; + this.objectTypes = normalizeObjectTypes(objectTypes); + } + + public static MetadataListConstraints orNone(MetadataListConstraints constraints) { + return constraints == null ? NONE : constraints; + } + + public String getFilter() { + return filter; + } + + public Integer getLimit() { + return limit; + } + + public Integer getOffset() { + return offset; + } + + public List getObjectTypes() { + return objectTypes.isEmpty() ? null : objectTypes; + } + + public boolean hasFilter() { + return !filter.isEmpty(); + } + + public boolean hasLimit() { + return limit != null; + } + + public boolean hasOffset() { + return offset != null; + } + + public boolean hasObjectTypes() { + return !objectTypes.isEmpty(); + } + + public MetadataListConstraints withoutPaging() { + return new MetadataListConstraints(filter, null, null, objectTypes); + } + + public String fuzzyLikePattern() { + if (filter.isEmpty()) { + return "%%"; + } + StringBuilder builder = new StringBuilder(filter.length() * 2 + 2); + builder.append('%'); + for (int i = 0; i < filter.length(); i++) { + char ch = filter.charAt(i); + if (ch == '\\' || ch == '%' || ch == '_') { + builder.append('\\'); + } + builder.append(ch); + builder.append('%'); + } + return builder.toString(); + } + + public boolean includesTableLikeTypes() { + if (objectTypes.isEmpty()) { + return true; + } + return objectTypes.contains("TABLE") + || objectTypes.contains("VIEW") + || objectTypes.contains("MATERIALIZED_VIEW"); + } + + public boolean tableTypeAllowed(String tableType) { + if (objectTypes.isEmpty()) { + return true; + } + return objectTypes.contains(normalizeTableType(tableType)); + } + + public boolean objectTypeAllowed(String objectType) { + if (objectTypes.isEmpty()) { + return true; + } + String normalized = normalizeObjectType(objectType); + return !normalized.isEmpty() && objectTypes.contains(normalized); + } + + public boolean nameMatches(String name) { + if (filter.isEmpty()) { + return true; + } + if (name == null) { + return false; + } + String candidate = name.toLowerCase(Locale.ROOT); + return candidate.contains(filter) || (filter.length() >= 2 && fuzzySubsequenceMatches(candidate, filter)); + } + + public List filterTables(List tables) { + List result = new ArrayList<>(); + int skipped = 0; + int max = limit == null ? Integer.MAX_VALUE : limit; + int start = offset == null ? 0 : offset; + for (TableInfo table : tables) { + if (!nameMatches(table.getName()) || !tableTypeAllowed(table.getTable_type())) { + continue; + } + if (skipped++ < start) { + continue; + } + if (result.size() >= max) { + break; + } + result.add(table); + } + return result; + } + + public List filterObjects(List objects) { + List result = new ArrayList<>(); + int skipped = 0; + int max = limit == null ? Integer.MAX_VALUE : limit; + int start = offset == null ? 0 : offset; + for (ObjectInfo object : objects) { + if (!nameMatches(object.getName()) || !objectTypeAllowed(object.getObject_type())) { + continue; + } + if (skipped++ < start) { + continue; + } + if (result.size() >= max) { + break; + } + result.add(object); + } + return result; + } + + private static String normalizeFilter(String value) { + if (value == null) { + return ""; + } + return value.trim().toLowerCase(Locale.ROOT); + } + + private static List normalizeObjectTypes(List values) { + if (values == null || values.isEmpty()) { + return Collections.emptyList(); + } + List result = new ArrayList<>(); + for (String value : values) { + String normalized = normalizeObjectType(value); + if (!normalized.isEmpty() && !result.contains(normalized)) { + result.add(normalized); + } + } + Collections.sort(result); + return Collections.unmodifiableList(result); + } + + private static String normalizeTableType(String value) { + String normalized = normalizeObjectType(value); + return normalized.isEmpty() ? "TABLE" : normalized; + } + + private static String normalizeObjectType(String value) { + if (value == null || value.trim().isEmpty()) { + return ""; + } + String upper = value.trim().toUpperCase(Locale.ROOT).replace(' ', '_'); + if (upper.contains("MATERIALIZED") && upper.contains("VIEW")) { + return "MATERIALIZED_VIEW"; + } + if (upper.equals("BASE_TABLE") || upper.contains("TABLE")) { + return "TABLE"; + } + if (upper.contains("VIEW")) { + return "VIEW"; + } + return upper; + } + + private static boolean fuzzySubsequenceMatches(String candidate, String expected) { + int cursor = 0; + for (int i = 0; i < expected.length(); i++) { + char ch = expected.charAt(i); + cursor = candidate.indexOf(ch, cursor); + if (cursor < 0) { + return false; + } + cursor++; + } + return true; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/MetadataSqlSupport.java b/agents/common/src/main/java/com/dbx/agent/MetadataSqlSupport.java new file mode 100644 index 0000000000..c5c85e0c72 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/MetadataSqlSupport.java @@ -0,0 +1,95 @@ +package com.dbx.agent; + +import java.sql.PreparedStatement; +import java.util.Collections; +import java.util.List; +import java.util.Locale; + +public final class MetadataSqlSupport { + private MetadataSqlSupport() { + } + + public static void appendNameFilter(StringBuilder sql, List args, String column, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (!normalized.hasFilter()) { + return; + } + sql.append(" AND UPPER(").append(column).append(") LIKE ? ESCAPE '\\\\'"); + args.add(normalized.fuzzyLikePattern().toUpperCase(Locale.ROOT)); + } + + public static void appendLimitOffset(StringBuilder sql, List args, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (!normalized.hasLimit()) { + return; + } + sql.append(" LIMIT ?"); + args.add(normalized.getLimit()); + if (normalized.hasOffset()) { + sql.append(" OFFSET ?"); + args.add(normalized.getOffset()); + } + } + + public static void appendOffsetFetch(StringBuilder sql, List args, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (normalized.hasOffset()) { + sql.append(" OFFSET ? ROWS"); + args.add(normalized.getOffset()); + } + if (normalized.hasLimit()) { + sql.append(normalized.hasOffset() ? " FETCH NEXT ? ROWS ONLY" : " FETCH FIRST ? ROWS ONLY"); + args.add(normalized.getLimit()); + } + } + + public static void appendLiteralLimitOffset(StringBuilder sql, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (!normalized.hasLimit()) { + return; + } + sql.append(" LIMIT ").append(normalized.getLimit()); + if (normalized.hasOffset()) { + sql.append(" OFFSET ").append(normalized.getOffset()); + } + } + + public static void appendLiteralOffsetFetch(StringBuilder sql, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (normalized.hasOffset()) { + sql.append(" OFFSET ").append(normalized.getOffset()).append(" ROWS"); + } + if (normalized.hasLimit()) { + sql.append(normalized.hasOffset() ? " FETCH NEXT " : " FETCH FIRST ") + .append(normalized.getLimit()) + .append(" ROWS ONLY"); + } + } + + public static void appendLiteralSkipFirst(StringBuilder sql, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (normalized.hasOffset()) { + sql.append("SKIP ").append(normalized.getOffset()).append(' '); + } + if (normalized.hasLimit()) { + sql.append("FIRST ").append(normalized.getLimit()).append(' '); + } + } + + public static String placeholders(int count) { + return String.join(", ", Collections.nCopies(count, "?")); + } + + public static void bind(PreparedStatement stmt, List args) throws Exception { + for (int index = 0; index < args.size(); index += 1) { + Object arg = args.get(index); + if (arg instanceof Integer) { + stmt.setInt(index + 1, (Integer) arg); + } else if (arg == null) { + stmt.setObject(index + 1, null); + } else { + stmt.setString(index + 1, String.valueOf(arg)); + } + } + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java b/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java index 478dd6b4cb..386e1fd6c7 100644 --- a/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java +++ b/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java @@ -70,24 +70,42 @@ public final class StandardJdbcMetadata { } public List listTables(Connection conn, JdbcAgentProfile profile, String configuredDatabase, String schema) { + return listTables(conn, profile, configuredDatabase, schema, MetadataListConstraints.NONE); + } + + public List listTables( + Connection conn, + JdbcAgentProfile profile, + String configuredDatabase, + String schema, + MetadataListConstraints constraints + ) { return unchecked(() -> { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (!normalized.includesTableLikeTypes()) { + return Collections.emptyList(); + } DatabaseMetaData meta = conn.getMetaData(); List result = new ArrayList<>(); - appendTables(result, meta, profile, null, schema); + appendTables(result, meta, profile, null, schema, normalized); if (result.isEmpty() && profile.getCatalogFallbackEnabled() && !configuredDatabase.trim().isEmpty()) { - appendTables(result, meta, profile, configuredDatabase, schema); + appendTables(result, meta, profile, configuredDatabase, schema, normalized); } result.sort(Comparator.comparing(TableInfo::getName)); - return result; + return normalized.filterTables(result); }); } public List listObjects(List tables, String schema) { + return listObjects(tables, schema, MetadataListConstraints.NONE); + } + + public List listObjects(List tables, String schema, MetadataListConstraints constraints) { List result = new ArrayList<>(); for (TableInfo table : tables) { result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); } - return result; + return MetadataListConstraints.orNone(constraints).filterObjects(result); } public List listDataTypes(Connection conn) { @@ -270,8 +288,19 @@ public final class StandardJdbcMetadata { return profile.schemaSwitchSql(schema, quote); } - private void appendTables(List result, DatabaseMetaData meta, JdbcAgentProfile profile, String catalog, String schema) throws Exception { - String[] tableTypes = getDriverTableTypes(meta, profile); + private void appendTables( + List result, + DatabaseMetaData meta, + JdbcAgentProfile profile, + String catalog, + String schema, + MetadataListConstraints constraints + ) throws Exception { + String[] tableTypes = constrainedTableTypes(getDriverTableTypes(meta, profile), constraints); + if (tableTypes.length == 0) { + return; + } + // JDBC has no portable metadata limit/offset. Use table types for safe pushdown and filter/page locally. try (ResultSet rs = meta.getTables(catalog, blankToNull(schema), "%", tableTypes)) { while (rs.next()) { result.add(new TableInfo( @@ -283,6 +312,20 @@ public final class StandardJdbcMetadata { } } + private static String[] constrainedTableTypes(String[] tableTypes, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (!normalized.hasObjectTypes()) { + return tableTypes; + } + List result = new ArrayList<>(); + for (String tableType : tableTypes) { + if (normalized.tableTypeAllowed(tableType)) { + result.add(tableType); + } + } + return result.toArray(new String[0]); + } + private void appendCompletionSchemas( List result, DatabaseMetaData meta, diff --git a/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java b/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java index 048d1c7b75..e538ad4208 100644 --- a/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java +++ b/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java @@ -146,6 +146,36 @@ class CommonJavaCompatibilityTest { assertEquals("app", agent.lastSchema); } + @Test + void jsonRpcServerAppliesConstrainedTableMetadataRequests() { + MetadataConstraintAgent agent = new MetadataConstraintAgent(); + JsonRpcServer server = new JsonRpcServer(agent); + + String response = server.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"" + AgentProtocol.METHOD_LIST_TABLES + "\",\"params\":{\"schema\":\"app\",\"filter\":\"us\",\"limit\":1,\"offset\":1,\"object_types\":[\"TABLE\"]}}" + ); + + JsonArray result = JsonParser.parseString(response).getAsJsonObject().getAsJsonArray("result"); + assertEquals(1, result.size()); + assertEquals("user_settings", result.get(0).getAsJsonObject().get("name").getAsString()); + assertEquals("app", agent.lastSchema); + } + + @Test + void jsonRpcServerAppliesConstrainedObjectMetadataRequests() { + MetadataConstraintAgent agent = new MetadataConstraintAgent(); + JsonRpcServer server = new JsonRpcServer(agent); + + String response = server.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":5,\"method\":\"" + AgentProtocol.METHOD_LIST_OBJECTS + "\",\"params\":{\"schema\":\"app\",\"filter\":\"fn\",\"limit\":1,\"offset\":1,\"object_types\":[\"FUNCTION\"]}}" + ); + + JsonArray result = JsonParser.parseString(response).getAsJsonObject().getAsJsonArray("result"); + assertEquals(1, result.size()); + assertEquals("fetch_name", result.get(0).getAsJsonObject().get("name").getAsString()); + assertEquals("FUNCTION", result.get(0).getAsJsonObject().get("object_type").getAsString()); + } + @Test void jsonRpcServerDispatchesTableReadSessionMethods() { TableReadDispatchAgent agent = new TableReadDispatchAgent(); @@ -253,6 +283,25 @@ class CommonJavaCompatibilityTest { ); } + @Test + void databaseAgentDefaultConstraintsFilterLegacyMetadataOverrides() { + DatabaseAgent agent = new LegacyObjectTypeAgent(); + + List tables = agent.listTables( + "public", + new MetadataListConstraints("us", 1, 1, Collections.singletonList("TABLE")) + ); + assertEquals(1, tables.size()); + assertEquals("user_settings", tables.get(0).getName()); + + List objects = agent.listObjects( + "public", + new MetadataListConstraints("us", 1, 0, Collections.singletonList("VIEW")) + ); + assertEquals(1, objects.size()); + assertEquals("usage_view", objects.get(0).getName()); + } + @Test void executesTransactionsOneByOneWhenJdbcDriverDoesNotSupportTransactions() { List calls = new ArrayList<>(); @@ -394,6 +443,33 @@ class CommonJavaCompatibilityTest { } } + private static final class LegacyObjectTypeAgent extends MinimalAgent { + @Override + public List listTables(String schema) { + return Arrays.asList( + new TableInfo("orders", "TABLE"), + new TableInfo("usage_view", "VIEW"), + new TableInfo("users", "TABLE"), + new TableInfo("user_settings", "TABLE") + ); + } + + @Override + public List listTables(String schema, List objectTypes) { + List result = listTables(schema); + if (objectTypes == null || objectTypes.isEmpty()) { + return result; + } + List filtered = new ArrayList<>(); + for (TableInfo table : result) { + if (objectTypes.contains(table.getTable_type())) { + filtered.add(table); + } + } + return filtered; + } + } + private static final class PreciseNumberAgent extends MinimalAgent { @Override public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { @@ -472,6 +548,44 @@ class CommonJavaCompatibilityTest { } } + private static final class MetadataConstraintAgent extends MinimalAgent { + private String lastSchema = ""; + + @Override + public Connection getConnection() { + return proxy(Connection.class, (method, args) -> { + if ("isClosed".equals(method.getName())) { + return false; + } + if ("isValid".equals(method.getName())) { + return true; + } + return defaultValue(method.getReturnType()); + }); + } + + @Override + public List listTables(String schema) { + lastSchema = schema; + return Arrays.asList( + new TableInfo("orders", "TABLE"), + new TableInfo("users", "TABLE"), + new TableInfo("usage_view", "VIEW"), + new TableInfo("user_settings", "TABLE") + ); + } + + @Override + public List listObjects(String schema) { + return Arrays.asList( + new ObjectInfo("orders", "TABLE", schema, null), + new ObjectInfo("find_user", "FUNCTION", schema, null), + new ObjectInfo("fetch_name", "FUNCTION", schema, null), + new ObjectInfo("cleanup_user", "PROCEDURE", schema, null) + ); + } + } + private static final class TableReadDispatchAgent extends MinimalAgent { private String lastSql; private String lastSchema; diff --git a/agents/common/src/test/java/com/dbx/agent/MetadataConstraintCoverageTest.java b/agents/common/src/test/java/com/dbx/agent/MetadataConstraintCoverageTest.java new file mode 100644 index 0000000000..be8498017f --- /dev/null +++ b/agents/common/src/test/java/com/dbx/agent/MetadataConstraintCoverageTest.java @@ -0,0 +1,114 @@ +package com.dbx.agent; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.regex.Pattern; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class MetadataConstraintCoverageTest { + private static final Pattern JAVA_METADATA_OVERRIDE = Pattern.compile( + "public\\s+List\\s*<\\s*(?:TableInfo|ObjectInfo)\\s*>\\s+list(?:Tables|Objects)\\s*\\(\\s*String\\s+schema" + ); + private static final Set VALID_STRATEGIES = new HashSet<>(Arrays.asList( + "native-pushdown", + "shared-fallback", + "intentional-fallback" + )); + + @Test + void everyCustomMetadataDriverIsListedInCoverageMatrix() throws Exception { + Path agentsRoot = agentsRoot(); + Map coverage = readCoverageMatrix(agentsRoot.resolve("metadata-constraint-coverage.tsv")); + Set discovered = discoverCustomMetadataDrivers(agentsRoot.resolve("drivers")); + + assertEquals(discovered, coverage.keySet()); + } + + @Test + void coverageMatrixUsesKnownStrategiesAndReasons() throws Exception { + Path agentsRoot = agentsRoot(); + Path matrix = agentsRoot.resolve("metadata-constraint-coverage.tsv"); + for (String line : Files.readAllLines(matrix, StandardCharsets.UTF_8)) { + if (line.startsWith("driver\t") || line.trim().isEmpty()) { + continue; + } + String[] parts = line.split("\t", -1); + assertEquals(4, parts.length, line); + assertTrue(VALID_STRATEGIES.contains(parts[1]), line); + assertFalse(parts[3].trim().isEmpty(), line); + } + } + + private static Set discoverCustomMetadataDrivers(Path driversRoot) throws IOException { + Set result = new HashSet<>(); + try (Stream files = Files.walk(driversRoot)) { + files.filter(Files::isRegularFile).forEach(file -> { + String fileName = file.getFileName().toString(); + if (fileName.endsWith(".java") && hasJavaMetadataOverride(file)) { + result.add(driversRoot.relativize(file).getName(0).toString()); + } else if ("main.go".equals(fileName) && hasGoMetadataDispatcher(file)) { + result.add(driversRoot.relativize(file).getName(0).toString()); + } + }); + } + return result; + } + + private static boolean hasJavaMetadataOverride(Path file) { + try { + return JAVA_METADATA_OVERRIDE.matcher(readUtf8(file)).find(); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + private static boolean hasGoMetadataDispatcher(Path file) { + try { + String source = readUtf8(file); + return source.contains("\"list_tables\"") || source.contains("\"list_objects\""); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + private static Map readCoverageMatrix(Path matrix) throws IOException { + Map result = new HashMap<>(); + for (String line : Files.readAllLines(matrix, StandardCharsets.UTF_8)) { + if (line.startsWith("driver\t") || line.trim().isEmpty()) { + continue; + } + String[] parts = line.split("\t", -1); + result.put(parts[0], parts[1]); + } + return result; + } + + private static String readUtf8(Path file) throws IOException { + return new String(Files.readAllBytes(file), StandardCharsets.UTF_8); + } + + private static Path agentsRoot() { + Path current = Paths.get("").toAbsolutePath(); + while (current != null) { + if (Files.isDirectory(current.resolve("drivers")) && Files.isDirectory(current.resolve("common"))) { + return current; + } + current = current.getParent(); + } + throw new IllegalStateException("Unable to locate agents root"); + } +} diff --git a/agents/drivers/bigquery/src/main/java/com/dbx/agent/bigquery/BigQueryAgent.java b/agents/drivers/bigquery/src/main/java/com/dbx/agent/bigquery/BigQueryAgent.java index 11feaefd29..71d64bfccd 100644 --- a/agents/drivers/bigquery/src/main/java/com/dbx/agent/bigquery/BigQueryAgent.java +++ b/agents/drivers/bigquery/src/main/java/com/dbx/agent/bigquery/BigQueryAgent.java @@ -8,12 +8,15 @@ import com.dbx.agent.ForeignKeyInfo; import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcIdentifiers; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.MetadataSqlSupport; import com.dbx.agent.TableInfo; import com.dbx.agent.TriggerInfo; import java.sql.ResultSet; import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Locale; public final class BigQueryAgent extends AbstractJdbcAgent { @@ -70,6 +73,56 @@ public final class BigQueryAgent extends AbstractJdbcAgent { }); } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + if (!normalized.includesTableLikeTypes()) { + return List.of(); + } + if (hasLikeWildcard(normalized.getFilter())) { + return normalized.filterTables(listTables(schema)); + } + try { + return queryConstrainedTables(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterTables(listTables(schema)); + } + } + + private List queryConstrainedTables(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List args = new ArrayList<>(); + String quotedSchema = JdbcIdentifiers.INSTANCE.backtick(schema); + StringBuilder sql = new StringBuilder("SELECT table_name, table_type FROM ") + .append(quotedSchema) + .append(".INFORMATION_SCHEMA.TABLES WHERE 1 = 1"); + appendBigQueryTableTypePredicate(sql, args, constraints); + if (constraints.hasFilter()) { + sql.append(" AND UPPER(table_name) LIKE ?"); + args.add(constraints.fuzzyLikePattern().toUpperCase(Locale.ROOT)); + } + sql.append(" ORDER BY table_name"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo( + rs.getString("table_name"), + normalizeTableType(rs.getString("table_type")), + null + )); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } + @Override public List getColumns(String schema, String table) { return unchecked(() -> { @@ -136,7 +189,43 @@ public final class BigQueryAgent extends AbstractJdbcAgent { } private static String normalizeTableType(String type) { - return "BASE TABLE".equals(type) ? "TABLE" : type; + if ("BASE TABLE".equals(type)) { + return "TABLE"; + } + if ("MATERIALIZED VIEW".equals(type)) { + return "MATERIALIZED_VIEW"; + } + return type; + } + + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() && !constraints.hasLimit() && !constraints.hasOffset() && !constraints.hasObjectTypes(); + } + + private static void appendBigQueryTableTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List types = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + types.add("BASE TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + types.add("VIEW"); + } + if (constraints.tableTypeAllowed("MATERIALIZED_VIEW")) { + types.add("MATERIALIZED VIEW"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND table_type IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + + private static boolean hasLikeWildcard(String value) { + return value != null && (value.contains("%") || value.contains("_") || value.contains("\\")); } public static void main(String[] args) { diff --git a/agents/drivers/bigquery/src/test/java/com/dbx/agent/bigquery/BigQueryAgentMetadataTest.java b/agents/drivers/bigquery/src/test/java/com/dbx/agent/bigquery/BigQueryAgentMetadataTest.java index d0340b5571..5a6e80de74 100644 --- a/agents/drivers/bigquery/src/test/java/com/dbx/agent/bigquery/BigQueryAgentMetadataTest.java +++ b/agents/drivers/bigquery/src/test/java/com/dbx/agent/bigquery/BigQueryAgentMetadataTest.java @@ -2,9 +2,11 @@ package com.dbx.agent.bigquery; import com.dbx.agent.test.TestSupport; import com.dbx.agent.ConnectParams; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.test.JdbcAgentFake; import com.dbx.agent.test.JdbcMetadataSqlFake; import java.sql.Connection; +import java.util.List; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -46,4 +48,18 @@ class BigQueryAgentMetadataTest { BigQueryAgent.buildUrl(params) ); } + + @Test + void constrainedTableMetadataPushesFilterTypesAndPaging() { + BigQueryAgent agent = new BigQueryAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listTables("demo`dataset", new MetadataListConstraints("ord", 25, 50, List.of("TABLE", "VIEW"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + Assertions.assertTrue(sql.contains("FROM `demo``dataset`.INFORMATION_SCHEMA.TABLES"), sql); + Assertions.assertTrue(sql.contains("table_type IN (?, ?)"), sql); + Assertions.assertTrue(sql.contains("UPPER(table_name) LIKE ?"), sql); + Assertions.assertTrue(sql.endsWith("LIMIT 25 OFFSET 50"), sql); + } } diff --git a/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java b/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java index fbb4eec7c3..10cb9331cd 100644 --- a/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java +++ b/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java @@ -10,6 +10,7 @@ import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcExecutor; import com.dbx.agent.JdbcIdentifiers; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.ObjectInfo; import com.dbx.agent.ObjectSource; import com.dbx.agent.QueryPageOptions; @@ -140,7 +141,9 @@ public final class DamengAgent extends BaseDatabaseAgent { @Override public List listTables(String schema) { - return listTables(schema, null); + // Keep this call on the legacy object-type overload; the common agent + // also exposes a constraints overload for paged metadata listing. + return listTables(schema, (List) null); } @Override @@ -162,6 +165,60 @@ public final class DamengAgent extends BaseDatabaseAgent { }); } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + return queryConstrainedTables(schema, normalized); + } + + private List queryConstrainedTables(String schema, MetadataListConstraints constraints) { + if (!constraints.includesTableLikeTypes()) { + return List.of(); + } + try { + return unchecked(() -> { + MetadataQuery query = buildConstrainedTablesQuery(schema, constraints); + List result = new ArrayList<>(); + try (PreparedStatement stmt = requireConnected().prepareStatement(query.sql())) { + bind(stmt, query.args()); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString("TABLE_NAME"), normalizeObjectType(rs.getString("TABLE_TYPE")), rs.getString("COMMENTS"))); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } catch (RuntimeException e) { + // Restricted Dameng catalog views vary by version; fall back to the + // legacy metadata path so browsing still works when pushdown fails. + return constraints.filterTables(listTables(schema)); + } + } + + static MetadataQuery buildConstrainedTablesQuery(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder(""" + SELECT o.OBJECT_NAME AS TABLE_NAME, + CASE WHEN o.OBJECT_TYPE = 'MATERIALIZED VIEW' THEN 'MATERIALIZED_VIEW' ELSE o.OBJECT_TYPE END AS TABLE_TYPE, + c.COMMENTS + FROM ALL_OBJECTS o + LEFT JOIN ALL_TAB_COMMENTS c ON c.OWNER = o.OWNER AND c.TABLE_NAME = o.OBJECT_NAME + WHERE o.OWNER = ? + """.stripIndent().trim()); + args.add(schema); + appendDamengObjectTypePredicate(sql, args, normalized, true); + sql.append(" AND (o.OBJECT_TYPE <> 'TABLE' OR o.OBJECT_NAME NOT LIKE 'MTAB$_%')"); + appendNameFilter(sql, args, "o.OBJECT_NAME", normalized); + sql.append(" ORDER BY o.OBJECT_NAME"); + appendLimitOffset(sql, args, normalized); + return new MetadataQuery(sql.toString(), args); + } + private void loadTableOrView(String schema, String tableType, Map tablesByName) { if (!loadTableOrViewFromAllObjects(schema, tableType, tablesByName)) { loadTableOrViewFromComments(schema, tableType, tablesByName); @@ -307,6 +364,111 @@ public final class DamengAgent extends BaseDatabaseAgent { return false; } + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() && !constraints.hasLimit() && !constraints.hasOffset() && !constraints.hasObjectTypes(); + } + + private static boolean includesSupportedObjectTypes(MetadataListConstraints constraints) { + return constraints.includesTableLikeTypes() + || constraints.objectTypeAllowed("PROCEDURE") + || constraints.objectTypeAllowed("FUNCTION"); + } + + private static void appendDamengObjectTypePredicate( + StringBuilder sql, + List args, + MetadataListConstraints constraints, + boolean tableOnly + ) { + List objectTypes = tableOnly ? damengTableObjectTypes(constraints) : damengObjectTypes(constraints); + if (objectTypes.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND o.OBJECT_TYPE IN (").append(placeholders(objectTypes.size())).append(")"); + args.addAll(objectTypes); + } + + private static List damengTableObjectTypes(MetadataListConstraints constraints) { + List result = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + result.add("TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + result.add("VIEW"); + } + if (constraints.tableTypeAllowed("MATERIALIZED_VIEW")) { + result.add("MATERIALIZED VIEW"); + } + return result; + } + + private static List damengObjectTypes(MetadataListConstraints constraints) { + List result = damengTableObjectTypes(constraints); + if (constraints.objectTypeAllowed("PROCEDURE")) { + result.add("PROCEDURE"); + } + if (constraints.objectTypeAllowed("FUNCTION")) { + result.add("FUNCTION"); + } + return result; + } + + private static void appendNameFilter(StringBuilder sql, List args, String column, MetadataListConstraints constraints) { + if (!constraints.hasFilter()) { + return; + } + sql.append(" AND UPPER(").append(column).append(") LIKE ? ESCAPE '\\\\'"); + args.add(constraints.fuzzyLikePattern().toUpperCase(Locale.ROOT)); + } + + private static void appendLimitOffset(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasLimit()) { + return; + } + sql.append(" LIMIT ?"); + args.add(constraints.getLimit()); + if (constraints.hasOffset()) { + sql.append(" OFFSET ?"); + args.add(constraints.getOffset()); + } + } + + private static String placeholders(int count) { + return String.join(", ", java.util.Collections.nCopies(count, "?")); + } + + private static void bind(PreparedStatement stmt, List args) throws Exception { + for (int index = 0; index < args.size(); index += 1) { + Object arg = args.get(index); + if (arg instanceof Integer) { + stmt.setInt(index + 1, (Integer) arg); + } else if (arg == null) { + stmt.setObject(index + 1, null); + } else { + stmt.setString(index + 1, String.valueOf(arg)); + } + } + } + + static final class MetadataQuery { + private final String sql; + private final List args; + + MetadataQuery(String sql, List args) { + this.sql = sql; + this.args = args; + } + + String sql() { + return sql; + } + + List args() { + return args; + } + } + private static String normalizeObjectType(String value) { String upper = value == null ? "" : value.trim().toUpperCase(Locale.ROOT).replace(' ', '_').replace('-', '_'); if (upper.contains("MATERIALIZED") && upper.contains("VIEW")) { @@ -349,6 +511,71 @@ public final class DamengAgent extends BaseDatabaseAgent { }); } + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listObjects(schema); + } + return queryConstrainedObjects(schema, normalized); + } + + private List queryConstrainedObjects(String schema, MetadataListConstraints constraints) { + if (!includesSupportedObjectTypes(constraints)) { + return List.of(); + } + try { + return unchecked(() -> { + MetadataQuery query = buildConstrainedObjectsQuery(schema, constraints); + List result = new ArrayList<>(); + try (PreparedStatement stmt = requireConnected().prepareStatement(query.sql())) { + bind(stmt, query.args()); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo( + rs.getString("OBJECT_NAME"), + normalizeObjectType(rs.getString("OBJECT_TYPE")), + schema, + rs.getString("COMMENTS") + )); + } + } + } + return constraints.withoutPaging().filterObjects(result); + }); + } catch (RuntimeException e) { + // Keep restricted/older Dameng catalogs usable even if SQL pushdown + // is unavailable for a specific connection. + return constraints.filterObjects(listObjects(schema)); + } + } + + static MetadataQuery buildConstrainedObjectsQuery(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder(""" + SELECT o.OBJECT_NAME, + CASE WHEN o.OBJECT_TYPE = 'MATERIALIZED VIEW' THEN 'MATERIALIZED_VIEW' ELSE o.OBJECT_TYPE END AS OBJECT_TYPE, + c.COMMENTS + FROM ALL_OBJECTS o + LEFT JOIN ALL_TAB_COMMENTS c ON c.OWNER = o.OWNER AND c.TABLE_NAME = o.OBJECT_NAME + WHERE o.OWNER = ? + """.stripIndent().trim()); + args.add(schema); + appendDamengObjectTypePredicate(sql, args, normalized, false); + sql.append(" AND (o.OBJECT_TYPE <> 'TABLE' OR o.OBJECT_NAME NOT LIKE 'MTAB$_%')"); + appendNameFilter(sql, args, "o.OBJECT_NAME", normalized); + sql.append(" ORDER BY CASE o.OBJECT_TYPE") + .append(" WHEN 'TABLE' THEN 0") + .append(" WHEN 'VIEW' THEN 1") + .append(" WHEN 'MATERIALIZED VIEW' THEN 2") + .append(" WHEN 'PROCEDURE' THEN 3") + .append(" WHEN 'FUNCTION' THEN 4") + .append(" ELSE 9 END, o.OBJECT_NAME"); + appendLimitOffset(sql, args, normalized); + return new MetadataQuery(sql.toString(), args); + } + @Override public ObjectSource getObjectSource(String schema, String name, String objectType) { return unchecked(() -> { diff --git a/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java index 9670eb5464..afa21b72a1 100644 --- a/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java +++ b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java @@ -1,7 +1,14 @@ package com.dbx.agent.dameng; import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; class DamengAgentTest extends JdbcFakeExecutionBehaviorTest { @Override @@ -13,4 +20,42 @@ class DamengAgentTest extends JdbcFakeExecutionBehaviorTest { protected String resultSetSql() { return "CALL SP_SAMPLE()"; } + + @Test + void constrainedTableQueryPushesFilterTypeAndPagingToDameng() { + DamengAgent.MetadataQuery query = DamengAgent.buildConstrainedTablesQuery( + "APP", + new MetadataListConstraints("ord", 50, 100, List.of("TABLE", "VIEW")) + ); + + assertTrue(query.sql().contains("FROM ALL_OBJECTS o")); + assertTrue(query.sql().contains("o.OBJECT_TYPE IN (?, ?)")); + assertTrue(query.sql().contains("UPPER(o.OBJECT_NAME) LIKE ? ESCAPE '\\\\'")); + assertTrue(query.sql().endsWith("LIMIT ? OFFSET ?")); + assertEquals(List.of("APP", "TABLE", "VIEW", "%O%R%D%", 50, 100), query.args()); + } + + @Test + void constrainedTableQueryMapsMaterializedViewTypeToDamengCatalogName() { + DamengAgent.MetadataQuery query = DamengAgent.buildConstrainedTablesQuery( + "APP", + new MetadataListConstraints(null, 20, null, List.of("MATERIALIZED_VIEW")) + ); + + assertTrue(query.sql().contains("MATERIALIZED_VIEW")); + assertEquals(List.of("APP", "MATERIALIZED VIEW", 20), query.args()); + } + + @Test + void constrainedObjectQueryPushesRoutineOnlySearchToDameng() { + DamengAgent.MetadataQuery query = DamengAgent.buildConstrainedObjectsQuery( + "APP", + new MetadataListConstraints("sync", 20, null, List.of("PROCEDURE", "FUNCTION")) + ); + + assertTrue(query.sql().contains("o.OBJECT_TYPE IN (?, ?)")); + assertTrue(query.sql().contains("WHEN 'PROCEDURE' THEN 3")); + assertTrue(query.sql().endsWith("LIMIT ?")); + assertEquals(List.of("APP", "PROCEDURE", "FUNCTION", "%S%Y%N%C%", 20), query.args()); + } } diff --git a/agents/drivers/databend/src/main/java/com/dbx/agent/databend/DatabendAgent.java b/agents/drivers/databend/src/main/java/com/dbx/agent/databend/DatabendAgent.java index e70e195d29..48c00e4786 100644 --- a/agents/drivers/databend/src/main/java/com/dbx/agent/databend/DatabendAgent.java +++ b/agents/drivers/databend/src/main/java/com/dbx/agent/databend/DatabendAgent.java @@ -3,11 +3,13 @@ package com.dbx.agent.databend; import com.dbx.agent.ConfiguredJdbcAgent; import com.dbx.agent.JdbcAgentProfile; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.ObjectInfo; import com.dbx.agent.ObjectSource; import java.sql.ResultSet; import java.sql.SQLException; +import java.sql.PreparedStatement; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -54,6 +56,23 @@ public final class DatabendAgent extends ConfiguredJdbcAgent { }); } + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + boolean includeTables = normalized.includesTableLikeTypes(); + boolean includeProcedures = includesProcedures(normalized); + if (includeProcedures && !includeTables) { + return listProcedures(schema, normalized); + } + if (includeTables && !includeProcedures) { + return super.listObjects(schema, normalized); + } + if (!includeTables && !includeProcedures) { + return List.of(); + } + return normalized.filterObjects(listObjects(schema)); + } + @Override public ObjectSource getObjectSource(String schema, String name, String objectType) { if (!"PROCEDURE".equalsIgnoreCase(objectType)) { @@ -80,6 +99,53 @@ public final class DatabendAgent extends ConfiguredJdbcAgent { } } + private List listProcedures(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + useSchema(schema); + List result = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT name, comment FROM system.procedures"); + List args = new ArrayList<>(); + if (constraints.hasFilter()) { + sql.append(" WHERE UPPER(name) LIKE ? ESCAPE '\\\\'"); + args.add(constraints.fuzzyLikePattern().toUpperCase(Locale.ROOT)); + } + sql.append(" ORDER BY name"); + if (constraints.hasLimit()) { + sql.append(" LIMIT ?"); + args.add(constraints.getLimit()); + if (constraints.hasOffset()) { + sql.append(" OFFSET ?"); + args.add(constraints.getOffset()); + } + } + try (PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString("name"), "PROCEDURE", schema, rs.getString("comment"))); + } + } + } + MetadataListConstraints guard = constraints.hasLimit() ? constraints.withoutPaging() : constraints; + return guard.filterObjects(result); + }); + } + + private static boolean includesProcedures(MetadataListConstraints constraints) { + return !constraints.hasObjectTypes() || constraints.objectTypeAllowed("PROCEDURE"); + } + + private static void bind(PreparedStatement stmt, List args) throws SQLException { + for (int index = 0; index < args.size(); index += 1) { + Object arg = args.get(index); + if (arg instanceof Integer) { + stmt.setInt(index + 1, (Integer) arg); + } else { + stmt.setString(index + 1, String.valueOf(arg)); + } + } + } + private ProcedureMetadata findProcedure(String name) throws SQLException { try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement( "SELECT arguments, comment FROM system.procedures WHERE name = ? ORDER BY procedure_id LIMIT 1" diff --git a/agents/drivers/databend/src/test/java/com/dbx/agent/databend/DatabendAgentTest.java b/agents/drivers/databend/src/test/java/com/dbx/agent/databend/DatabendAgentTest.java index 327aa848c2..d875129281 100644 --- a/agents/drivers/databend/src/test/java/com/dbx/agent/databend/DatabendAgentTest.java +++ b/agents/drivers/databend/src/test/java/com/dbx/agent/databend/DatabendAgentTest.java @@ -2,10 +2,24 @@ package com.dbx.agent.databend; import com.dbx.agent.ConnectParams; import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.ObjectInfo; import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import com.dbx.agent.test.TestSupport; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; + class DatabendAgentTest extends JdbcFakeExecutionBehaviorTest { @Override protected DatabaseAgent createAgent() { @@ -51,4 +65,132 @@ class DatabendAgentTest extends JdbcFakeExecutionBehaviorTest { DatabendAgent.DATABEND_PROFILE.schemaSwitchSql("analytics\"prod") ); } + + @Test + void constrainedListObjectsKeepsDatabendProcedures() { + List sql = new ArrayList<>(); + DatabendAgent agent = new DatabendAgent(); + TestSupport.setPrivateConnection(agent, databendMetadataConnection(sql, resultSet( + new String[]{"name", "comment"}, + new Object[][]{ + {"refresh_orders", "proc comment"}, + {"format_order", null} + } + ))); + + List objects = agent.listObjects( + "default", + new MetadataListConstraints("format", 1, null, List.of("PROCEDURE")) + ); + + Assertions.assertEquals(1, objects.size()); + Assertions.assertEquals("format_order", objects.get(0).getName()); + Assertions.assertEquals("PROCEDURE", objects.get(0).getObject_type()); + Assertions.assertTrue(sql.stream().anyMatch(value -> value.contains("FROM system.procedures")), sql.toString()); + Assertions.assertTrue(sql.stream().anyMatch(value -> value.contains("UPPER(name) LIKE ?")), sql.toString()); + Assertions.assertTrue(sql.stream().anyMatch(value -> value.contains("LIMIT ?")), sql.toString()); + } + + private static Connection databendMetadataConnection(List sql, ResultSet procedures) { + DatabaseMetaData metadata = proxy(DatabaseMetaData.class, (method, args) -> { + if ("getTableTypes".equals(method.getName()) || "getTables".equals(method.getName())) { + return emptyResultSet(); + } + return defaultValue(method.getReturnType()); + }); + Statement statement = proxy(Statement.class, (method, args) -> { + if ("execute".equals(method.getName()) || "close".equals(method.getName())) { + return "execute".equals(method.getName()) ? false : null; + } + if ("executeQuery".equals(method.getName())) { + return procedures; + } + return defaultValue(method.getReturnType()); + }); + PreparedStatement preparedStatement = proxy(PreparedStatement.class, (method, args) -> { + if ("executeQuery".equals(method.getName())) { + return procedures; + } + if ("setString".equals(method.getName()) || "setInt".equals(method.getName()) || "close".equals(method.getName())) { + return null; + } + return defaultValue(method.getReturnType()); + }); + return proxy(Connection.class, (method, args) -> { + if ("getMetaData".equals(method.getName())) { + return metadata; + } + if ("createStatement".equals(method.getName())) { + return statement; + } + if ("prepareStatement".equals(method.getName())) { + sql.add(String.valueOf(args[0])); + return preparedStatement; + } + if ("isClosed".equals(method.getName())) { + return false; + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSet emptyResultSet() { + return resultSet(new String[]{"VALUE"}, new Object[][]{}); + } + + private static ResultSet resultSet(String[] columns, Object[][] rows) { + int[] index = {-1}; + return proxy(ResultSet.class, (method, args) -> { + switch (method.getName()) { + case "next": + index[0] += 1; + return index[0] < rows.length; + case "getString": + Object value = columnValue(columns, rows[index[0]], args[0]); + return value == null ? null : String.valueOf(value); + case "close": + return null; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static Object columnValue(String[] columns, Object[] row, Object key) { + if (key instanceof Number) { + return row[((Number) key).intValue() - 1]; + } + for (int i = 0; i < columns.length; i++) { + if (columns[i].equalsIgnoreCase(String.valueOf(key))) { + return row[i]; + } + } + return null; + } + + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + return handler.handle(method, args == null ? new Object[0] : args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (type == Boolean.TYPE) return false; + if (type == Byte.TYPE) return (byte) 0; + if (type == Short.TYPE) return (short) 0; + if (type == Integer.TYPE) return 0; + if (type == Long.TYPE) return 0L; + if (type == Float.TYPE) return 0f; + if (type == Double.TYPE) return 0d; + if (type == Character.TYPE) return (char) 0; + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args) throws Throwable; + } } diff --git a/agents/drivers/db2/src/main/java/com/dbx/agent/db2/Db2Agent.java b/agents/drivers/db2/src/main/java/com/dbx/agent/db2/Db2Agent.java index 613050547f..047b11cb00 100644 --- a/agents/drivers/db2/src/main/java/com/dbx/agent/db2/Db2Agent.java +++ b/agents/drivers/db2/src/main/java/com/dbx/agent/db2/Db2Agent.java @@ -10,6 +10,8 @@ import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcExecutor; import com.dbx.agent.JdbcIdentifiers; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.MetadataSqlSupport; import com.dbx.agent.ObjectInfo; import com.dbx.agent.ObjectSource; import com.dbx.agent.QueryResult; @@ -114,6 +116,44 @@ public final class Db2Agent extends BaseDatabaseAgent { }); } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + if (!normalized.includesTableLikeTypes()) { + return List.of(); + } + try { + return queryConstrainedTables(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterTables(listTables(schema)); + } + } + + private List queryConstrainedTables(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT TABNAME, TYPE FROM SYSCAT.TABLES WHERE TABSCHEMA = ?"); + args.add(schema); + appendDb2TableTypePredicate(sql, args, constraints); + MetadataSqlSupport.appendNameFilter(sql, args, "TABNAME", constraints); + sql.append(" ORDER BY TABNAME"); + MetadataSqlSupport.appendLiteralOffsetFetch(sql, constraints); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1).trim(), db2TableType(rs.getString(2)), null)); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } + @Override public List listObjects(String schema) { return unchecked(() -> { @@ -135,6 +175,63 @@ public final class Db2Agent extends BaseDatabaseAgent { }); } + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listObjects(schema); + } + if (!includesSupportedObjects(normalized)) { + return List.of(); + } + try { + return queryConstrainedObjects(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterObjects(listObjects(schema)); + } + } + + private List queryConstrainedObjects(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List branches = new ArrayList<>(); + List args = new ArrayList<>(); + if (constraints.includesTableLikeTypes()) { + StringBuilder tableSql = new StringBuilder( + "SELECT TABNAME AS OBJECT_NAME, CASE TYPE WHEN 'T' THEN 'TABLE' WHEN 'V' THEN 'VIEW' ELSE TYPE END AS OBJECT_TYPE FROM SYSCAT.TABLES WHERE TABSCHEMA = ?" + ); + args.add(schema); + appendDb2TableTypePredicate(tableSql, args, constraints); + MetadataSqlSupport.appendNameFilter(tableSql, args, "TABNAME", constraints); + branches.add(tableSql.toString()); + } + if (constraints.objectTypeAllowed("PROCEDURE")) { + StringBuilder procedureSql = new StringBuilder( + "SELECT PROCNAME AS OBJECT_NAME, 'PROCEDURE' AS OBJECT_TYPE FROM SYSCAT.PROCEDURES WHERE PROCSCHEMA = ?" + ); + args.add(schema); + MetadataSqlSupport.appendNameFilter(procedureSql, args, "PROCNAME", constraints); + branches.add(procedureSql.toString()); + } + if (branches.isEmpty()) { + return List.of(); + } + StringBuilder sql = new StringBuilder("SELECT OBJECT_NAME, OBJECT_TYPE FROM (") + .append(String.join(" UNION ALL ", branches)) + .append(") metadata_objects ORDER BY CASE OBJECT_TYPE WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'PROCEDURE' THEN 2 ELSE 9 END, OBJECT_NAME"); + MetadataSqlSupport.appendLiteralOffsetFetch(sql, constraints); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1).trim(), rs.getString(2), schema, null)); + } + } + } + return constraints.withoutPaging().filterObjects(result); + }); + } + @Override public ObjectSource getObjectSource(String schema, String name, String objectType) { return unchecked(() -> { @@ -331,6 +428,39 @@ public final class Db2Agent extends BaseDatabaseAgent { return url + ":" + extraParams + (extraParams.endsWith(";") ? "" : ";"); } + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() && !constraints.hasLimit() && !constraints.hasOffset() && !constraints.hasObjectTypes(); + } + + private static boolean includesSupportedObjects(MetadataListConstraints constraints) { + return constraints.includesTableLikeTypes() || constraints.objectTypeAllowed("PROCEDURE"); + } + + private static void appendDb2TableTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + List types = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + types.add("T"); + } + if (constraints.tableTypeAllowed("VIEW")) { + types.add("V"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND TYPE IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + + private static String db2TableType(String value) { + String db2Type = value == null ? "" : value.trim(); + return switch (db2Type) { + case "T" -> "TABLE"; + case "V" -> "VIEW"; + default -> db2Type; + }; + } + private static String trimDb2UrlParams(String urlParams) { String value = urlParams == null ? "" : urlParams.trim(); while (value.startsWith("?") || value.startsWith("&") || value.startsWith(":") || value.startsWith(";")) { diff --git a/agents/drivers/db2/src/test/java/com/dbx/agent/db2/Db2AgentTest.java b/agents/drivers/db2/src/test/java/com/dbx/agent/db2/Db2AgentTest.java index e36e528caa..1eeb38d106 100644 --- a/agents/drivers/db2/src/test/java/com/dbx/agent/db2/Db2AgentTest.java +++ b/agents/drivers/db2/src/test/java/com/dbx/agent/db2/Db2AgentTest.java @@ -2,7 +2,9 @@ package com.dbx.agent.db2; import com.dbx.agent.DatabaseAgent; import com.dbx.agent.ConnectParams; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import com.dbx.agent.test.JdbcMetadataSqlFake; import com.dbx.agent.test.TestSupport; import org.junit.jupiter.api.Test; @@ -20,6 +22,7 @@ import java.util.concurrent.atomic.AtomicReference; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; class Db2AgentTest extends JdbcFakeExecutionBehaviorTest { @Override @@ -59,6 +62,35 @@ class Db2AgentTest extends JdbcFakeExecutionBehaviorTest { assertFalse(executedSql.get().contains("OWNERTYPE")); } + @Test + void constrainedTableMetadataUsesDb2CatalogPushdown() { + Db2Agent agent = new Db2Agent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listTables("APP", new MetadataListConstraints("ord", 25, 50, List.of("TABLE"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + assertTrue(sql.contains("FROM SYSCAT.TABLES")); + assertTrue(sql.contains("TYPE IN (?)")); + assertTrue(sql.contains("UPPER(TABNAME) LIKE ? ESCAPE '\\\\'")); + assertTrue(sql.endsWith("OFFSET 50 ROWS FETCH NEXT 25 ROWS ONLY")); + assertEquals(Arrays.asList("param:1=APP", "param:2=T", "param:3=%O%R%D%"), JdbcMetadataSqlFake.statements.subList(1, 4)); + } + + @Test + void constrainedObjectMetadataUsesDb2CatalogPushdown() { + Db2Agent agent = new Db2Agent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listObjects("APP", new MetadataListConstraints("sync", 10, null, List.of("PROCEDURE"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + assertTrue(sql.contains("FROM SYSCAT.PROCEDURES")); + assertTrue(sql.contains("ORDER BY CASE OBJECT_TYPE")); + assertTrue(sql.endsWith("FETCH FIRST 10 ROWS ONLY")); + assertEquals(Arrays.asList("param:1=APP", "param:2=%S%Y%N%C%"), JdbcMetadataSqlFake.statements.subList(1, 3)); + } + private static Connection connection(AtomicReference executedSql, ResultSet resultSet) { Statement statement = proxy(Statement.class, new MethodHandler() { @Override diff --git a/agents/drivers/gbase8a/src/main/java/com/dbx/agent/gbase8a/Gbase8aAgent.java b/agents/drivers/gbase8a/src/main/java/com/dbx/agent/gbase8a/Gbase8aAgent.java index 374a26eed2..6d7b28dead 100644 --- a/agents/drivers/gbase8a/src/main/java/com/dbx/agent/gbase8a/Gbase8aAgent.java +++ b/agents/drivers/gbase8a/src/main/java/com/dbx/agent/gbase8a/Gbase8aAgent.java @@ -4,6 +4,7 @@ import com.dbx.agent.ConfiguredJdbcAgent; import com.dbx.agent.ExecuteQueryOptions; import com.dbx.agent.JdbcAgentProfile; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.ObjectInfo; import com.dbx.agent.QueryPageOptions; import com.dbx.agent.QueryPageResult; @@ -55,18 +56,34 @@ public final class Gbase8aAgent extends ConfiguredJdbcAgent { @Override public List listTables(String schema) { + return queryTables(schema, MetadataListConstraints.NONE); + } + + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + return queryTables(schema, MetadataListConstraints.orNone(constraints)); + } + + private List queryTables(String schema, MetadataListConstraints constraints) { return unchecked(() -> { - List result = new ArrayList<>(); - String sql; - if (schema != null && !schema.trim().isEmpty()) { - sql = "SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA = ? ORDER BY TABLE_NAME"; - } else { - sql = "SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA NOT IN ('information_schema', 'performance_schema', 'gclusterdb', 'gctmpdb') ORDER BY TABLE_SCHEMA, TABLE_NAME"; + if (!constraints.includesTableLikeTypes()) { + return List.of(); } - try (PreparedStatement stmt = requireConnection().prepareStatement(sql)) { - if (schema != null && !schema.trim().isEmpty()) { - stmt.setString(1, schema); - } + List result = new ArrayList<>(); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE "); + if (hasSchema(schema)) { + sql.append("TABLE_SCHEMA = ?"); + args.add(schema); + } else { + sql.append("TABLE_SCHEMA NOT IN ('information_schema', 'performance_schema', 'gclusterdb', 'gctmpdb')"); + } + appendGbase8aTableTypePredicate(sql, args, constraints); + appendNameFilter(sql, args, "TABLE_NAME", constraints); + sql.append(hasSchema(schema) ? " ORDER BY TABLE_NAME" : " ORDER BY TABLE_SCHEMA, TABLE_NAME"); + appendLimitOffset(sql, args, constraints); + try (PreparedStatement stmt = requireConnection().prepareStatement(sql.toString())) { + bind(stmt, args); try (ResultSet rs = stmt.executeQuery()) { while (rs.next()) { String tableType = rs.getString("TABLE_TYPE"); @@ -78,32 +95,63 @@ public final class Gbase8aAgent extends ConfiguredJdbcAgent { } } result.sort(Comparator.comparing(TableInfo::getName)); - return result; + MetadataListConstraints guard = constraints.hasLimit() ? constraints.withoutPaging() : constraints; + return guard.filterTables(result); }); } @Override public List listObjects(String schema) { + return queryObjects(schema, MetadataListConstraints.NONE); + } + + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + return queryObjects(schema, MetadataListConstraints.orNone(constraints)); + } + + private List queryObjects(String schema, MetadataListConstraints constraints) { return unchecked(() -> { - List result = StandardJdbcMetadata.INSTANCE.listObjects(listTables(schema), schema); - // GBase 8a's table metadata does not surface routines, so sidebar objects must load them explicitly. - appendRoutines(result, schema); - return result; + boolean includeTables = constraints.includesTableLikeTypes(); + boolean includeRoutines = includesRoutineTypes(constraints); + if (!includeTables && !includeRoutines) { + return List.of(); + } + List result = new ArrayList<>(); + if (includeTables) { + MetadataListConstraints tableConstraints = includeRoutines + ? new MetadataListConstraints(constraints.getFilter(), null, null, tableLikeObjectTypes(constraints)) + : constraints; + result.addAll(StandardJdbcMetadata.INSTANCE.listObjects(queryTables(schema, tableConstraints), schema)); + } + if (includeRoutines) { + MetadataListConstraints routineConstraints = includeTables + ? new MetadataListConstraints(constraints.getFilter(), null, null, routineObjectTypes(constraints)) + : constraints; + appendRoutines(result, schema, routineConstraints); + } + boolean singlePagedSource = includeTables != includeRoutines && constraints.hasLimit(); + MetadataListConstraints guard = singlePagedSource ? constraints.withoutPaging() : constraints; + return guard.filterObjects(result); }); } - private void appendRoutines(List result, String schema) throws Exception { - String sql; - boolean hasSchema = schema != null && !schema.trim().isEmpty(); + private void appendRoutines(List result, String schema, MetadataListConstraints constraints) throws Exception { + StringBuilder sql = new StringBuilder(); + List args = new ArrayList<>(); + boolean hasSchema = hasSchema(schema); if (hasSchema) { - sql = "SELECT ROUTINE_NAME, ROUTINE_TYPE, ROUTINE_COMMENT FROM information_schema.ROUTINES WHERE ROUTINE_SCHEMA = ? ORDER BY ROUTINE_NAME"; + sql.append("SELECT ROUTINE_NAME, ROUTINE_TYPE, ROUTINE_COMMENT FROM information_schema.ROUTINES WHERE ROUTINE_SCHEMA = ?"); + args.add(schema); } else { - sql = "SELECT ROUTINE_SCHEMA, ROUTINE_NAME, ROUTINE_TYPE, ROUTINE_COMMENT FROM information_schema.ROUTINES WHERE ROUTINE_SCHEMA NOT IN ('information_schema', 'performance_schema', 'gclusterdb', 'gctmpdb') ORDER BY ROUTINE_SCHEMA, ROUTINE_NAME"; + sql.append("SELECT ROUTINE_SCHEMA, ROUTINE_NAME, ROUTINE_TYPE, ROUTINE_COMMENT FROM information_schema.ROUTINES WHERE ROUTINE_SCHEMA NOT IN ('information_schema', 'performance_schema', 'gclusterdb', 'gctmpdb')"); } - try (PreparedStatement stmt = requireConnection().prepareStatement(sql)) { - if (hasSchema) { - stmt.setString(1, schema); - } + appendGbase8aRoutineTypePredicate(sql, args, constraints); + appendNameFilter(sql, args, "ROUTINE_NAME", constraints); + sql.append(hasSchema ? " ORDER BY ROUTINE_NAME" : " ORDER BY ROUTINE_SCHEMA, ROUTINE_NAME"); + appendLimitOffset(sql, args, constraints); + try (PreparedStatement stmt = requireConnection().prepareStatement(sql.toString())) { + bind(stmt, args); try (ResultSet rs = stmt.executeQuery()) { while (rs.next()) { String routineSchema = hasSchema ? schema : rs.getString("ROUTINE_SCHEMA"); @@ -118,6 +166,126 @@ public final class Gbase8aAgent extends ConfiguredJdbcAgent { } } + private static void appendGbase8aTableTypePredicate( + StringBuilder sql, + List args, + MetadataListConstraints constraints + ) { + if (!constraints.hasObjectTypes()) { + return; + } + List tableTypes = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + tableTypes.add("BASE TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + tableTypes.add("VIEW"); + } + if (tableTypes.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND TABLE_TYPE IN (").append(placeholders(tableTypes.size())).append(")"); + args.addAll(tableTypes); + } + + private static void appendGbase8aRoutineTypePredicate( + StringBuilder sql, + List args, + MetadataListConstraints constraints + ) { + if (!constraints.hasObjectTypes()) { + return; + } + List routineTypes = new ArrayList<>(); + if (constraints.objectTypeAllowed("PROCEDURE")) { + routineTypes.add("PROCEDURE"); + } + if (constraints.objectTypeAllowed("FUNCTION")) { + routineTypes.add("FUNCTION"); + } + if (routineTypes.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND ROUTINE_TYPE IN (").append(placeholders(routineTypes.size())).append(")"); + args.addAll(routineTypes); + } + + private static void appendNameFilter(StringBuilder sql, List args, String column, MetadataListConstraints constraints) { + if (!constraints.hasFilter()) { + return; + } + sql.append(" AND UPPER(").append(column).append(") LIKE ? ESCAPE '\\\\'"); + args.add(constraints.fuzzyLikePattern().toUpperCase(Locale.ROOT)); + } + + private static void appendLimitOffset(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasLimit()) { + return; + } + sql.append(" LIMIT ?"); + args.add(constraints.getLimit()); + if (constraints.hasOffset()) { + sql.append(" OFFSET ?"); + args.add(constraints.getOffset()); + } + } + + private static boolean hasSchema(String schema) { + return schema != null && !schema.trim().isEmpty(); + } + + private static boolean includesRoutineTypes(MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return true; + } + return constraints.objectTypeAllowed("PROCEDURE") || constraints.objectTypeAllowed("FUNCTION"); + } + + private static List tableLikeObjectTypes(MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return null; + } + List result = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + result.add("TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + result.add("VIEW"); + } + return result; + } + + private static List routineObjectTypes(MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return null; + } + List result = new ArrayList<>(); + if (constraints.objectTypeAllowed("PROCEDURE")) { + result.add("PROCEDURE"); + } + if (constraints.objectTypeAllowed("FUNCTION")) { + result.add("FUNCTION"); + } + return result; + } + + private static String placeholders(int count) { + return String.join(", ", java.util.Collections.nCopies(count, "?")); + } + + private static void bind(PreparedStatement stmt, List args) throws Exception { + for (int index = 0; index < args.size(); index += 1) { + Object arg = args.get(index); + if (arg instanceof Integer) { + stmt.setInt(index + 1, (Integer) arg); + } else { + stmt.setString(index + 1, String.valueOf(arg)); + } + } + } + private static String normalizeRoutineType(String routineType) { if (routineType == null || routineType.trim().isEmpty()) { return "PROCEDURE"; diff --git a/agents/drivers/gbase8a/src/test/java/com/dbx/agent/gbase8a/Gbase8aAgentTest.java b/agents/drivers/gbase8a/src/test/java/com/dbx/agent/gbase8a/Gbase8aAgentTest.java index f6f281e62a..a936d7a2b6 100644 --- a/agents/drivers/gbase8a/src/test/java/com/dbx/agent/gbase8a/Gbase8aAgentTest.java +++ b/agents/drivers/gbase8a/src/test/java/com/dbx/agent/gbase8a/Gbase8aAgentTest.java @@ -1,6 +1,8 @@ package com.dbx.agent.gbase8a; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.ObjectInfo; +import com.dbx.agent.TableInfo; import com.dbx.agent.test.TestSupport; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -46,6 +48,53 @@ class Gbase8aAgentTest { Assertions.assertTrue(sql.get(1).contains("ROUTINE_SCHEMA = ?"), sql.get(1)); } + @Test + void constrainedListTablesPushesFilterTypeAndLimit() { + List sql = new ArrayList<>(); + Gbase8aAgent agent = new Gbase8aAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, + resultSet( + new String[]{"TABLE_NAME", "TABLE_TYPE"}, + new Object[][]{{"user_order", "BASE TABLE"}} + ) + )); + + List tables = agent.listTables( + "app", + new MetadataListConstraints("user", 1, 1, List.of("TABLE")) + ); + + Assertions.assertEquals(1, tables.size()); + Assertions.assertEquals("user_order", tables.get(0).getName()); + Assertions.assertTrue(sql.get(0).contains("TABLE_TYPE IN (?)"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("UPPER(TABLE_NAME) LIKE ?"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("LIMIT ?"), sql.get(0)); + } + + @Test + void constrainedListObjectsKeepsCustomRoutineMetadata() { + List sql = new ArrayList<>(); + Gbase8aAgent agent = new Gbase8aAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, + resultSet( + new String[]{"ROUTINE_NAME", "ROUTINE_TYPE", "ROUTINE_COMMENT"}, + new Object[][]{{"refresh_orders", "PROCEDURE", "proc comment"}, {"format_order", "FUNCTION", null}} + ) + )); + + List objects = agent.listObjects( + "app", + new MetadataListConstraints("format", 1, null, List.of("FUNCTION")) + ); + + Assertions.assertEquals(1, objects.size()); + Assertions.assertEquals("format_order", objects.get(0).getName()); + Assertions.assertEquals("FUNCTION", objects.get(0).getObject_type()); + Assertions.assertTrue(sql.get(0).contains("ROUTINE_TYPE IN (?)"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("UPPER(ROUTINE_NAME) LIKE ?"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("LIMIT ?"), sql.get(0)); + } + private static Connection preparedConnection(List sql, ResultSet... resultSets) { int[] resultSetIndex = {0}; PreparedStatement statement = proxy(PreparedStatement.class, (method, args) -> { @@ -54,7 +103,7 @@ class Gbase8aAgentTest { resultSetIndex[0] += 1; return resultSets[current]; } - if ("setString".equals(method.getName()) || "close".equals(method.getName())) { + if ("setString".equals(method.getName()) || "setInt".equals(method.getName()) || "close".equals(method.getName())) { return null; } return defaultValue(method.getReturnType()); diff --git a/agents/drivers/gbase8s/src/main/java/com/dbx/agent/gbase8s/Gbase8sAgent.java b/agents/drivers/gbase8s/src/main/java/com/dbx/agent/gbase8s/Gbase8sAgent.java index 6824fca345..1db8a20eb9 100644 --- a/agents/drivers/gbase8s/src/main/java/com/dbx/agent/gbase8s/Gbase8sAgent.java +++ b/agents/drivers/gbase8s/src/main/java/com/dbx/agent/gbase8s/Gbase8sAgent.java @@ -6,6 +6,7 @@ import com.dbx.agent.DatabaseInfo; import com.dbx.agent.ExecuteQueryOptions; import com.dbx.agent.JdbcAgentProfile; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.QueryResult; import com.dbx.agent.TableInfo; import java.sql.Connection; @@ -185,6 +186,60 @@ public final class Gbase8sAgent extends ConfiguredJdbcAgent { } } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + return queryConstrainedTables(schema, normalized); + } + + private List queryConstrainedTables(String schema, MetadataListConstraints constraints) { + if (!constraints.includesTableLikeTypes()) { + return List.of(); + } + try { + List result = new ArrayList<>(); + List args = new ArrayList<>(); + String owner = trim(schema); + // GBase 8s follows Informix-style SKIP/FIRST pagination in the SELECT list. + StringBuilder sql = new StringBuilder("SELECT "); + if (constraints.hasOffset()) { + sql.append("SKIP ").append(constraints.getOffset()).append(' '); + } + if (constraints.hasLimit()) { + sql.append("FIRST ").append(constraints.getLimit()).append(' '); + } + sql.append("tabname, tabtype FROM systables WHERE tabid >= 100"); + appendGbase8sTableTypePredicate(sql, constraints); + if (!owner.isEmpty()) { + sql.append(" AND owner = ?"); + args.add(owner); + } + if (constraints.hasFilter()) { + sql.append(" AND UPPER(tabname) LIKE ? ESCAPE '\\\\'"); + args.add(constraints.fuzzyLikePattern().toUpperCase(Locale.ROOT)); + } + sql.append(" ORDER BY tabname"); + try (PreparedStatement stmt = requireConnection().prepareStatement(sql.toString())) { + bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo( + trim(rs.getString("tabname")), + tableType(rs.getString("tabtype")) + )); + } + } + } + result.sort(Comparator.comparing(TableInfo::getName)); + return constraints.withoutPaging().filterTables(result); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + public static void main(String[] args) { new JsonRpcServer(new Gbase8sAgent()).run(); } @@ -281,6 +336,38 @@ public final class Gbase8sAgent extends ConfiguredJdbcAgent { return "V".equalsIgnoreCase(trim(tabtype)) ? "VIEW" : "TABLE"; } + private static void appendGbase8sTableTypePredicate(StringBuilder sql, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + sql.append(" AND tabtype IN ('T', 'V')"); + return; + } + List tabTypes = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + tabTypes.add("'T'"); + } + if (constraints.tableTypeAllowed("VIEW")) { + tabTypes.add("'V'"); + } + if (tabTypes.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND tabtype IN (").append(String.join(", ", tabTypes)).append(")"); + } + + private static void bind(PreparedStatement stmt, List args) throws Exception { + for (int index = 0; index < args.size(); index += 1) { + stmt.setString(index + 1, String.valueOf(args.get(index))); + } + } + + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() + && !constraints.hasLimit() + && !constraints.hasOffset() + && !constraints.hasObjectTypes(); + } + private static String trim(String value) { return value == null ? "" : value.trim(); } diff --git a/agents/drivers/gbase8s/src/test/java/com/dbx/agent/gbase8s/Gbase8sAgentTest.java b/agents/drivers/gbase8s/src/test/java/com/dbx/agent/gbase8s/Gbase8sAgentTest.java index 32183511a2..726873c4ab 100644 --- a/agents/drivers/gbase8s/src/test/java/com/dbx/agent/gbase8s/Gbase8sAgentTest.java +++ b/agents/drivers/gbase8s/src/test/java/com/dbx/agent/gbase8s/Gbase8sAgentTest.java @@ -1,9 +1,21 @@ package com.dbx.agent.gbase8s; import com.dbx.agent.ConnectParams; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.TableInfo; +import com.dbx.agent.test.TestSupport; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.List; + class Gbase8sAgentTest { @Test void declaresGbase8sProfile() { @@ -95,4 +107,108 @@ class Gbase8sAgentTest { Assertions.assertEquals("jdbc:gbasedbt-sqli://db.example.com:20013/app:GBASEDBTSERVER=gbase01", url); } + + @Test + void constrainedListTablesUsesGbase8sSystemTableQuery() { + List sql = new ArrayList<>(); + Gbase8sAgent agent = new Gbase8sAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet( + new String[]{"tabname", "tabtype"}, + new Object[][]{ + {"user_order", "T"} + } + ))); + + List tables = agent.listTables( + "app", + new MetadataListConstraints("user", 1, 1, List.of("TABLE")) + ); + + Assertions.assertEquals(1, tables.size()); + Assertions.assertEquals("user_order", tables.get(0).getName()); + Assertions.assertTrue(sql.get(0).contains("FROM systables"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("SELECT SKIP 1 FIRST 1"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("UPPER(tabname) LIKE ?"), sql.get(0)); + } + + private static Connection preparedConnection(List sql, ResultSet resultSet) { + PreparedStatement statement = proxy(PreparedStatement.class, (method, args) -> { + if ("executeQuery".equals(method.getName())) { + return resultSet; + } + if ("setString".equals(method.getName()) || "close".equals(method.getName())) { + return null; + } + return defaultValue(method.getReturnType()); + }); + return proxy(Connection.class, (method, args) -> { + if ("getCatalog".equals(method.getName())) { + return "appdb"; + } + if ("prepareStatement".equals(method.getName())) { + sql.add(String.valueOf(args[0])); + return statement; + } + if ("isClosed".equals(method.getName())) { + return false; + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSet resultSet(String[] columns, Object[][] rows) { + int[] index = {-1}; + return proxy(ResultSet.class, (method, args) -> { + switch (method.getName()) { + case "next": + index[0] += 1; + return index[0] < rows.length; + case "getString": + Object value = columnValue(columns, rows[index[0]], args[0]); + return value == null ? null : String.valueOf(value); + case "close": + return null; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static Object columnValue(String[] columns, Object[] row, Object key) { + if (key instanceof Number) { + return row[((Number) key).intValue() - 1]; + } + for (int i = 0; i < columns.length; i++) { + if (columns[i].equalsIgnoreCase(String.valueOf(key))) { + return row[i]; + } + } + return null; + } + + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + return handler.handle(method, args == null ? new Object[0] : args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (type == Boolean.TYPE) return false; + if (type == Byte.TYPE) return (byte) 0; + if (type == Short.TYPE) return (short) 0; + if (type == Integer.TYPE) return 0; + if (type == Long.TYPE) return 0L; + if (type == Float.TYPE) return 0f; + if (type == Double.TYPE) return 0d; + if (type == Character.TYPE) return (char) 0; + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args) throws Throwable; + } } diff --git a/agents/drivers/goldendb/src/main/java/com/dbx/agent/goldendb/GoldendbAgent.java b/agents/drivers/goldendb/src/main/java/com/dbx/agent/goldendb/GoldendbAgent.java index 04590b0cc9..4f7596f5ba 100644 --- a/agents/drivers/goldendb/src/main/java/com/dbx/agent/goldendb/GoldendbAgent.java +++ b/agents/drivers/goldendb/src/main/java/com/dbx/agent/goldendb/GoldendbAgent.java @@ -10,6 +10,8 @@ import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcExecutor; import com.dbx.agent.JdbcIdentifiers; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.MetadataSqlSupport; import com.dbx.agent.ObjectInfo; import com.dbx.agent.ObjectSource; import com.dbx.agent.QueryResult; @@ -99,6 +101,44 @@ public final class GoldendbAgent extends BaseDatabaseAgent { }); } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + if (!normalized.includesTableLikeTypes()) { + return List.of(); + } + try { + return queryConstrainedTables(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterTables(listTables(schema)); + } + } + + private List queryConstrainedTables(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT TABLE_NAME, TABLE_TYPE, TABLE_COMMENT FROM information_schema.TABLES WHERE TABLE_SCHEMA = ?"); + args.add(schema); + appendInformationSchemaTableTypePredicate(sql, args, constraints); + MetadataSqlSupport.appendNameFilter(sql, args, "TABLE_NAME", constraints); + sql.append(" ORDER BY TABLE_NAME"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString("TABLE_NAME"), normalizeTableType(rs.getString("TABLE_TYPE")), emptyToNull(rs.getString("TABLE_COMMENT")))); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } + @Override public List listObjects(String schema) { return unchecked(() -> { @@ -121,6 +161,60 @@ public final class GoldendbAgent extends BaseDatabaseAgent { }); } + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listObjects(schema); + } + if (!includesSupportedObjects(normalized)) { + return List.of(); + } + try { + return queryConstrainedObjects(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterObjects(listObjects(schema)); + } + } + + private List queryConstrainedObjects(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List branches = new ArrayList<>(); + List args = new ArrayList<>(); + if (constraints.includesTableLikeTypes()) { + StringBuilder tableSql = new StringBuilder("SELECT TABLE_NAME AS OBJECT_NAME, TABLE_TYPE AS OBJECT_TYPE, TABLE_COMMENT AS OBJECT_COMMENT FROM information_schema.TABLES WHERE TABLE_SCHEMA = ?"); + args.add(schema); + appendInformationSchemaTableTypePredicate(tableSql, args, constraints); + MetadataSqlSupport.appendNameFilter(tableSql, args, "TABLE_NAME", constraints); + branches.add(tableSql.toString()); + } + if (constraints.objectTypeAllowed("PROCEDURE") || constraints.objectTypeAllowed("FUNCTION")) { + StringBuilder routineSql = new StringBuilder("SELECT ROUTINE_NAME AS OBJECT_NAME, ROUTINE_TYPE AS OBJECT_TYPE, NULL AS OBJECT_COMMENT FROM information_schema.ROUTINES WHERE ROUTINE_SCHEMA = ?"); + args.add(schema); + appendRoutineTypePredicate(routineSql, args, constraints); + MetadataSqlSupport.appendNameFilter(routineSql, args, "ROUTINE_NAME", constraints); + branches.add(routineSql.toString()); + } + if (branches.isEmpty()) { + return List.of(); + } + StringBuilder sql = new StringBuilder("SELECT OBJECT_NAME, OBJECT_TYPE, OBJECT_COMMENT FROM (") + .append(String.join(" UNION ALL ", branches)) + .append(") metadata_objects ORDER BY CASE OBJECT_TYPE WHEN 'BASE TABLE' THEN 0 WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'PROCEDURE' THEN 2 WHEN 'FUNCTION' THEN 3 ELSE 9 END, OBJECT_NAME"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString("OBJECT_NAME"), normalizeTableType(rs.getString("OBJECT_TYPE")), schema, emptyToNull(rs.getString("OBJECT_COMMENT")))); + } + } + } + return constraints.withoutPaging().filterObjects(result); + }); + } + @Override public ObjectSource getObjectSource(String schema, String name, String objectType) { return unchecked(() -> { @@ -379,6 +473,58 @@ public final class GoldendbAgent extends BaseDatabaseAgent { + "?useSSL=false&allowPublicKeyRetrieval=true"; } + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() && !constraints.hasLimit() && !constraints.hasOffset() && !constraints.hasObjectTypes(); + } + + private static boolean includesSupportedObjects(MetadataListConstraints constraints) { + return constraints.includesTableLikeTypes() + || constraints.objectTypeAllowed("PROCEDURE") + || constraints.objectTypeAllowed("FUNCTION"); + } + + private static void appendInformationSchemaTableTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List types = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + types.add("BASE TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + types.add("VIEW"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND TABLE_TYPE IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + + private static void appendRoutineTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List types = new ArrayList<>(); + if (constraints.objectTypeAllowed("PROCEDURE")) { + types.add("PROCEDURE"); + } + if (constraints.objectTypeAllowed("FUNCTION")) { + types.add("FUNCTION"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND ROUTINE_TYPE IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + + private static String normalizeTableType(String tableType) { + return "BASE TABLE".equals(tableType) ? "TABLE" : tableType; + } + private static String emptyToNull(String value) { return value == null || value.isEmpty() ? null : value; } diff --git a/agents/drivers/goldendb/src/test/java/com/dbx/agent/goldendb/GoldendbAgentMetadataTest.java b/agents/drivers/goldendb/src/test/java/com/dbx/agent/goldendb/GoldendbAgentMetadataTest.java index ba3d9e1ac6..6ec7a08cc7 100644 --- a/agents/drivers/goldendb/src/test/java/com/dbx/agent/goldendb/GoldendbAgentMetadataTest.java +++ b/agents/drivers/goldendb/src/test/java/com/dbx/agent/goldendb/GoldendbAgentMetadataTest.java @@ -1,5 +1,6 @@ package com.dbx.agent.goldendb; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.test.TestSupport; import com.dbx.agent.test.JdbcAgentFake; import com.dbx.agent.test.JdbcMetadataSqlFake; @@ -22,4 +23,32 @@ class GoldendbAgentMetadataTest { JdbcMetadataSqlFake.statements ); } + + @Test + void constrainedTableMetadataPushesFilterTypesAndPaging() { + GoldendbAgent agent = new GoldendbAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listTables("app", new MetadataListConstraints("ord", 25, 50, List.of("TABLE", "VIEW"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + Assertions.assertTrue(sql.contains("FROM information_schema.TABLES"), sql); + Assertions.assertTrue(sql.contains("TABLE_TYPE IN (?, ?)"), sql); + Assertions.assertTrue(sql.contains("UPPER(TABLE_NAME) LIKE ? ESCAPE '\\\\'"), sql); + Assertions.assertTrue(sql.endsWith("LIMIT 25 OFFSET 50"), sql); + } + + @Test + void constrainedObjectMetadataPushesRoutineTypesAndPaging() { + GoldendbAgent agent = new GoldendbAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listObjects("app", new MetadataListConstraints("sync", 10, null, List.of("PROCEDURE", "FUNCTION"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + Assertions.assertTrue(sql.contains("FROM information_schema.ROUTINES"), sql); + Assertions.assertTrue(sql.contains("ROUTINE_TYPE IN (?, ?)"), sql); + Assertions.assertTrue(sql.contains("ORDER BY CASE OBJECT_TYPE"), sql); + Assertions.assertTrue(sql.endsWith("LIMIT 10"), sql); + } } diff --git a/agents/drivers/h2/src/main/java/com/dbx/agent/h2/H2Agent.java b/agents/drivers/h2/src/main/java/com/dbx/agent/h2/H2Agent.java index 7ec4374e3b..af12099534 100644 --- a/agents/drivers/h2/src/main/java/com/dbx/agent/h2/H2Agent.java +++ b/agents/drivers/h2/src/main/java/com/dbx/agent/h2/H2Agent.java @@ -9,6 +9,8 @@ import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcExecutor; import com.dbx.agent.JdbcIdentifiers; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.MetadataSqlSupport; import com.dbx.agent.ObjectInfo; import com.dbx.agent.ObjectSource; import com.dbx.agent.TableInfo; @@ -93,6 +95,45 @@ public final class H2Agent extends AbstractJdbcAgent { }); } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + if (!normalized.includesTableLikeTypes()) { + return List.of(); + } + try { + return queryConstrainedTables(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterTables(listTables(schema)); + } + } + + private List queryConstrainedTables(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + String effectiveSchema = resolveSchema(schema); + List result = new ArrayList<>(); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT TABLE_NAME, TABLE_TYPE FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ?"); + args.add(effectiveSchema); + appendH2TableTypePredicate(sql, args, constraints); + MetadataSqlSupport.appendNameFilter(sql, args, "TABLE_NAME", constraints); + sql.append(" ORDER BY TABLE_NAME"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (var stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString("TABLE_NAME"), normalizeTableType(rs.getString("TABLE_TYPE")), null)); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } + @Override public List listObjects(String schema) { return unchecked(() -> { @@ -116,6 +157,61 @@ public final class H2Agent extends AbstractJdbcAgent { }); } + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listObjects(schema); + } + if (!includesSupportedObjects(normalized)) { + return List.of(); + } + try { + return queryConstrainedObjects(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterObjects(listObjects(schema)); + } + } + + private List queryConstrainedObjects(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + String effectiveSchema = resolveSchema(schema); + List result = new ArrayList<>(); + List branches = new ArrayList<>(); + List args = new ArrayList<>(); + if (constraints.includesTableLikeTypes()) { + StringBuilder tableSql = new StringBuilder("SELECT TABLE_NAME AS OBJECT_NAME, TABLE_TYPE AS OBJECT_TYPE FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ?"); + args.add(effectiveSchema); + appendH2TableTypePredicate(tableSql, args, constraints); + MetadataSqlSupport.appendNameFilter(tableSql, args, "TABLE_NAME", constraints); + branches.add(tableSql.toString()); + } + if (constraints.objectTypeAllowed("PROCEDURE") || constraints.objectTypeAllowed("FUNCTION")) { + StringBuilder routineSql = new StringBuilder("SELECT ROUTINE_NAME AS OBJECT_NAME, ROUTINE_TYPE AS OBJECT_TYPE FROM INFORMATION_SCHEMA.ROUTINES WHERE ROUTINE_SCHEMA = ?"); + args.add(effectiveSchema); + appendRoutineTypePredicate(routineSql, args, constraints); + MetadataSqlSupport.appendNameFilter(routineSql, args, "ROUTINE_NAME", constraints); + branches.add(routineSql.toString()); + } + if (branches.isEmpty()) { + return List.of(); + } + StringBuilder sql = new StringBuilder("SELECT OBJECT_NAME, OBJECT_TYPE FROM (") + .append(String.join(" UNION ALL ", branches)) + .append(") metadata_objects ORDER BY CASE OBJECT_TYPE WHEN 'BASE TABLE' THEN 0 WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'PROCEDURE' THEN 2 WHEN 'FUNCTION' THEN 3 ELSE 9 END, OBJECT_NAME"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (var stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), normalizeTableType(rs.getString(2)), schema, null)); + } + } + } + return constraints.withoutPaging().filterObjects(result); + }); + } + @Override public ObjectSource getObjectSource(String schema, String name, String objectType) { return unchecked(() -> { @@ -346,6 +442,58 @@ public final class H2Agent extends AbstractJdbcAgent { return "jdbc:h2:tcp://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); } + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() && !constraints.hasLimit() && !constraints.hasOffset() && !constraints.hasObjectTypes(); + } + + private static boolean includesSupportedObjects(MetadataListConstraints constraints) { + return constraints.includesTableLikeTypes() + || constraints.objectTypeAllowed("PROCEDURE") + || constraints.objectTypeAllowed("FUNCTION"); + } + + private static void appendH2TableTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List types = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + types.add("BASE TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + types.add("VIEW"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND TABLE_TYPE IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + + private static void appendRoutineTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List types = new ArrayList<>(); + if (constraints.objectTypeAllowed("PROCEDURE")) { + types.add("PROCEDURE"); + } + if (constraints.objectTypeAllowed("FUNCTION")) { + types.add("FUNCTION"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND ROUTINE_TYPE IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + + private static String normalizeTableType(String tableType) { + return "BASE TABLE".equals(tableType) ? "TABLE" : tableType; + } + private static String resolveSchema(String schema) { if ("PUBLIC".equalsIgnoreCase(schema) || "INFORMATION_SCHEMA".equalsIgnoreCase(schema)) { return schema.toUpperCase(Locale.ROOT); diff --git a/agents/drivers/h2/src/test/java/com/dbx/agent/h2/H2AgentTest.java b/agents/drivers/h2/src/test/java/com/dbx/agent/h2/H2AgentTest.java index a3a1ffec5d..350b1688b2 100644 --- a/agents/drivers/h2/src/test/java/com/dbx/agent/h2/H2AgentTest.java +++ b/agents/drivers/h2/src/test/java/com/dbx/agent/h2/H2AgentTest.java @@ -4,7 +4,9 @@ import com.dbx.agent.BaseDatabaseAgent; import com.dbx.agent.ConnectParams; import com.dbx.agent.DatabaseAgent; import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; import com.dbx.agent.test.JdbcExecutionBehaviorTest; import com.dbx.agent.test.JdbcMetadataBehaviorTest; import java.util.List; @@ -132,6 +134,23 @@ class H2MetadataBehaviorTest extends JdbcMetadataBehaviorTest { protected List expectedColumnsInOrder() { return List.of("ID", "NAME", "CREATED_AT"); } + + @Test + void constrainedTableMetadataFiltersTypesAndPages() { + withAgent("dbx-agent-h2-constrained-metadata", agent -> { + for (String sql : metadataFixtureSql()) { + agent.executeQuery(sql, null, new ExecuteQueryOptions()); + } + + List tables = agent.listTables( + metadataSchema(), + new MetadataListConstraints("table", 1, 1, List.of("TABLE")) + ); + + Assertions.assertEquals(1, tables.size()); + Assertions.assertEquals("BETA_TABLE", tables.get(0).getName()); + }); + } } final class H2AgentTestSupport { diff --git a/agents/drivers/informix/src/main/java/com/dbx/agent/informix/InformixAgent.java b/agents/drivers/informix/src/main/java/com/dbx/agent/informix/InformixAgent.java index d4651ff971..e8a5f5cf9f 100644 --- a/agents/drivers/informix/src/main/java/com/dbx/agent/informix/InformixAgent.java +++ b/agents/drivers/informix/src/main/java/com/dbx/agent/informix/InformixAgent.java @@ -9,6 +9,8 @@ import com.dbx.agent.ForeignKeyInfo; import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcExecutor; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.MetadataSqlSupport; import com.dbx.agent.ObjectInfo; import com.dbx.agent.ObjectSource; import com.dbx.agent.QueryResult; @@ -16,6 +18,7 @@ import com.dbx.agent.TableInfo; import com.dbx.agent.TriggerInfo; import java.sql.Connection; import java.sql.DriverManager; +import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; import java.util.ArrayList; @@ -205,6 +208,45 @@ public final class InformixAgent extends BaseDatabaseAgent { }); } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + if (!normalized.includesTableLikeTypes()) { + return List.of(); + } + try { + return queryConstrainedTables(normalized); + } catch (RuntimeException e) { + return normalized.filterTables(listTables(schema)); + } + } + + private List queryConstrainedTables(MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT "); + MetadataSqlSupport.appendLiteralSkipFirst(sql, constraints); + sql.append("tabname, CASE tabtype WHEN 'T' THEN 'TABLE' WHEN 'V' THEN 'VIEW' ELSE tabtype END ") + .append("FROM systables WHERE tabid >= 100"); + appendInformixTableTypePredicate(sql, constraints); + MetadataSqlSupport.appendNameFilter(sql, args, "tabname", constraints); + sql.append(" ORDER BY tabname"); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1).trim(), rs.getString(2).trim(), null)); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } + @Override public List listObjects(String schema) { return unchecked(() -> { @@ -234,6 +276,63 @@ public final class InformixAgent extends BaseDatabaseAgent { }); } + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listObjects(schema); + } + if (!includesSupportedObjects(normalized)) { + return List.of(); + } + try { + return queryConstrainedObjects(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterObjects(listObjects(schema)); + } + } + + private List queryConstrainedObjects(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List branches = new ArrayList<>(); + List args = new ArrayList<>(); + if (constraints.includesTableLikeTypes()) { + StringBuilder tableSql = new StringBuilder("SELECT tabname AS object_name, CASE tabtype WHEN 'T' THEN 'TABLE' WHEN 'V' THEN 'VIEW' ELSE tabtype END AS object_type, 0 AS object_order FROM systables WHERE tabid >= 100"); + appendInformixTableTypePredicate(tableSql, constraints); + MetadataSqlSupport.appendNameFilter(tableSql, args, "tabname", constraints); + branches.add(tableSql.toString()); + } + if (constraints.objectTypeAllowed("FUNCTION")) { + StringBuilder functionSql = new StringBuilder("SELECT procname AS object_name, 'FUNCTION' AS object_type, 1 AS object_order FROM sysprocedures WHERE owner != 'informix' AND isproc = 'f'"); + MetadataSqlSupport.appendNameFilter(functionSql, args, "procname", constraints); + branches.add(functionSql.toString()); + } + if (constraints.objectTypeAllowed("PROCEDURE")) { + StringBuilder procedureSql = new StringBuilder("SELECT procname AS object_name, 'PROCEDURE' AS object_type, 2 AS object_order FROM sysprocedures WHERE owner != 'informix' AND isproc = 't'"); + MetadataSqlSupport.appendNameFilter(procedureSql, args, "procname", constraints); + branches.add(procedureSql.toString()); + } + if (branches.isEmpty()) { + return List.of(); + } + StringBuilder sql = new StringBuilder("SELECT "); + MetadataSqlSupport.appendLiteralSkipFirst(sql, constraints); + sql.append("object_name, object_type FROM (") + .append(String.join(" UNION ALL ", branches)) + .append(") metadata_objects ORDER BY object_order, object_name"); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1).trim(), rs.getString(2).trim(), schema, null)); + } + } + } + return constraints.withoutPaging().filterObjects(result); + }); + } + @Override public ObjectSource getObjectSource(String schema, String name, String objectType) { return unchecked(() -> { @@ -394,6 +493,34 @@ public final class InformixAgent extends BaseDatabaseAgent { }); } + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() && !constraints.hasLimit() && !constraints.hasOffset() && !constraints.hasObjectTypes(); + } + + private static boolean includesSupportedObjects(MetadataListConstraints constraints) { + return constraints.includesTableLikeTypes() + || constraints.objectTypeAllowed("PROCEDURE") + || constraints.objectTypeAllowed("FUNCTION"); + } + + private static void appendInformixTableTypePredicate(StringBuilder sql, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List types = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + types.add("'T'"); + } + if (constraints.tableTypeAllowed("VIEW")) { + types.add("'V'"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND tabtype IN (").append(String.join(", ", types)).append(")"); + } + private static String trimStart(String value, char... chars) { int index = 0; while (index < value.length() && contains(chars, value.charAt(index))) { diff --git a/agents/drivers/informix/src/test/java/com/dbx/agent/informix/InformixAgentTest.java b/agents/drivers/informix/src/test/java/com/dbx/agent/informix/InformixAgentTest.java index 795cc8246d..ea99fbba31 100644 --- a/agents/drivers/informix/src/test/java/com/dbx/agent/informix/InformixAgentTest.java +++ b/agents/drivers/informix/src/test/java/com/dbx/agent/informix/InformixAgentTest.java @@ -1,7 +1,11 @@ package com.dbx.agent.informix; import com.dbx.agent.ConnectParams; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import com.dbx.agent.test.TestSupport; import java.util.Arrays; +import java.util.List; import java.util.Set; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -125,4 +129,33 @@ class InformixAgentTest { void listsDatabasesFromSysmasterCatalog() { Assertions.assertEquals("SELECT name FROM sysmaster:sysdatabases ORDER BY name", InformixAgent.databaseCatalogSql()); } + + @Test + void constrainedTableMetadataUsesInformixSkipFirstPushdown() { + InformixAgent agent = new InformixAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listTables("stores", new MetadataListConstraints("ord", 25, 50, List.of("TABLE"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + Assertions.assertTrue(sql.startsWith("SELECT SKIP 50 FIRST 25 tabname"), sql); + Assertions.assertTrue(sql.contains("tabtype IN ('T')"), sql); + Assertions.assertTrue(sql.contains("UPPER(tabname) LIKE ? ESCAPE '\\\\'"), sql); + Assertions.assertTrue(sql.endsWith("ORDER BY tabname"), sql); + } + + @Test + void constrainedObjectMetadataUsesInformixUnionPushdown() { + InformixAgent agent = new InformixAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listObjects("stores", new MetadataListConstraints("sync", 10, null, List.of("PROCEDURE", "FUNCTION"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + Assertions.assertTrue(sql.startsWith("SELECT FIRST 10 object_name, object_type FROM ("), sql); + Assertions.assertTrue(sql.contains("FROM sysprocedures"), sql); + Assertions.assertTrue(sql.contains("isproc = 'f'"), sql); + Assertions.assertTrue(sql.contains("isproc = 't'"), sql); + Assertions.assertTrue(sql.endsWith("ORDER BY object_order, object_name"), sql); + } } diff --git a/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java b/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java index 4e14388126..3fd839a103 100644 --- a/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java +++ b/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java @@ -7,6 +7,8 @@ import com.dbx.agent.ForeignKeyInfo; import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcIdentifiers; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.MetadataSqlSupport; import com.dbx.agent.ObjectInfo; import com.dbx.agent.ObjectSource; import com.dbx.agent.PostgresLikeAgent; @@ -125,6 +127,84 @@ public final class KingbaseAgent extends PostgresLikeAgent { }); } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + if (!normalized.includesTableLikeTypes()) { + return List.of(); + } + try { + return isMysqlCompatMode() + ? queryMysqlCompatTables(schema, normalized) + : queryRegularTables(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterTables(listTables(schema)); + } + } + + private List queryRegularTables(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT c.relname AS table_name, ") + .append("CASE c.relkind ") + .append("WHEN 'r' THEN 'TABLE' ") + .append("WHEN 'p' THEN 'TABLE' ") + .append("WHEN 'v' THEN 'VIEW' ") + .append("WHEN 'm' THEN 'MATERIALIZED_VIEW' ") + .append("WHEN 'f' THEN 'FOREIGN_TABLE' ") + .append("ELSE 'TABLE' END AS table_type, ") + .append("d.description AS table_comment ") + .append("FROM sys_catalog.sys_class c ") + .append("JOIN sys_catalog.sys_namespace n ON n.oid = c.relnamespace ") + .append("LEFT JOIN sys_catalog.sys_description d ON d.objoid = c.oid AND d.objsubid = 0 ") + .append("WHERE n.nspname = ? AND c.relkind IN ('r','p','v','m','f')"); + args.add(effectiveSchema(schema)); + appendRelkindPredicate(sql, args, constraints); + MetadataSqlSupport.appendNameFilter(sql, args, "c.relname", constraints); + sql.append(" ORDER BY c.relname"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo( + rs.getString("table_name"), + normalizeTableType(rs.getString("table_type")), + rs.getString("table_comment") + )); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } + + private List queryMysqlCompatTables(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT table_name, table_type FROM information_schema.tables WHERE table_schema = ?"); + args.add(effectiveSchema(schema)); + appendMysqlCompatTableTypePredicate(sql, args, constraints); + MetadataSqlSupport.appendNameFilter(sql, args, "table_name", constraints); + sql.append(" ORDER BY table_name"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1), normalizeTableType(rs.getString(2)))); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } + @Override public List listObjects(String schema) { return unchecked(() -> { @@ -162,6 +242,80 @@ public final class KingbaseAgent extends PostgresLikeAgent { }); } + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listObjects(schema); + } + if (!includesSupportedObjects(normalized)) { + return List.of(); + } + try { + return isMysqlCompatMode() + ? normalized.filterObjects(toObjects(queryMysqlCompatTables(schema, normalized), effectiveSchema(schema))) + : queryRegularObjects(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterObjects(listObjects(schema)); + } + } + + private List queryRegularObjects(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + String effectiveSchema = effectiveSchema(schema); + List result = new ArrayList<>(); + List branches = new ArrayList<>(); + List args = new ArrayList<>(); + if (constraints.includesTableLikeTypes()) { + StringBuilder tableSql = new StringBuilder("SELECT c.relname AS object_name, ") + .append("CASE c.relkind WHEN 'r' THEN 'TABLE' WHEN 'p' THEN 'TABLE' WHEN 'v' THEN 'VIEW' WHEN 'm' THEN 'MATERIALIZED_VIEW' WHEN 'f' THEN 'FOREIGN_TABLE' ELSE 'TABLE' END AS object_type, ") + .append("d.description AS object_comment ") + .append("FROM sys_catalog.sys_class c ") + .append("JOIN sys_catalog.sys_namespace n ON n.oid = c.relnamespace ") + .append("LEFT JOIN sys_catalog.sys_description d ON d.objoid = c.oid AND d.objsubid = 0 ") + .append("WHERE n.nspname = ? AND c.relkind IN ('r','p','v','m','f')"); + args.add(effectiveSchema); + appendRelkindPredicate(tableSql, args, constraints); + MetadataSqlSupport.appendNameFilter(tableSql, args, "c.relname", constraints); + branches.add(tableSql.toString()); + } + if (constraints.objectTypeAllowed("PROCEDURE") || constraints.objectTypeAllowed("FUNCTION")) { + StringBuilder routineSql = new StringBuilder("SELECT p.proname AS object_name, ") + .append("CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, ") + .append("d.description AS object_comment ") + .append("FROM sys_catalog.sys_proc p ") + .append("JOIN sys_catalog.sys_namespace n ON n.oid = p.pronamespace ") + .append("LEFT JOIN sys_catalog.sys_description d ON d.objoid = p.oid AND d.objsubid = 0 ") + .append("WHERE n.nspname = ?"); + args.add(effectiveSchema); + appendRoutineKindPredicate(routineSql, args, constraints); + MetadataSqlSupport.appendNameFilter(routineSql, args, "p.proname", constraints); + branches.add(routineSql.toString()); + } + if (branches.isEmpty()) { + return List.of(); + } + StringBuilder sql = new StringBuilder("SELECT object_name, object_type, object_comment FROM (") + .append(String.join(" UNION ALL ", branches)) + .append(") metadata_objects ORDER BY CASE object_type WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'MATERIALIZED_VIEW' THEN 2 WHEN 'FOREIGN_TABLE' THEN 3 WHEN 'PROCEDURE' THEN 4 WHEN 'FUNCTION' THEN 5 ELSE 9 END, object_name"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo( + rs.getString("object_name"), + rs.getString("object_type"), + effectiveSchema, + rs.getString("object_comment") + )); + } + } + } + return constraints.withoutPaging().filterObjects(result); + }); + } + @Override public ObjectSource getObjectSource(String schema, String name, String objectType) { if ("FUNCTION".equalsIgnoreCase(objectType) || "PROCEDURE".equalsIgnoreCase(objectType)) { @@ -465,6 +619,83 @@ public final class KingbaseAgent extends PostgresLikeAgent { }); } + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() && !constraints.hasLimit() && !constraints.hasOffset() && !constraints.hasObjectTypes(); + } + + private static boolean includesSupportedObjects(MetadataListConstraints constraints) { + return constraints.includesTableLikeTypes() + || constraints.objectTypeAllowed("PROCEDURE") + || constraints.objectTypeAllowed("FUNCTION"); + } + + private static void appendRelkindPredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List kinds = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + kinds.add("r"); + kinds.add("p"); + } + if (constraints.tableTypeAllowed("VIEW")) { + kinds.add("v"); + } + if (constraints.tableTypeAllowed("MATERIALIZED_VIEW")) { + kinds.add("m"); + } + if (kinds.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND c.relkind IN (").append(MetadataSqlSupport.placeholders(kinds.size())).append(")"); + args.addAll(kinds); + } + + private static void appendMysqlCompatTableTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + sql.append(" AND table_type IN ('BASE TABLE', 'VIEW')"); + return; + } + List types = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + types.add("BASE TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + types.add("VIEW"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND table_type IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + + private static void appendRoutineKindPredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + List kinds = new ArrayList<>(); + if (!constraints.hasObjectTypes() || constraints.objectTypeAllowed("PROCEDURE")) { + kinds.add("p"); + } + if (!constraints.hasObjectTypes() || constraints.objectTypeAllowed("FUNCTION")) { + kinds.add("f"); + } + if (kinds.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND p.prokind IN (").append(MetadataSqlSupport.placeholders(kinds.size())).append(")"); + args.addAll(kinds); + } + + private static List toObjects(List tables, String schema) { + List result = new ArrayList<>(); + for (TableInfo table : tables) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + return result; + } + private String effectiveSchema(String schema) { if (schema != null && !schema.trim().isEmpty()) { return schema; diff --git a/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java b/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java index ede7f4109e..eef795655c 100644 --- a/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java +++ b/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java @@ -4,6 +4,7 @@ import com.dbx.agent.ColumnInfo; import com.dbx.agent.DatabaseAgent; import com.dbx.agent.DatabaseInfo; import com.dbx.agent.IndexInfo; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.ObjectInfo; import com.dbx.agent.ObjectSource; import com.dbx.agent.TableInfo; @@ -153,6 +154,58 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { Assertions.assertTrue(sql.get(1).contains("p.prokind IN ('p','f')"), sql.get(1)); } + @Test + void constrainedRegularTableMetadataPushesFilterTypesAndPaging() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet( + new String[]{"table_name", "table_type", "table_comment"}, + new Object[][]{} + ))); + + agent.listTables("public", new MetadataListConstraints("ord", 30, 60, List.of("TABLE", "VIEW"))); + + Assertions.assertTrue(sql.get(0).contains("FROM sys_catalog.sys_class"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("c.relkind IN (?, ?, ?)"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("UPPER(c.relname) LIKE ? ESCAPE '\\\\'"), sql.get(0)); + Assertions.assertTrue(sql.get(0).endsWith("LIMIT 30 OFFSET 60"), sql.get(0)); + } + + @Test + void constrainedRegularObjectMetadataPushesRoutineTypesAndPaging() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet( + new String[]{"object_name", "object_type", "object_comment"}, + new Object[][]{} + ))); + + agent.listObjects("public", new MetadataListConstraints("sync", 10, null, List.of("PROCEDURE", "FUNCTION"))); + + Assertions.assertTrue(sql.get(0).contains("FROM sys_catalog.sys_proc"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("p.prokind IN (?, ?)"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("ORDER BY CASE object_type"), sql.get(0)); + Assertions.assertTrue(sql.get(0).endsWith("LIMIT 10"), sql.get(0)); + } + + @Test + void constrainedMysqlCompatTableMetadataPushesInformationSchemaPaging() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + agent.setMysqlCompatMode(true); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet( + new String[]{"table_name", "table_type"}, + new Object[][]{} + ))); + + agent.listTables("PUBLIC", new MetadataListConstraints("ord", 20, 40, List.of("VIEW"))); + + Assertions.assertTrue(sql.get(0).contains("FROM information_schema.tables"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("table_type IN (?)"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("UPPER(table_name) LIKE ? ESCAPE '\\\\'"), sql.get(0)); + Assertions.assertTrue(sql.get(0).endsWith("LIMIT 20 OFFSET 40"), sql.get(0)); + } + @Test void regularRoutineSourceUsesKingbaseFunctionDefinition() { List sql = new ArrayList<>(); diff --git a/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java b/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java index 683a19f887..d8419092fa 100644 --- a/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java +++ b/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java @@ -9,6 +9,7 @@ import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcAgentProfile; import com.dbx.agent.JdbcIdentifiers; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.ObjectInfo; import com.dbx.agent.TableInfo; import com.dbx.agent.TriggerInfo; @@ -81,39 +82,79 @@ public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent { @Override public List listTables(String schema) { + return queryTables(schema, MetadataListConstraints.NONE); + } + + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + return queryTables(schema, MetadataListConstraints.orNone(constraints)); + } + + private List queryTables(String schema, MetadataListConstraints constraints) { return unchecked(() -> { String owner = normalizeSchema(schema); - String sql = """ + List objectTypes = oceanBaseTableTypes(constraints); + if (objectTypes.isEmpty()) { + return List.of(); + } + String baseSql = """ SELECT o.OBJECT_NAME, CASE o.OBJECT_TYPE WHEN 'VIEW' THEN 'VIEW' ELSE 'TABLE' END AS TABLE_TYPE, c.COMMENTS FROM ALL_OBJECTS o LEFT JOIN ALL_TAB_COMMENTS c ON c.OWNER = o.OWNER AND c.TABLE_NAME = o.OBJECT_NAME - WHERE o.OWNER = ? AND o.OBJECT_TYPE IN ('TABLE', 'VIEW') - ORDER BY o.OBJECT_NAME + WHERE o.OWNER = ? AND o.OBJECT_TYPE IN (%s) """.stripIndent().trim(); + MetadataSql query = oceanBaseMetadataSql( + String.format(baseSql, placeholders(objectTypes.size())), + "OBJECT_NAME, TABLE_TYPE, COMMENTS", + "o.OBJECT_NAME", + "ORDER BY OBJECT_NAME", + owner, + objectTypes, + constraints + ); List result = new ArrayList<>(); - try (var stmt = requireConnection().prepareStatement(sql)) { - stmt.setString(1, owner); + try (var stmt = requireConnection().prepareStatement(query.sql)) { + bind(stmt, query.args); try (ResultSet rs = stmt.executeQuery()) { while (rs.next()) { result.add(new TableInfo(rs.getString(1), rs.getString(2), rs.getString(3))); } } } - return result; + return constraints.withoutPaging().filterTables(result); }); } @Override public List listObjects(String schema) { + return queryObjects(schema, MetadataListConstraints.NONE); + } + + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + return queryObjects(schema, MetadataListConstraints.orNone(constraints)); + } + + private List queryObjects(String schema, MetadataListConstraints constraints) { return unchecked(() -> { String owner = normalizeSchema(schema); - String sql = """ + List objectTypes = oceanBaseObjectTypes(constraints); + if (objectTypes.isEmpty()) { + return List.of(); + } + String baseSql = """ SELECT OBJECT_NAME, OBJECT_TYPE FROM ALL_OBJECTS - WHERE OWNER = ? AND OBJECT_TYPE IN ('TABLE', 'VIEW', 'PROCEDURE', 'FUNCTION', 'PACKAGE', 'SEQUENCE', 'SYNONYM') + WHERE OWNER = ? AND OBJECT_TYPE IN (%s) + """.stripIndent().trim(); + MetadataSql query = oceanBaseMetadataSql( + String.format(baseSql, placeholders(objectTypes.size())), + "OBJECT_NAME, OBJECT_TYPE", + "OBJECT_NAME", + """ ORDER BY CASE OBJECT_TYPE WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 @@ -123,21 +164,113 @@ public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent { WHEN 'SEQUENCE' THEN 5 ELSE 6 END, OBJECT_NAME - """.stripIndent().trim(); + """.stripIndent().trim(), + owner, + objectTypes, + constraints + ); List result = new ArrayList<>(); - try (var stmt = requireConnection().prepareStatement(sql)) { - stmt.setString(1, owner); + try (var stmt = requireConnection().prepareStatement(query.sql)) { + bind(stmt, query.args); try (ResultSet rs = stmt.executeQuery()) { while (rs.next()) { result.add(new ObjectInfo(rs.getString(1), rs.getString(2), owner, null)); } } } - return result; + return constraints.withoutPaging().filterObjects(result); }); } + private static MetadataSql oceanBaseMetadataSql( + String baseSql, + String selectList, + String nameColumn, + String orderSql, + String owner, + List objectTypes, + MetadataListConstraints constraints + ) { + List args = new ArrayList<>(); + args.add(owner); + args.addAll(objectTypes); + String sql = baseSql; + if (constraints.hasFilter()) { + sql += " AND UPPER(" + nameColumn + ") LIKE ? ESCAPE '\\'"; + args.add(constraints.fuzzyLikePattern().toUpperCase(Locale.ROOT)); + } + sql += "\n" + orderSql; + if (constraints.hasLimit()) { + // OceanBase Oracle mode is safest with the classic ordered ROWNUM wrapper for paged metadata. + int offset = constraints.getOffset() == null ? 0 : constraints.getOffset(); + sql = "SELECT " + selectList + "\nFROM (\n SELECT DBX_Q.*, ROWNUM AS DBX_RN\n FROM (\n" + + sql + + "\n ) DBX_Q\n WHERE ROWNUM <= ?\n)\nWHERE DBX_RN > ?"; + args.add(offset + constraints.getLimit()); + args.add(offset); + } else if (constraints.hasOffset()) { + sql = "SELECT " + selectList + "\nFROM (\n SELECT DBX_Q.*, ROWNUM AS DBX_RN\n FROM (\n" + + sql + + "\n ) DBX_Q\n)\nWHERE DBX_RN > ?"; + args.add(constraints.getOffset()); + } + return new MetadataSql(sql, args); + } + + private static List oceanBaseTableTypes(MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return List.of("TABLE", "VIEW"); + } + List result = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + result.add("TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + result.add("VIEW"); + } + return result; + } + + private static List oceanBaseObjectTypes(MetadataListConstraints constraints) { + List supported = List.of("TABLE", "VIEW", "PROCEDURE", "FUNCTION", "PACKAGE", "SEQUENCE", "SYNONYM"); + if (!constraints.hasObjectTypes()) { + return supported; + } + List result = new ArrayList<>(); + for (String objectType : supported) { + if (constraints.objectTypeAllowed(objectType)) { + result.add(objectType); + } + } + return result; + } + + private static String placeholders(int count) { + return String.join(", ", java.util.Collections.nCopies(count, "?")); + } + + private static void bind(java.sql.PreparedStatement stmt, List args) throws SQLException { + for (int index = 0; index < args.size(); index += 1) { + Object arg = args.get(index); + if (arg instanceof Integer) { + stmt.setInt(index + 1, (Integer) arg); + } else { + stmt.setString(index + 1, String.valueOf(arg)); + } + } + } + + private static final class MetadataSql { + private final String sql; + private final List args; + + private MetadataSql(String sql, List args) { + this.sql = sql; + this.args = args; + } + } + @Override public List getColumns(String schema, String table) { return unchecked(() -> { diff --git a/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java b/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java index 1fd9b0901a..d5f47b6807 100644 --- a/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java +++ b/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java @@ -1,9 +1,22 @@ package com.dbx.agent.oceanbaseoracle; import com.dbx.agent.ConnectParams; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.TableInfo; +import com.dbx.agent.test.TestSupport; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.List; + class OceanBaseOracleAgentTest { @Test void buildsOceanBaseJdbcUrl() { @@ -56,4 +69,131 @@ class OceanBaseOracleAgentTest { OceanBaseOracleAgent.buildUrl(params) ); } + + @Test + void constrainedListTablesUsesOceanBaseOracleMetadataSql() { + List sql = new ArrayList<>(); + OceanBaseOracleAgent agent = new OceanBaseOracleAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet( + new String[]{"OBJECT_NAME", "TABLE_TYPE", "COMMENTS"}, + new Object[][]{ + {"USER_SETTINGS", "TABLE", null} + } + ))); + + List tables = agent.listTables( + "APP", + new MetadataListConstraints("user", 1, 1, List.of("TABLE")) + ); + + Assertions.assertEquals(1, tables.size()); + Assertions.assertEquals("USER_SETTINGS", tables.get(0).getName()); + Assertions.assertTrue(sql.get(0).contains("ALL_OBJECTS"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("UPPER(o.OBJECT_NAME) LIKE ?"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("ROWNUM <= ?"), sql.get(0)); + } + + @Test + void constrainedListObjectsUsesOceanBaseOracleMetadataSql() { + List sql = new ArrayList<>(); + OceanBaseOracleAgent agent = new OceanBaseOracleAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet( + new String[]{"OBJECT_NAME", "OBJECT_TYPE"}, + new Object[][]{ + {"FORMAT_USER", "FUNCTION"} + } + ))); + + List objects = agent.listObjects( + "APP", + new MetadataListConstraints("user", 1, 1, List.of("FUNCTION")) + ); + + Assertions.assertEquals(1, objects.size()); + Assertions.assertEquals("FORMAT_USER", objects.get(0).getName()); + Assertions.assertEquals("FUNCTION", objects.get(0).getObject_type()); + Assertions.assertTrue(sql.get(0).contains("OBJECT_TYPE IN (?)"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("ROWNUM <= ?"), sql.get(0)); + } + + private static Connection preparedConnection(List sql, ResultSet... resultSets) { + int[] resultSetIndex = {0}; + PreparedStatement statement = proxy(PreparedStatement.class, (method, args) -> { + if ("executeQuery".equals(method.getName())) { + int current = Math.min(resultSetIndex[0], resultSets.length - 1); + resultSetIndex[0] += 1; + return resultSets[current]; + } + if ("setString".equals(method.getName()) || "setInt".equals(method.getName()) || "close".equals(method.getName())) { + return null; + } + return defaultValue(method.getReturnType()); + }); + return proxy(Connection.class, (method, args) -> { + if ("prepareStatement".equals(method.getName())) { + sql.add(String.valueOf(args[0])); + return statement; + } + if ("isClosed".equals(method.getName())) { + return false; + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSet resultSet(String[] columns, Object[][] rows) { + int[] index = {-1}; + return proxy(ResultSet.class, (method, args) -> { + switch (method.getName()) { + case "next": + index[0] += 1; + return index[0] < rows.length; + case "getString": + Object value = columnValue(columns, rows[index[0]], args[0]); + return value == null ? null : String.valueOf(value); + case "close": + return null; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static Object columnValue(String[] columns, Object[] row, Object key) { + if (key instanceof Number) { + return row[((Number) key).intValue() - 1]; + } + for (int i = 0; i < columns.length; i++) { + if (columns[i].equalsIgnoreCase(String.valueOf(key))) { + return row[i]; + } + } + return null; + } + + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + return handler.handle(method, args == null ? new Object[0] : args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (type == Boolean.TYPE) return false; + if (type == Byte.TYPE) return (byte) 0; + if (type == Short.TYPE) return (short) 0; + if (type == Integer.TYPE) return 0; + if (type == Long.TYPE) return 0L; + if (type == Float.TYPE) return 0f; + if (type == Double.TYPE) return 0d; + if (type == Character.TYPE) return (char) 0; + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args) throws Throwable; + } } diff --git a/agents/drivers/oracle-go/main.go b/agents/drivers/oracle-go/main.go index 440ecac439..1f5b2020b8 100644 --- a/agents/drivers/oracle-go/main.go +++ b/agents/drivers/oracle-go/main.go @@ -11,6 +11,7 @@ import ( "net/url" "os" "regexp" + "sort" "strconv" "strings" "time" @@ -42,7 +43,7 @@ ORDER BY CASE WHEN username = SYS_CONTEXT('USERENV', 'SESSION_USER') THEN 1 ELSE 2 END, username` -const oracleListTablesSQL = ` +const oracleListTablesBaseSQL = ` SELECT OBJECT_NAME, TABLE_TYPE, COMMENTS FROM ( SELECT t.TABLE_NAME AS OBJECT_NAME, @@ -58,9 +59,10 @@ SELECT o.OBJECT_NAME, FROM ALL_OBJECTS o WHERE o.OWNER = :2 AND o.OBJECT_TYPE = 'VIEW' -) -ORDER BY OBJECT_NAME` -const oracleListObjectsSQL = ` +)` +const oracleListTablesOrderSQL = `ORDER BY OBJECT_NAME` +const oracleListTablesSQL = oracleListTablesBaseSQL + "\n" + oracleListTablesOrderSQL +const oracleListObjectsBaseSQL = ` SELECT OBJECT_NAME, OBJECT_TYPE, COMMENTS FROM ( SELECT t.TABLE_NAME AS OBJECT_NAME, @@ -71,19 +73,21 @@ WHERE t.OWNER = :1 AND t.NESTED = 'NO' UNION ALL SELECT o.OBJECT_NAME, - o.OBJECT_TYPE, + CASE o.OBJECT_TYPE WHEN 'PACKAGE BODY' THEN 'PACKAGE_BODY' ELSE o.OBJECT_TYPE END AS OBJECT_TYPE, CAST(NULL AS VARCHAR2(4000)) AS COMMENTS FROM ALL_OBJECTS o WHERE o.OWNER = :2 - AND o.OBJECT_TYPE IN ('VIEW', 'PROCEDURE', 'FUNCTION', 'PACKAGE') -) -ORDER BY CASE OBJECT_TYPE + AND o.OBJECT_TYPE IN ('VIEW', 'PROCEDURE', 'FUNCTION', 'PACKAGE', 'PACKAGE BODY') +)` +const oracleListObjectsOrderSQL = `ORDER BY CASE OBJECT_TYPE WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'PROCEDURE' THEN 2 WHEN 'FUNCTION' THEN 3 - ELSE 4 + WHEN 'PACKAGE' THEN 4 + ELSE 5 END, OBJECT_NAME` +const oracleListObjectsSQL = oracleListObjectsBaseSQL + "\n" + oracleListObjectsOrderSQL type request struct { ID json.RawMessage `json:"id"` @@ -188,6 +192,13 @@ type tableInfo struct { Comment *string `json:"comment"` } +type metadataListConstraints struct { + Filter string + Limit int + Offset int + ObjectTypes []string +} + type objectInfo struct { Name string `json:"name"` ObjectType string `json:"object_type"` @@ -336,11 +347,11 @@ func (s *server) dispatch(method string, params map[string]json.RawMessage) (any return result, false, err case "list_tables": schema := stringParam(params, "schema") - result, err := s.listTables(schema) + result, err := s.listTables(schema, metadataListConstraintsFromParams(params)) return result, false, err case "list_objects": schema := stringParam(params, "schema") - result, err := s.listObjects(schema) + result, err := s.listObjects(schema, metadataListConstraintsFromParams(params)) return result, false, err case "get_columns": schema := stringParam(params, "schema") @@ -675,12 +686,145 @@ func (s *server) normalizeSchema(schema string) (string, error) { return strings.ToUpper(schema), nil } -func (s *server) listTables(schema string) ([]tableInfo, error) { +type oracleMetadataListQuery struct { + SQL string + Args []any +} + +func metadataListConstraintsFromParams(params map[string]json.RawMessage) metadataListConstraints { + objectTypes := stringSliceParam(params, "object_types") + if len(objectTypes) == 0 { + objectTypes = stringSliceParam(params, "objectTypes") + } + limit := intParam(params, "limit") + offset := intParam(params, "offset") + if limit < 0 { + limit = 0 + } + if offset < 0 { + offset = 0 + } + return metadataListConstraints{ + Filter: stringParam(params, "filter"), + Limit: limit, + Offset: offset, + ObjectTypes: objectTypes, + } +} + +func oracleListTablesQuery(schema string, constraints metadataListConstraints) oracleMetadataListQuery { + return oracleConstrainedMetadataListQuery( + oracleListTablesBaseSQL, + "OBJECT_NAME, TABLE_TYPE, COMMENTS", + "TABLE_TYPE", + oracleListTablesOrderSQL, + []any{schema, schema}, + constraints, + ) +} + +func oracleListObjectsQuery(schema string, constraints metadataListConstraints) oracleMetadataListQuery { + return oracleConstrainedMetadataListQuery( + oracleListObjectsBaseSQL, + "OBJECT_NAME, OBJECT_TYPE, COMMENTS", + "OBJECT_TYPE", + oracleListObjectsOrderSQL, + []any{schema, schema}, + constraints, + ) +} + +func oracleConstrainedMetadataListQuery(baseSQL, selectList, typeColumn, orderSQL string, baseArgs []any, constraints metadataListConstraints) oracleMetadataListQuery { + args := append([]any{}, baseArgs...) + where := make([]string, 0, 2) + if filter := strings.TrimSpace(constraints.Filter); filter != "" { + args = append(args, strings.ToUpper(oracleFuzzyLikePattern(filter))) + where = append(where, fmt.Sprintf("UPPER(OBJECT_NAME) LIKE :%d ESCAPE '\\'", len(args))) + } + if objectTypes := normalizedMetadataObjectTypes(constraints.ObjectTypes); len(objectTypes) > 0 { + placeholders := make([]string, 0, len(objectTypes)) + for _, objectType := range objectTypes { + args = append(args, objectType) + placeholders = append(placeholders, fmt.Sprintf(":%d", len(args))) + } + where = append(where, fmt.Sprintf("%s IN (%s)", typeColumn, strings.Join(placeholders, ","))) + } + + sqlText := fmt.Sprintf("SELECT %s\nFROM (\n%s\n)", selectList, baseSQL) + if len(where) > 0 { + sqlText += "\nWHERE " + strings.Join(where, " AND ") + } + sqlText += "\n" + orderSQL + + if constraints.Limit > 0 { + args = append(args, constraints.Offset+constraints.Limit) + maxRowParam := len(args) + args = append(args, constraints.Offset) + offsetParam := len(args) + sqlText = fmt.Sprintf( + "SELECT %s\nFROM (\n SELECT DBX_Q.*, ROWNUM AS DBX_RN\n FROM (\n%s\n ) DBX_Q\n WHERE ROWNUM <= :%d\n)\nWHERE DBX_RN > :%d", + selectList, + sqlText, + maxRowParam, + offsetParam, + ) + } else if constraints.Offset > 0 { + args = append(args, constraints.Offset) + offsetParam := len(args) + sqlText = fmt.Sprintf( + "SELECT %s\nFROM (\n SELECT DBX_Q.*, ROWNUM AS DBX_RN\n FROM (\n%s\n ) DBX_Q\n)\nWHERE DBX_RN > :%d", + selectList, + sqlText, + offsetParam, + ) + } + + return oracleMetadataListQuery{SQL: sqlText, Args: args} +} + +func normalizedMetadataObjectTypes(values []string) []string { + seen := map[string]bool{} + result := make([]string, 0, len(values)) + for _, value := range values { + normalized := strings.ToUpper(strings.TrimSpace(value)) + normalized = strings.ReplaceAll(normalized, "-", "_") + normalized = strings.ReplaceAll(normalized, " ", "_") + if normalized == "" || seen[normalized] { + continue + } + seen[normalized] = true + result = append(result, normalized) + } + sort.Strings(result) + return result +} + +func oracleFuzzyLikePattern(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "%%" + } + var builder strings.Builder + builder.Grow(len(value)*2 + 2) + builder.WriteByte('%') + for _, ch := range value { + switch ch { + case '\\', '%', '_': + builder.WriteByte('\\') + } + builder.WriteRune(ch) + builder.WriteByte('%') + } + return builder.String() +} + +func (s *server) listTables(schema string, constraints metadataListConstraints) ([]tableInfo, error) { schema, err := s.normalizeSchema(schema) if err != nil { return nil, err } - rows, err := s.queryRows(oracleListTablesSQL, []any{schema, schema}) + query := oracleListTablesQuery(schema, constraints) + rows, err := s.queryRows(query.SQL, query.Args) if err != nil { if isOraclePGALimitError(err) { return []tableInfo{}, nil @@ -699,12 +843,13 @@ func (s *server) listTables(schema string) ([]tableInfo, error) { return emptyIfNil(result), rows.Err() } -func (s *server) listObjects(schema string) ([]objectInfo, error) { +func (s *server) listObjects(schema string, constraints metadataListConstraints) ([]objectInfo, error) { schema, err := s.normalizeSchema(schema) if err != nil { return nil, err } - rows, err := s.queryRows(oracleListObjectsSQL, []any{schema, schema}) + query := oracleListObjectsQuery(schema, constraints) + rows, err := s.queryRows(query.SQL, query.Args) if err != nil { if isOraclePGALimitError(err) { return []objectInfo{}, nil diff --git a/agents/drivers/oracle-go/main_test.go b/agents/drivers/oracle-go/main_test.go index 7c5d9e8760..0187593455 100644 --- a/agents/drivers/oracle-go/main_test.go +++ b/agents/drivers/oracle-go/main_test.go @@ -358,6 +358,32 @@ func TestListTablesSQLUsesSplitDictionaryQuery(t *testing.T) { } } +func TestListTablesQueryAppliesMetadataConstraints(t *testing.T) { + query := oracleListTablesQuery("APP", metadataListConstraints{ + Filter: "u_r", + Limit: 501, + Offset: 10, + ObjectTypes: []string{"view", "TABLE", "TABLE"}, + }) + sqlText := strings.ToUpper(query.SQL) + + if !strings.Contains(sqlText, "UPPER(OBJECT_NAME) LIKE :3 ESCAPE '\\'") { + t.Fatalf("table listing should push filter predicate, got: %s", query.SQL) + } + if !strings.Contains(sqlText, "TABLE_TYPE IN (:4,:5)") { + t.Fatalf("table listing should push table type predicate, got: %s", query.SQL) + } + if !strings.Contains(sqlText, "ROWNUM <= :6") || !strings.Contains(sqlText, "DBX_RN > :7") { + t.Fatalf("table listing should use rownum pagination, got: %s", query.SQL) + } + if len(query.Args) != 7 { + t.Fatalf("unexpected args: %#v", query.Args) + } + if query.Args[0] != "APP" || query.Args[1] != "APP" || query.Args[2] != "%U%\\_%R%" || query.Args[3] != "TABLE" || query.Args[4] != "VIEW" || query.Args[5] != 511 || query.Args[6] != 10 { + t.Fatalf("constraints args were not normalized: %#v", query.Args) + } +} + func TestListObjectsSQLUsesSplitDictionaryQuery(t *testing.T) { sqlText := strings.ToUpper(oracleListObjectsSQL) @@ -370,6 +396,42 @@ func TestListObjectsSQLUsesSplitDictionaryQuery(t *testing.T) { if strings.Contains(sqlText, "ALL_TAB_COMMENTS") { t.Fatalf("object listing should not load comments during refresh, got: %s", oracleListObjectsSQL) } + if !strings.Contains(sqlText, "'PACKAGE BODY'") || !strings.Contains(sqlText, "PACKAGE_BODY") { + t.Fatalf("object listing should include package bodies with normalized type, got: %s", oracleListObjectsSQL) + } +} + +func TestListObjectsQueryAppliesMetadataConstraints(t *testing.T) { + query := oracleListObjectsQuery("APP", metadataListConstraints{ + Filter: "pkg%", + Limit: 25, + ObjectTypes: []string{"FUNCTION", "package"}, + }) + sqlText := strings.ToUpper(query.SQL) + + if !strings.Contains(sqlText, "UPPER(OBJECT_NAME) LIKE :3 ESCAPE '\\'") { + t.Fatalf("object listing should push filter predicate, got: %s", query.SQL) + } + if !strings.Contains(sqlText, "OBJECT_TYPE IN (:4,:5)") { + t.Fatalf("object listing should push object type predicate, got: %s", query.SQL) + } + if !strings.Contains(sqlText, "ROWNUM <= :6") || !strings.Contains(sqlText, "DBX_RN > :7") { + t.Fatalf("object listing should use rownum pagination, got: %s", query.SQL) + } + if len(query.Args) != 7 { + t.Fatalf("unexpected args: %#v", query.Args) + } + if query.Args[2] != "%P%K%G%\\%%" || query.Args[3] != "FUNCTION" || query.Args[4] != "PACKAGE" || query.Args[5] != 25 || query.Args[6] != 0 { + t.Fatalf("object constraints args were not normalized: %#v", query.Args) + } +} + +func TestOracleFuzzyLikePatternEscapesSpecialCharacters(t *testing.T) { + got := oracleFuzzyLikePattern(`a_%\b`) + want := `%a%\_%\%%\\%b%` + if got != want { + t.Fatalf("oracleFuzzyLikePattern() = %q, want %q", got, want) + } } func TestIsOraclePGALimitError(t *testing.T) { diff --git a/agents/drivers/snowflake/src/main/java/com/dbx/agent/snowflake/SnowflakeAgent.java b/agents/drivers/snowflake/src/main/java/com/dbx/agent/snowflake/SnowflakeAgent.java index da11603db8..608ff57b4c 100644 --- a/agents/drivers/snowflake/src/main/java/com/dbx/agent/snowflake/SnowflakeAgent.java +++ b/agents/drivers/snowflake/src/main/java/com/dbx/agent/snowflake/SnowflakeAgent.java @@ -8,6 +8,8 @@ import com.dbx.agent.ForeignKeyInfo; import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcIdentifiers; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.MetadataSqlSupport; import com.dbx.agent.ObjectInfo; import com.dbx.agent.ObjectSource; import com.dbx.agent.TableInfo; @@ -86,6 +88,48 @@ public final class SnowflakeAgent extends AbstractJdbcAgent { }); } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + if (!normalized.includesTableLikeTypes()) { + return List.of(); + } + try { + return queryConstrainedTables(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterTables(listTables(schema)); + } + } + + private List queryConstrainedTables(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT TABLE_NAME, TABLE_TYPE FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ?"); + args.add(schema); + appendSnowflakeTableTypePredicate(sql, args, constraints); + MetadataSqlSupport.appendNameFilter(sql, args, "TABLE_NAME", constraints); + sql.append(" ORDER BY TABLE_NAME"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo( + rs.getString("TABLE_NAME"), + normalizeTableType(rs.getString("TABLE_TYPE")), + null + )); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } + @Override public List listObjects(String schema) { return unchecked(() -> { @@ -110,6 +154,64 @@ public final class SnowflakeAgent extends AbstractJdbcAgent { }); } + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listObjects(schema); + } + if (!includesSupportedObjects(normalized)) { + return List.of(); + } + try { + return queryConstrainedObjects(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterObjects(listObjects(schema)); + } + } + + private List queryConstrainedObjects(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List branches = new ArrayList<>(); + List args = new ArrayList<>(); + if (constraints.includesTableLikeTypes()) { + StringBuilder tableSql = new StringBuilder("SELECT TABLE_NAME AS OBJECT_NAME, TABLE_TYPE AS OBJECT_TYPE FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ?"); + args.add(schema); + appendSnowflakeTableTypePredicate(tableSql, args, constraints); + MetadataSqlSupport.appendNameFilter(tableSql, args, "TABLE_NAME", constraints); + branches.add(tableSql.toString()); + } + if (constraints.objectTypeAllowed("PROCEDURE")) { + StringBuilder procedureSql = new StringBuilder("SELECT PROCEDURE_NAME AS OBJECT_NAME, 'PROCEDURE' AS OBJECT_TYPE FROM INFORMATION_SCHEMA.PROCEDURES WHERE PROCEDURE_SCHEMA = ?"); + args.add(schema); + MetadataSqlSupport.appendNameFilter(procedureSql, args, "PROCEDURE_NAME", constraints); + branches.add(procedureSql.toString()); + } + if (branches.isEmpty()) { + return List.of(); + } + StringBuilder sql = new StringBuilder("SELECT OBJECT_NAME, OBJECT_TYPE FROM (") + .append(String.join(" UNION ALL ", branches)) + .append(") metadata_objects ORDER BY CASE OBJECT_TYPE WHEN 'BASE TABLE' THEN 0 WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'MATERIALIZED VIEW' THEN 2 WHEN 'PROCEDURE' THEN 3 ELSE 9 END, OBJECT_NAME"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo( + rs.getString("OBJECT_NAME"), + normalizeTableType(rs.getString("OBJECT_TYPE")), + schema, + null + )); + } + } + } + return constraints.withoutPaging().filterObjects(result); + }); + } + @Override public ObjectSource getObjectSource(String schema, String name, String objectType) { return unchecked(() -> { @@ -224,8 +326,44 @@ public final class SnowflakeAgent extends AbstractJdbcAgent { return "jdbc:snowflake://" + params.getHost() + ":" + params.getPort() + "/?db=" + params.getDatabase(); } + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() && !constraints.hasLimit() && !constraints.hasOffset() && !constraints.hasObjectTypes(); + } + + private static boolean includesSupportedObjects(MetadataListConstraints constraints) { + return constraints.includesTableLikeTypes() || constraints.objectTypeAllowed("PROCEDURE"); + } + + private static void appendSnowflakeTableTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List types = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + types.add("BASE TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + types.add("VIEW"); + } + if (constraints.tableTypeAllowed("MATERIALIZED_VIEW")) { + types.add("MATERIALIZED VIEW"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND TABLE_TYPE IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + private static String normalizeTableType(String tableType) { - return "BASE TABLE".equals(tableType) ? "TABLE" : tableType; + if ("BASE TABLE".equals(tableType)) { + return "TABLE"; + } + if ("MATERIALIZED VIEW".equals(tableType)) { + return "MATERIALIZED_VIEW"; + } + return tableType; } private static String blankToNull(String value) { diff --git a/agents/drivers/snowflake/src/test/java/com/dbx/agent/snowflake/SnowflakeAgentMetadataTest.java b/agents/drivers/snowflake/src/test/java/com/dbx/agent/snowflake/SnowflakeAgentMetadataTest.java index 847d48878d..8708a7f2e6 100644 --- a/agents/drivers/snowflake/src/test/java/com/dbx/agent/snowflake/SnowflakeAgentMetadataTest.java +++ b/agents/drivers/snowflake/src/test/java/com/dbx/agent/snowflake/SnowflakeAgentMetadataTest.java @@ -1,8 +1,10 @@ package com.dbx.agent.snowflake; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.test.TestSupport; import com.dbx.agent.test.JdbcAgentFake; import com.dbx.agent.test.JdbcMetadataSqlFake; +import java.util.List; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -26,4 +28,32 @@ class SnowflakeAgentMetadataTest { ) ); } + + @Test + void constrainedTableMetadataPushesFilterTypesAndPaging() { + SnowflakeAgent agent = new SnowflakeAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listTables("PUBLIC", new MetadataListConstraints("ord", 25, 50, List.of("TABLE", "VIEW"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + Assertions.assertTrue(sql.contains("FROM INFORMATION_SCHEMA.TABLES"), sql); + Assertions.assertTrue(sql.contains("TABLE_TYPE IN (?, ?)"), sql); + Assertions.assertTrue(sql.contains("UPPER(TABLE_NAME) LIKE ? ESCAPE '\\\\'"), sql); + Assertions.assertTrue(sql.endsWith("LIMIT 25 OFFSET 50"), sql); + } + + @Test + void constrainedObjectMetadataPushesProcedureSearchAndPaging() { + SnowflakeAgent agent = new SnowflakeAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listObjects("PUBLIC", new MetadataListConstraints("sync", 10, null, List.of("PROCEDURE"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + Assertions.assertTrue(sql.contains("FROM INFORMATION_SCHEMA.PROCEDURES"), sql); + Assertions.assertTrue(sql.contains("UPPER(PROCEDURE_NAME) LIKE ? ESCAPE '\\\\'"), sql); + Assertions.assertTrue(sql.contains("ORDER BY CASE OBJECT_TYPE"), sql); + Assertions.assertTrue(sql.endsWith("LIMIT 10"), sql); + } } diff --git a/agents/drivers/sundb/src/main/java/com/dbx/agent/sundb/SundbAgent.java b/agents/drivers/sundb/src/main/java/com/dbx/agent/sundb/SundbAgent.java index 9a23eb897b..e36620f86c 100644 --- a/agents/drivers/sundb/src/main/java/com/dbx/agent/sundb/SundbAgent.java +++ b/agents/drivers/sundb/src/main/java/com/dbx/agent/sundb/SundbAgent.java @@ -10,6 +10,8 @@ import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcExecutor; import com.dbx.agent.JdbcIdentifiers; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.MetadataSqlSupport; import com.dbx.agent.ObjectInfo; import com.dbx.agent.ObjectSource; import com.dbx.agent.QueryResult; @@ -103,6 +105,44 @@ public final class SundbAgent extends BaseDatabaseAgent { }); } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + if (!normalized.includesTableLikeTypes()) { + return List.of(); + } + try { + return queryConstrainedTables(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterTables(listTables(schema)); + } + } + + private List queryConstrainedTables(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA = ?"); + args.add(schema); + appendInformationSchemaTableTypePredicate(sql, args, constraints); + MetadataSqlSupport.appendNameFilter(sql, args, "TABLE_NAME", constraints); + sql.append(" ORDER BY TABLE_NAME"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (var stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString("TABLE_NAME"), normalizeTableType(rs.getString("TABLE_TYPE")), null)); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } + @Override public List listObjects(String schema) { return unchecked(() -> { @@ -125,6 +165,60 @@ public final class SundbAgent extends BaseDatabaseAgent { }); } + @Override + public List listObjects(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listObjects(schema); + } + if (!includesSupportedObjects(normalized)) { + return List.of(); + } + try { + return queryConstrainedObjects(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterObjects(listObjects(schema)); + } + } + + private List queryConstrainedObjects(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List branches = new ArrayList<>(); + List args = new ArrayList<>(); + if (constraints.includesTableLikeTypes()) { + StringBuilder tableSql = new StringBuilder("SELECT TABLE_NAME AS OBJECT_NAME, TABLE_TYPE AS OBJECT_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA = ?"); + args.add(schema); + appendInformationSchemaTableTypePredicate(tableSql, args, constraints); + MetadataSqlSupport.appendNameFilter(tableSql, args, "TABLE_NAME", constraints); + branches.add(tableSql.toString()); + } + if (constraints.objectTypeAllowed("PROCEDURE") || constraints.objectTypeAllowed("FUNCTION")) { + StringBuilder routineSql = new StringBuilder("SELECT ROUTINE_NAME AS OBJECT_NAME, ROUTINE_TYPE AS OBJECT_TYPE FROM information_schema.ROUTINES WHERE ROUTINE_SCHEMA = ?"); + args.add(schema); + appendRoutineTypePredicate(routineSql, args, constraints); + MetadataSqlSupport.appendNameFilter(routineSql, args, "ROUTINE_NAME", constraints); + branches.add(routineSql.toString()); + } + if (branches.isEmpty()) { + return List.of(); + } + StringBuilder sql = new StringBuilder("SELECT OBJECT_NAME, OBJECT_TYPE FROM (") + .append(String.join(" UNION ALL ", branches)) + .append(") metadata_objects ORDER BY CASE OBJECT_TYPE WHEN 'BASE TABLE' THEN 0 WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'PROCEDURE' THEN 2 WHEN 'FUNCTION' THEN 3 ELSE 9 END, OBJECT_NAME"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (var stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), normalizeTableType(rs.getString(2)), schema, null)); + } + } + } + return constraints.withoutPaging().filterObjects(result); + }); + } + @Override public ObjectSource getObjectSource(String schema, String name, String objectType) { return unchecked(() -> { @@ -344,6 +438,58 @@ public final class SundbAgent extends BaseDatabaseAgent { return "jdbc:sundb://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); } + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() && !constraints.hasLimit() && !constraints.hasOffset() && !constraints.hasObjectTypes(); + } + + private static boolean includesSupportedObjects(MetadataListConstraints constraints) { + return constraints.includesTableLikeTypes() + || constraints.objectTypeAllowed("PROCEDURE") + || constraints.objectTypeAllowed("FUNCTION"); + } + + private static void appendInformationSchemaTableTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List types = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + types.add("BASE TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + types.add("VIEW"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND TABLE_TYPE IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + + private static void appendRoutineTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List types = new ArrayList<>(); + if (constraints.objectTypeAllowed("PROCEDURE")) { + types.add("PROCEDURE"); + } + if (constraints.objectTypeAllowed("FUNCTION")) { + types.add("FUNCTION"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND ROUTINE_TYPE IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + + private static String normalizeTableType(String tableType) { + return "BASE TABLE".equals(tableType) ? "TABLE" : tableType; + } + private static Integer integerOrNull(ResultSet rs, String column) throws Exception { Object value = rs.getObject(column); return value instanceof Integer ? (Integer) value : null; diff --git a/agents/drivers/sundb/src/test/java/com/dbx/agent/sundb/SundbAgentMetadataTest.java b/agents/drivers/sundb/src/test/java/com/dbx/agent/sundb/SundbAgentMetadataTest.java index 943e7953f2..2fd7c91f9f 100644 --- a/agents/drivers/sundb/src/test/java/com/dbx/agent/sundb/SundbAgentMetadataTest.java +++ b/agents/drivers/sundb/src/test/java/com/dbx/agent/sundb/SundbAgentMetadataTest.java @@ -1,9 +1,11 @@ package com.dbx.agent.sundb; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.test.TestSupport; import com.dbx.agent.test.JdbcAgentFake; import com.dbx.agent.test.JdbcMetadataSqlFake; import java.util.Collections; +import java.util.List; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -20,4 +22,32 @@ class SundbAgentMetadataTest { JdbcMetadataSqlFake.statements ); } + + @Test + void constrainedTableMetadataPushesFilterTypesAndPaging() { + SundbAgent agent = new SundbAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listTables("app", new MetadataListConstraints("ord", 25, 50, List.of("TABLE", "VIEW"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + Assertions.assertTrue(sql.contains("FROM information_schema.TABLES"), sql); + Assertions.assertTrue(sql.contains("TABLE_TYPE IN (?, ?)"), sql); + Assertions.assertTrue(sql.contains("UPPER(TABLE_NAME) LIKE ? ESCAPE '\\\\'"), sql); + Assertions.assertTrue(sql.endsWith("LIMIT 25 OFFSET 50"), sql); + } + + @Test + void constrainedObjectMetadataPushesRoutineTypesAndPaging() { + SundbAgent agent = new SundbAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listObjects("app", new MetadataListConstraints("sync", 10, null, List.of("PROCEDURE", "FUNCTION"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + Assertions.assertTrue(sql.contains("FROM information_schema.ROUTINES"), sql); + Assertions.assertTrue(sql.contains("ROUTINE_TYPE IN (?, ?)"), sql); + Assertions.assertTrue(sql.contains("ORDER BY CASE OBJECT_TYPE"), sql); + Assertions.assertTrue(sql.endsWith("LIMIT 10"), sql); + } } diff --git a/agents/drivers/trino/src/main/java/com/dbx/agent/trino/TrinoAgent.java b/agents/drivers/trino/src/main/java/com/dbx/agent/trino/TrinoAgent.java index 6ffdad8ace..e0e348a79f 100644 --- a/agents/drivers/trino/src/main/java/com/dbx/agent/trino/TrinoAgent.java +++ b/agents/drivers/trino/src/main/java/com/dbx/agent/trino/TrinoAgent.java @@ -8,6 +8,8 @@ import com.dbx.agent.ForeignKeyInfo; import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcIdentifiers; import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.MetadataListConstraints; +import com.dbx.agent.MetadataSqlSupport; import com.dbx.agent.TableInfo; import com.dbx.agent.TriggerInfo; import java.sql.Connection; @@ -87,6 +89,48 @@ public final class TrinoAgent extends AbstractJdbcAgent { }); } + @Override + public List listTables(String schema, MetadataListConstraints constraints) { + MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints); + if (isUnconstrained(normalized)) { + return listTables(schema); + } + if (!normalized.includesTableLikeTypes()) { + return List.of(); + } + try { + return queryConstrainedTables(schema, normalized); + } catch (RuntimeException e) { + return normalized.filterTables(listTables(schema)); + } + } + + private List queryConstrainedTables(String schema, MetadataListConstraints constraints) { + return unchecked(() -> { + List result = new ArrayList<>(); + List args = new ArrayList<>(); + StringBuilder sql = new StringBuilder("SELECT table_name, table_type FROM information_schema.tables WHERE table_schema = ?"); + args.add(schema); + appendTrinoTableTypePredicate(sql, args, constraints); + MetadataSqlSupport.appendNameFilter(sql, args, "table_name", constraints); + sql.append(" ORDER BY table_name"); + MetadataSqlSupport.appendLiteralLimitOffset(sql, constraints); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) { + MetadataSqlSupport.bind(stmt, args); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo( + rs.getString("table_name"), + normalizeTableType(rs.getString("table_type")), + null + )); + } + } + } + return constraints.withoutPaging().filterTables(result); + }); + } + @Override public List getColumns(String schema, String table) { Connection conn = requireConnected(); @@ -188,6 +232,32 @@ public final class TrinoAgent extends AbstractJdbcAgent { return "jdbc:trino://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); } + private static boolean isUnconstrained(MetadataListConstraints constraints) { + return !constraints.hasFilter() && !constraints.hasLimit() && !constraints.hasOffset() && !constraints.hasObjectTypes(); + } + + private static void appendTrinoTableTypePredicate(StringBuilder sql, List args, MetadataListConstraints constraints) { + if (!constraints.hasObjectTypes()) { + return; + } + List types = new ArrayList<>(); + if (constraints.tableTypeAllowed("TABLE")) { + types.add("BASE TABLE"); + } + if (constraints.tableTypeAllowed("VIEW")) { + types.add("VIEW"); + } + if (constraints.tableTypeAllowed("MATERIALIZED_VIEW")) { + types.add("MATERIALIZED VIEW"); + } + if (types.isEmpty()) { + sql.append(" AND 1 = 0"); + return; + } + sql.append(" AND table_type IN (").append(MetadataSqlSupport.placeholders(types.size())).append(")"); + args.addAll(types); + } + private static String quoteIdentifier(String identifier) { return "\"" + identifier.replace("\"", "\"\"") + "\""; } @@ -255,6 +325,9 @@ public final class TrinoAgent extends AbstractJdbcAgent { if ("BASE TABLE".equals(type)) { return "TABLE"; } + if ("MATERIALIZED VIEW".equals(type)) { + return "MATERIALIZED_VIEW"; + } return type; } diff --git a/agents/drivers/trino/src/test/java/com/dbx/agent/trino/TrinoAgentTest.java b/agents/drivers/trino/src/test/java/com/dbx/agent/trino/TrinoAgentTest.java index f1d94db0f6..9b8943d4aa 100644 --- a/agents/drivers/trino/src/test/java/com/dbx/agent/trino/TrinoAgentTest.java +++ b/agents/drivers/trino/src/test/java/com/dbx/agent/trino/TrinoAgentTest.java @@ -1,7 +1,14 @@ package com.dbx.agent.trino; import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.MetadataListConstraints; import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import com.dbx.agent.test.TestSupport; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.List; class TrinoAgentTest extends JdbcFakeExecutionBehaviorTest { @Override @@ -13,4 +20,18 @@ class TrinoAgentTest extends JdbcFakeExecutionBehaviorTest { protected String resultSetSql() { return "CALL system.runtime.nodes()"; } + + @Test + void constrainedTableMetadataPushesFilterTypesAndPaging() { + TrinoAgent agent = new TrinoAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listTables("public", new MetadataListConstraints("ord", 25, 50, List.of("TABLE", "VIEW"))); + + String sql = JdbcMetadataSqlFake.statements.get(0); + Assertions.assertTrue(sql.contains("FROM information_schema.tables"), sql); + Assertions.assertTrue(sql.contains("table_type IN (?, ?)"), sql); + Assertions.assertTrue(sql.contains("UPPER(table_name) LIKE ? ESCAPE '\\\\'"), sql); + Assertions.assertTrue(sql.endsWith("LIMIT 25 OFFSET 50"), sql); + } } diff --git a/agents/drivers/xugu/main.go b/agents/drivers/xugu/main.go index d0eb7ee693..3ee35efd91 100644 --- a/agents/drivers/xugu/main.go +++ b/agents/drivers/xugu/main.go @@ -10,6 +10,7 @@ import ( "io" "os" "regexp" + "sort" "strconv" "strings" "time" @@ -196,6 +197,18 @@ type objectInfo struct { Comment *string `json:"comment"` } +type metadataListConstraints struct { + Filter string + Limit int + Offset int + ObjectTypes []string +} + +type xuguMetadataListQuery struct { + SQL string + Args []any +} + type columnInfo struct { Name string `json:"name"` DataType string `json:"data_type"` @@ -341,14 +354,14 @@ func (s *server) dispatch(method string, params map[string]json.RawMessage) (any return nil, false, err } schema := stringParam(params, "schema") - result, err := s.listTables(schema) + result, err := s.listTables(schema, metadataListConstraintsFromParams(params)) return result, false, err case "list_objects": if err := s.useDatabase(stringParam(params, "database")); err != nil { return nil, false, err } schema := stringParam(params, "schema") - result, err := s.listObjects(schema) + result, err := s.listObjects(schema, metadataListConstraintsFromParams(params)) return result, false, err case "list_data_types": return xuguDataTypes, false, nil @@ -793,22 +806,13 @@ func (s *server) normalizeSchema(schema string) (string, error) { return strings.ToUpper(schema), nil } -func (s *server) listTables(schema string) ([]tableInfo, error) { +func (s *server) listTables(schema string, constraints metadataListConstraints) ([]tableInfo, error) { schema, err := s.normalizeSchema(schema) if err != nil { return nil, err } - rows, err := s.queryRows(` -SELECT t.TABLE_NAME, 'TABLE' AS TABLE_TYPE, t.COMMENTS -FROM ALL_TABLES t -JOIN ALL_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID -WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) -UNION ALL -SELECT v.VIEW_NAME, 'VIEW' AS TABLE_TYPE, v.COMMENTS -FROM ALL_VIEWS v -JOIN ALL_SCHEMAS s ON s.DB_ID = v.DB_ID AND s.SCHEMA_ID = v.SCHEMA_ID -WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) -ORDER BY 2, 1`, []any{schema, schema}) + query := xuguListTablesQuery(schema, constraints) + rows, err := s.queryRows(query.SQL, query.Args) if err != nil { return nil, err } @@ -824,22 +828,13 @@ ORDER BY 2, 1`, []any{schema, schema}) return emptyIfNil(result), rows.Err() } -func (s *server) listObjects(schema string) ([]objectInfo, error) { +func (s *server) listObjects(schema string, constraints metadataListConstraints) ([]objectInfo, error) { schema, err := s.normalizeSchema(schema) if err != nil { return nil, err } - rows, err := s.queryRows(` -SELECT t.TABLE_NAME, 'TABLE' AS OBJECT_TYPE, t.COMMENTS -FROM ALL_TABLES t -JOIN ALL_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID -WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) -UNION ALL -SELECT v.VIEW_NAME, 'VIEW' AS OBJECT_TYPE, v.COMMENTS -FROM ALL_VIEWS v -JOIN ALL_SCHEMAS s ON s.DB_ID = v.DB_ID AND s.SCHEMA_ID = v.SCHEMA_ID -WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) -ORDER BY 2, 1`, []any{schema, schema}) + query := xuguListObjectsQuery(schema, constraints) + rows, err := s.queryRows(query.SQL, query.Args) if err != nil { return nil, err } @@ -856,6 +851,159 @@ ORDER BY 2, 1`, []any{schema, schema}) return emptyIfNil(result), rows.Err() } +func metadataListConstraintsFromParams(params map[string]json.RawMessage) metadataListConstraints { + objectTypes := stringSliceParam(params, "object_types") + if len(objectTypes) == 0 { + objectTypes = stringSliceParam(params, "objectTypes") + } + limit := intParam(params, "limit") + offset := intParam(params, "offset") + if limit < 0 { + limit = 0 + } + if offset < 0 { + offset = 0 + } + return metadataListConstraints{ + Filter: stringParam(params, "filter"), + Limit: limit, + Offset: offset, + ObjectTypes: objectTypes, + } +} + +func xuguListTablesQuery(schema string, constraints metadataListConstraints) xuguMetadataListQuery { + return xuguConstrainedMetadataListQuery( + ` +SELECT t.TABLE_NAME, 'TABLE' AS TABLE_TYPE, t.COMMENTS +FROM ALL_TABLES t +JOIN ALL_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) +UNION ALL +SELECT v.VIEW_NAME, 'VIEW' AS TABLE_TYPE, v.COMMENTS +FROM ALL_VIEWS v +JOIN ALL_SCHEMAS s ON s.DB_ID = v.DB_ID AND s.SCHEMA_ID = v.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?)`, + "TABLE_NAME, TABLE_TYPE, COMMENTS", + "TABLE_NAME", + "TABLE_TYPE", + []any{schema, schema}, + constraints, + ) +} + +func xuguListObjectsQuery(schema string, constraints metadataListConstraints) xuguMetadataListQuery { + return xuguConstrainedMetadataListQuery( + ` +SELECT t.TABLE_NAME AS OBJECT_NAME, 'TABLE' AS OBJECT_TYPE, t.COMMENTS +FROM ALL_TABLES t +JOIN ALL_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) +UNION ALL +SELECT v.VIEW_NAME AS OBJECT_NAME, 'VIEW' AS OBJECT_TYPE, v.COMMENTS +FROM ALL_VIEWS v +JOIN ALL_SCHEMAS s ON s.DB_ID = v.DB_ID AND s.SCHEMA_ID = v.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?)`, + "OBJECT_NAME, OBJECT_TYPE, COMMENTS", + "OBJECT_NAME", + "OBJECT_TYPE", + []any{schema, schema}, + constraints, + ) +} + +func xuguConstrainedMetadataListQuery(baseSQL, selectList, nameColumn, typeColumn string, baseArgs []any, constraints metadataListConstraints) xuguMetadataListQuery { + args := append([]any{}, baseArgs...) + where := make([]string, 0, 2) + if filter := strings.TrimSpace(constraints.Filter); filter != "" { + args = append(args, strings.ToUpper(xuguFuzzyLikePattern(filter))) + where = append(where, fmt.Sprintf("UPPER(%s) LIKE ? ESCAPE '\\'", nameColumn)) + } + if len(constraints.ObjectTypes) > 0 { + objectTypes := normalizedXuguObjectTypes(constraints.ObjectTypes) + if len(objectTypes) == 0 { + where = append(where, "1 = 0") + } else { + placeholders := make([]string, 0, len(objectTypes)) + for _, objectType := range objectTypes { + args = append(args, objectType) + placeholders = append(placeholders, "?") + } + where = append(where, fmt.Sprintf("%s IN (%s)", typeColumn, strings.Join(placeholders, ","))) + } + } + + sqlText := fmt.Sprintf("SELECT %s\nFROM (\n%s\n)", selectList, baseSQL) + if len(where) > 0 { + sqlText += "\nWHERE " + strings.Join(where, " AND ") + } + sqlText += fmt.Sprintf("\nORDER BY %s, %s", typeColumn, nameColumn) + + // Xugu documents ROWNUM as the safe pagination path when ORDER BY belongs + // to an inner query; LIMIT is not portable for this UNION metadata query. + if constraints.Limit > 0 { + args = append(args, constraints.Offset+constraints.Limit, constraints.Offset) + sqlText = fmt.Sprintf( + "SELECT %s\nFROM (\n SELECT DBX_Q.*, ROWNUM AS DBX_RN\n FROM (\n%s\n ) DBX_Q\n WHERE ROWNUM <= ?\n)\nWHERE DBX_RN > ?", + selectList, + sqlText, + ) + } else if constraints.Offset > 0 { + args = append(args, constraints.Offset) + sqlText = fmt.Sprintf( + "SELECT %s\nFROM (\n SELECT DBX_Q.*, ROWNUM AS DBX_RN\n FROM (\n%s\n ) DBX_Q\n)\nWHERE DBX_RN > ?", + selectList, + sqlText, + ) + } + + return xuguMetadataListQuery{SQL: sqlText, Args: args} +} + +func normalizedXuguObjectTypes(values []string) []string { + seen := map[string]bool{} + result := make([]string, 0, len(values)) + for _, value := range values { + normalized := strings.ToUpper(strings.TrimSpace(value)) + normalized = strings.ReplaceAll(normalized, "-", "_") + normalized = strings.ReplaceAll(normalized, " ", "_") + switch normalized { + case "TABLE", "BASE_TABLE": + normalized = "TABLE" + case "VIEW": + normalized = "VIEW" + default: + continue + } + if seen[normalized] { + continue + } + seen[normalized] = true + result = append(result, normalized) + } + sort.Strings(result) + return result +} + +func xuguFuzzyLikePattern(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "%%" + } + var builder strings.Builder + builder.Grow(len(value)*2 + 2) + builder.WriteByte('%') + for _, ch := range value { + switch ch { + case '\\', '%', '_': + builder.WriteByte('\\') + } + builder.WriteRune(ch) + builder.WriteByte('%') + } + return builder.String() +} + func (s *server) getColumns(schema, table string) ([]columnInfo, error) { schema, err := s.normalizeSchema(schema) if err != nil { @@ -1825,6 +1973,21 @@ func intParam(params map[string]json.RawMessage, key string) int { return value } +func stringSliceParam(params map[string]json.RawMessage, key string) []string { + if params == nil || len(params[key]) == 0 { + return nil + } + var values []string + if err := json.Unmarshal(params[key], &values); err == nil { + return values + } + var single string + if err := json.Unmarshal(params[key], &single); err == nil && strings.TrimSpace(single) != "" { + return []string{single} + } + return nil +} + func errorResponse(id json.RawMessage, err error) response { return response{JSONRPC: "2.0", ID: id, Error: &rpcError{Code: -1, Message: err.Error()}} } diff --git a/agents/drivers/xugu/main_test.go b/agents/drivers/xugu/main_test.go index 95b40e09a9..89f3f415e2 100644 --- a/agents/drivers/xugu/main_test.go +++ b/agents/drivers/xugu/main_test.go @@ -383,6 +383,73 @@ func TestXuguMetadataAccessErrorDetection(t *testing.T) { } } +func TestXuguListTablesQueryAppliesMetadataConstraints(t *testing.T) { + query := xuguListTablesQuery("APP", metadataListConstraints{ + Filter: "ord_", + ObjectTypes: []string{"view", "table", "VIEW"}, + Limit: 25, + Offset: 50, + }) + + for _, want := range []string{ + "UPPER(TABLE_NAME) LIKE ? ESCAPE '\\'", + "TABLE_TYPE IN (?,?)", + "ORDER BY TABLE_TYPE, TABLE_NAME", + "ROWNUM <= ?", + "DBX_RN > ?", + } { + if !strings.Contains(query.SQL, want) { + t.Fatalf("expected SQL to contain %q:\n%s", want, query.SQL) + } + } + + wantArgs := []any{"APP", "APP", `%O%R%D%\_%`, "TABLE", "VIEW", 75, 50} + assertArgs(t, query.Args, wantArgs) +} + +func TestXuguListObjectsQueryRejectsUnsupportedObjectTypes(t *testing.T) { + query := xuguListObjectsQuery("APP", metadataListConstraints{ + ObjectTypes: []string{"INDEX"}, + Limit: 10, + }) + + if !strings.Contains(query.SQL, "1 = 0") { + t.Fatalf("unsupported object type should produce empty-result predicate:\n%s", query.SQL) + } + + wantArgs := []any{"APP", "APP", 10, 0} + assertArgs(t, query.Args, wantArgs) +} + +func TestMetadataListConstraintsFromParams(t *testing.T) { + params := map[string]json.RawMessage{ + "filter": json.RawMessage(`"tab"`), + "limit": json.RawMessage(`30`), + "offset": json.RawMessage(`5`), + "object_types": json.RawMessage(`["TABLE","VIEW"]`), + } + + constraints := metadataListConstraintsFromParams(params) + if constraints.Filter != "tab" || constraints.Limit != 30 || constraints.Offset != 5 { + t.Fatalf("unexpected constraints: %+v", constraints) + } + if len(constraints.ObjectTypes) != 2 || constraints.ObjectTypes[0] != "TABLE" || constraints.ObjectTypes[1] != "VIEW" { + t.Fatalf("unexpected object types: %+v", constraints.ObjectTypes) + } +} + +func assertArgs(t *testing.T, got []any, want []any) { + t.Helper() + if len(got) != len(want) { + t.Fatalf("args length = %d, want %d: got=%#v want=%#v", len(got), len(want), got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("arg %d = %#v, want %#v; args=%#v", i, got[i], want[i], got) + } + } +} + func TestParseForeignKeyColumns(t *testing.T) { local, ref := parseForeignKeyColumns(`("C1","C2")("ID1","ID2")`) diff --git a/agents/metadata-constraint-coverage.tsv b/agents/metadata-constraint-coverage.tsv new file mode 100644 index 0000000000..7a00fc1079 --- /dev/null +++ b/agents/metadata-constraint-coverage.tsv @@ -0,0 +1,23 @@ +driver strategy scope reason +access intentional-fallback java-jdbc-metadata Uses Access JDBC metadata; no portable stable server-side paging API, common constraints filter locally. +bigquery native-pushdown java-sql Uses dataset INFORMATION_SCHEMA.TABLES with type, normal-name filter, stable order, and literal LIMIT/OFFSET. +cassandra intentional-fallback java-cql Uses Cassandra system_schema; portable LIKE/fuzzy paging is not safe across versions, common constraints filter locally. +dameng native-pushdown java-sql Uses ALL_OBJECTS with type, filter, stable order, and LIMIT/OFFSET with legacy fallback. +databend native-pushdown java-sql Uses system procedures for routine pushdown and shared table metadata constraints for table-like objects. +db2 native-pushdown java-sql Uses SYSCAT metadata with type, filter, stable order, and OFFSET/FETCH. +gbase8a native-pushdown java-sql Uses information_schema with type, filter, stable order, and LIMIT/OFFSET. +gbase8s native-pushdown java-sql Uses systables with type, filter, stable order, and SKIP/FIRST. +goldendb native-pushdown java-sql Uses information_schema with type, filter, stable order, and LIMIT/OFFSET. +h2 native-pushdown java-sql Uses INFORMATION_SCHEMA with type, filter, stable order, and LIMIT/OFFSET. +hive intentional-fallback java-sql Uses Hive metadata paths where stable server-side fuzzy paging is not portable, common constraints filter locally. +informix native-pushdown java-sql Uses systables/sysprocedures with type, filter, stable order, and SKIP/FIRST. +kingbase native-pushdown java-sql Uses sys_catalog/information_schema with type, filter, stable order, and LIMIT/OFFSET for both compatibility modes. +kylin intentional-fallback java-jdbc-metadata Uses JDBC metadata from Kylin driver; no portable server-side paging API, common constraints filter locally. +neo4j intentional-fallback java-graph-metadata Uses JDBC metadata with Cypher fallback; label discovery cannot guarantee relational-style filtered paging, common constraints filter locally. +oceanbase-oracle native-pushdown java-sql Uses Oracle-compatible metadata SQL with type, filter, stable order, and ROWNUM paging. +oracle-go native-pushdown native-go Uses Oracle-compatible metadata SQL with type, filter, stable order, and ROWNUM paging, with legacy profile guard. +snowflake native-pushdown java-sql Uses INFORMATION_SCHEMA with type, filter, stable order, and literal LIMIT/OFFSET. +sundb native-pushdown java-sql Uses information_schema with type, filter, stable order, and LIMIT/OFFSET. +tdengine intentional-fallback java-show-metadata Uses SHOW STABLES/TABLES; stable server-side fuzzy paging is not portable, common constraints filter locally. +trino native-pushdown java-sql Uses information_schema.tables with type, filter, stable order, and literal LIMIT/OFFSET. +xugu native-pushdown native-go Uses ALL_TABLES/ALL_VIEWS metadata with type, filter, stable order, and ROWNUM paging. diff --git a/agents/test-support/src/main/java/com/dbx/agent/test/JdbcMetadataSqlFake.java b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcMetadataSqlFake.java index 872a4c53e7..3482872554 100644 --- a/agents/test-support/src/main/java/com/dbx/agent/test/JdbcMetadataSqlFake.java +++ b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcMetadataSqlFake.java @@ -88,7 +88,7 @@ public final class JdbcMetadataSqlFake { if ("getUpdateCount".equals(name)) { return 0; } - if ("setString".equals(name)) { + if ("setString".equals(name) || "setInt".equals(name) || "setObject".equals(name)) { statements.add("param:" + args[0] + "=" + args[1]); return null; } diff --git a/apps/desktop/src/components/generate/DataGenerateDialog.vue b/apps/desktop/src/components/generate/DataGenerateDialog.vue index b32547b47e..efc6aedafe 100644 --- a/apps/desktop/src/components/generate/DataGenerateDialog.vue +++ b/apps/desktop/src/components/generate/DataGenerateDialog.vue @@ -480,6 +480,9 @@ async function startInsert() { } } perTable.push({ table: r.tableName, total: rowCount, ok, err: rowCount - ok, error: lastError || undefined }); + if (ok > 0) { + store.invalidateMetadataCache(cid, db, props.prefillSchema || undefined, r.tableName); + } } executeResults.value = perTable; currentStep.value = "result"; diff --git a/apps/desktop/src/components/import/TableImportDialog.vue b/apps/desktop/src/components/import/TableImportDialog.vue index 3b57652965..257d12e798 100644 --- a/apps/desktop/src/components/import/TableImportDialog.vue +++ b/apps/desktop/src/components/import/TableImportDialog.vue @@ -197,6 +197,7 @@ async function startImport() { }, ); toast(t("tableImport.success", { count: summary.rowsImported }), 2500); + store.invalidateMetadataCache(props.prefillConnectionId, props.prefillDatabase || "", props.prefillSchema || undefined, props.prefillTable); open.value = false; } catch (e: any) { errorMessage.value = String(e?.message || e); diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue index 72c658b33f..80528ba9cb 100644 --- a/apps/desktop/src/components/sidebar/TreeItem.vue +++ b/apps/desktop/src/components/sidebar/TreeItem.vue @@ -77,8 +77,9 @@ import { buildTableDeleteTemplate, buildTableInsertTemplate, buildTableSelectTem import { connectionFilePath, defaultSqliteBackupFileName, isMemorySqlitePath, sqliteBackupSourcePath } from "@/lib/connection/connectionFile"; import { revealPathInFileManager } from "@/lib/backend/tauri"; import { clearActiveTableReferencePayload, createTableReferencePayload, createTableReferenceDropEvent, setActiveTableReferencePayload, type QueryEditorTableReferencePayload } from "@/lib/editor/queryEditorTableDrop"; -import { editableRowIdentifierColumns, usesSyntheticRowIdKey } from "@/lib/table/tableEditing"; +import { usesSyntheticRowIdKey } from "@/lib/table/tableEditing"; import { tableOpenPageLimit } from "@/lib/table/tableOpenPageLimit"; +import { getCachedTableMetadata, loadTableMetadata, TABLE_METADATA_CACHE_TTL_MS, tableMetadataToDataTabMeta } from "@/lib/metadata/tableMetadataCache"; import { supportsDatabaseCreation, supportsDatabaseSearch, supportsFieldLineage, supportsObjectBrowserTreeNode, supportsSchemaDiagram, supportsSqlFileExecution, supportsTableImport, supportsTableTruncate, supportsTableStructureEditing, usesTreeSchemaMode } from "@/lib/database/databaseCapabilities"; import { copyNameForTreeNode, isDocumentBrowserTreeNode, objectSourceKindForTreeNode, shouldRunTreeNodeRowAction, treeNodeRowAction, treeNodeRowDoubleClickAction } from "@/lib/sidebar/treeNodeClick"; import { isCopySidebarSelectionShortcut, isEditSidebarConnectionShortcut, isPasteSidebarSelectionShortcut } from "@/lib/editor/keyboardShortcuts"; @@ -223,7 +224,7 @@ const useWindowsSidebarCommentFont = isWindows(); type StructureCopyFormat = "tsv" | "markdown"; type DuplicateStructureSource = TreeNode & { connectionId: string; database: string }; -const DATA_TAB_METADATA_TTL_MS = 30_000; +const DATA_TAB_METADATA_TTL_MS = TABLE_METADATA_CACHE_TTL_MS; const { getDatabaseOptions } = useDatabaseOptions(); const showVisibleDatabasesDialog = ref(false); const showVisibleSchemasDialog = ref(false); @@ -1121,6 +1122,9 @@ async function openData() { }); const tableSchema = connectionObjectTreeNodeSchema(config, node.database, node.schema); const tableType = node.type === "view" ? "VIEW" : node.type === "materialized_view" ? "MATERIALIZED_VIEW" : (node.tableType ?? "TABLE"); + const querySchema = config ? connectionObjectTreeQuerySchema(config, node.database, tableSchema) : (tableSchema ?? ""); + const effectiveDbType = effectiveDatabaseTypeForConnection(config); + const metadataDatabaseType = effectiveDbType || config?.db_type || ""; const isSameDataTableTab = (tab: (typeof queryStore.tabs)[number]) => tab.mode === "data" && tab.connectionId === node.connectionId && tab.database === node.database && (tab.schema || "") === (tableSchema || "") && (tab.tableMeta?.tableName || tab.title) === node.label; const existingSameTableTab = queryStore.tabs.find(isSameDataTableTab); const resetReusedDataTabState = (tab: (typeof queryStore.tabs)[number]) => { @@ -1186,7 +1190,21 @@ async function openData() { } const existingTableMeta = tab?.tableMeta; const existingTableMetaAgeMs = tab?.tableMetaUpdatedAt ? Date.now() - tab.tableMetaUpdatedAt : Number.POSITIVE_INFINITY; - const cachedTableMeta = existingTableMeta?.tableName === node.label && existingTableMeta.schema === tableSchema && existingTableMeta.tableType === tableType && existingTableMeta.columns.length > 0 && existingTableMetaAgeMs < DATA_TAB_METADATA_TTL_MS ? existingTableMeta : undefined; + const sharedCachedTableMeta = config + ? getCachedTableMetadata({ + connectionId: node.connectionId, + database: node.database, + schema: querySchema, + tableName: node.label, + tableType, + databaseType: metadataDatabaseType, + driverProfile: config.driver_profile || config.db_type, + }) + : undefined; + const tabCachedTableMeta = existingTableMeta?.tableName === node.label && existingTableMeta.schema === tableSchema && existingTableMeta.tableType === tableType && existingTableMeta.columns.length > 0 && existingTableMetaAgeMs < DATA_TAB_METADATA_TTL_MS ? existingTableMeta : undefined; + const cachedTableMeta = sharedCachedTableMeta ? tableMetadataToDataTabMeta(sharedCachedTableMeta.metadata, tableSchema) : tabCachedTableMeta; + const cachedTableMetaAgeMs = sharedCachedTableMeta?.ageMs ?? existingTableMetaAgeMs; + const cachedTableMetaSource = sharedCachedTableMeta ? "shared" : tabCachedTableMeta ? "tab" : undefined; queryStore.setTableMeta( tabId, cachedTableMeta ?? { @@ -1218,8 +1236,6 @@ async function openData() { logPhase("ensure-connected", { tabId }); if (!config) throw new Error("Connection config not found"); - const querySchema = connectionObjectTreeQuerySchema(config, node.database, tableSchema); - const effectiveDbType = effectiveDatabaseTypeForConnection(config); const limit = tableOpenPageLimit(settingsStore.editorSettings.pageSize); const refreshTableMetaInBackground = async () => { const metadataStartedAt = performance.now(); @@ -1231,30 +1247,34 @@ async function openData() { elapsed: elapsed(), }); try { - const nextColumns = await api.getColumns(node.connectionId, node.database, querySchema, node.label); - const indexes = await api.listIndexes(node.connectionId, node.database, querySchema, node.label).catch(() => []); + const loadedMetadata = await loadTableMetadata({ + connectionId: node.connectionId, + database: node.database, + schema: querySchema, + tableName: node.label, + tableType, + databaseType: metadataDatabaseType, + driverProfile: config.driver_profile || config.db_type, + traceLogger: (event) => console.debug("[DBX][openData:metadata:trace]", { sourceTraceId: traceId, ...event }), + }); if (!isCurrentDataTab()) { console.info("[DBX][openData:metadata:stale]", { traceId, tabId, - columnCount: nextColumns.length, + columnCount: loadedMetadata.metadata.columns.length, elapsed: elapsed(), }); return; } - const nextPrimaryKeys = editableRowIdentifierColumns(effectiveDbType, nextColumns, indexes, tableType); - queryStore.setTableMeta(tabId, { - schema: tableSchema, - tableName: node.label, - tableType, - columns: nextColumns, - primaryKeys: nextPrimaryKeys, - }); + const nextTableMeta = tableMetadataToDataTabMeta(loadedMetadata.metadata, tableSchema); + queryStore.setTableMeta(tabId, nextTableMeta); console.info("[DBX][openData:metadata:done]", { traceId, tabId, - columnCount: nextColumns.length, - primaryKeyCount: nextPrimaryKeys.length, + columnCount: nextTableMeta.columns.length, + primaryKeyCount: nextTableMeta.primaryKeys.length, + cacheStatus: loadedMetadata.cacheStatus, + ageMs: Math.round(loadedMetadata.ageMs), elapsed: elapsed(), metadataMs: Math.round(performance.now() - metadataStartedAt), }); @@ -1269,7 +1289,8 @@ async function openData() { tabId, columnCount: cachedTableMeta.columns.length, primaryKeyCount: cachedTableMeta.primaryKeys.length, - ageMs: Math.round(existingTableMetaAgeMs), + source: cachedTableMetaSource, + ageMs: Math.round(cachedTableMetaAgeMs), elapsed: elapsed(), }); } else { diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index c3bf8c24c1..1d391c40bb 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -518,13 +518,16 @@ export async function getTableComment(_connectionId: string, _database: string, throw new Error("Table comment lookup is not available in the web backend"); } -export async function listObjects(connectionId: string, database: string, schema: string, objectTypes?: SidebarObjectKind[]): Promise { +export async function listObjects(connectionId: string, database: string, schema: string, objectTypes?: SidebarObjectKind[], filter?: string, limit?: number, offset?: number): Promise { return get( `/api/schema/objects?${qs({ connection_id: connectionId, database, schema, object_types: objectTypes?.join(","), + filter, + limit, + offset, })}`, ); } diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 7cecf5286b..dd54d1d8cd 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -610,8 +610,8 @@ export async function getTableComment(connectionId: string, database: string, sc return invoke("get_table_comment", { connectionId, database, schema, table }); } -export async function listObjects(connectionId: string, database: string, schema: string, objectTypes?: SidebarObjectKind[]): Promise { - return invoke("list_objects", { connectionId, database, schema, objectTypes }); +export async function listObjects(connectionId: string, database: string, schema: string, objectTypes?: SidebarObjectKind[], filter?: string, limit?: number, offset?: number): Promise { + return invoke("list_objects", { connectionId, database, schema, objectTypes, filter, limit, offset }); } export async function listObjectStatistics(connectionId: string, database: string, schema: string): Promise { diff --git a/apps/desktop/src/lib/metadata/metadataLoadCoordinator.ts b/apps/desktop/src/lib/metadata/metadataLoadCoordinator.ts new file mode 100644 index 0000000000..892ad40477 --- /dev/null +++ b/apps/desktop/src/lib/metadata/metadataLoadCoordinator.ts @@ -0,0 +1,145 @@ +import { metadataScopeKey, type MetadataScopeInput } from "./metadataLoadScope"; + +export type MetadataLoadCacheStatus = "none" | "hit" | "miss" | "stale" | "refresh"; + +export interface MetadataLoadTraceDetails { + cacheStatus?: MetadataLoadCacheStatus; + resultCount?: number; + deduped?: boolean; + force?: boolean; + stale?: boolean; + superseded?: boolean; + error?: unknown; +} + +export interface MetadataLoadTraceEvent extends MetadataLoadTraceDetails { + event: string; + traceId: string; + scopeKey: string; + elapsedMs: number; +} + +export type MetadataLoadTraceLogger = (event: MetadataLoadTraceEvent) => void; + +export interface MetadataLoadTrace { + traceId: string; + scopeKey: string; + startedAt: number; + elapsedMs: () => number; + event: (event: string, details?: MetadataLoadTraceDetails) => MetadataLoadTraceEvent; +} + +export interface MetadataLoadCoordinatorEvent { + event: "start" | "dedupe" | "done" | "error"; + key: string; + kind: string; + active: number; + force: boolean; +} + +export type MetadataLoadCoordinatorLogger = (event: MetadataLoadCoordinatorEvent) => void; + +export interface MetadataLoadCoordinatorOptions { + force?: boolean; + kind?: string; +} + +interface InFlightLoad { + kind: string; + promise: Promise; +} + +function defaultTraceId() { + return Math.random().toString(16).slice(2, 10); +} + +function defaultNow() { + return typeof performance !== "undefined" && typeof performance.now === "function" ? performance.now() : Date.now(); +} + +export function createMetadataLoadTrace( + scope: MetadataScopeInput | string, + options?: { + traceId?: string; + now?: () => number; + }, +): MetadataLoadTrace { + const now = options?.now ?? defaultNow; + const startedAt = now(); + const scopeKey = typeof scope === "string" ? scope : metadataScopeKey(scope); + const traceId = options?.traceId ?? defaultTraceId(); + const elapsedMs = () => Math.max(0, Math.round(now() - startedAt)); + return { + traceId, + scopeKey, + startedAt, + elapsedMs, + event: (event, details = {}) => ({ + event, + traceId, + scopeKey, + elapsedMs: elapsedMs(), + ...details, + }), + }; +} + +export function logMetadataLoadTrace(logger: MetadataLoadTraceLogger | undefined, trace: MetadataLoadTrace, event: string, details?: MetadataLoadTraceDetails) { + logger?.(trace.event(event, details)); +} + +export class MetadataLoadCoordinator { + private readonly inFlight = new Map>(); + + constructor(private readonly logger?: MetadataLoadCoordinatorLogger) {} + + run(scope: MetadataScopeInput | string, load: () => Promise, options?: MetadataLoadCoordinatorOptions): Promise { + const key = typeof scope === "string" ? scope : metadataScopeKey(scope); + const kind = options?.kind ?? (typeof scope === "string" ? "metadata" : scope.kind); + const force = options?.force === true; + const existing = force ? undefined : this.inFlight.get(key); + if (existing) { + this.logger?.({ event: "dedupe", key, kind: existing.kind, active: this.inFlight.size, force: false }); + return existing.promise as Promise; + } + + this.logger?.({ event: "start", key, kind, active: this.inFlight.size + 1, force }); + const promise = Promise.resolve() + .then(load) + .then( + (value) => { + this.logger?.({ event: "done", key, kind, active: this.inFlight.size, force }); + return value; + }, + (error) => { + this.logger?.({ event: "error", key, kind, active: this.inFlight.size, force }); + throw error; + }, + ) + .finally(() => { + const current = this.inFlight.get(key); + if (current?.promise === promise) { + this.inFlight.delete(key); + } + }); + + if (!force) { + this.inFlight.set(key, { kind, promise }); + } + return promise; + } + + has(scope: MetadataScopeInput | string): boolean { + const key = typeof scope === "string" ? scope : metadataScopeKey(scope); + return this.inFlight.has(key); + } + + clear(scope?: MetadataScopeInput | string) { + if (scope === undefined) { + this.inFlight.clear(); + return; + } + const key = typeof scope === "string" ? scope : metadataScopeKey(scope); + this.inFlight.delete(key); + } +} diff --git a/apps/desktop/src/lib/metadata/metadataLoadScope.ts b/apps/desktop/src/lib/metadata/metadataLoadScope.ts new file mode 100644 index 0000000000..c807fc0404 --- /dev/null +++ b/apps/desktop/src/lib/metadata/metadataLoadScope.ts @@ -0,0 +1,86 @@ +export type MetadataScopeValue = string | number | boolean | null | undefined; + +export interface MetadataScopeInput { + kind: string; + connectionId?: string | null; + database?: string | null; + schema?: string | null; + nodeKind?: string | null; + tableName?: string | null; + tableType?: string | null; + objectTypes?: readonly (string | null | undefined)[] | null; + searchFilter?: string | null; + limit?: number | null; + offset?: number | null; + sidebarDisplayMode?: string | null; + driverProfile?: string | null; + extra?: Record | null; +} + +function normalizeString(value: string | null | undefined): string | undefined { + if (value == null) return undefined; + return value; +} + +function normalizeSearchFilter(value: string | null | undefined): string | undefined { + const normalized = value?.trim(); + return normalized ? normalized : undefined; +} + +function normalizeNumber(value: number | null | undefined): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined; +} + +function normalizeObjectTypes(values: readonly (string | null | undefined)[] | null | undefined): string[] | undefined { + if (!values?.length) return undefined; + const normalized = values + .filter((value): value is string => typeof value === "string" && value.trim().length > 0) + .map((value) => value.trim().toUpperCase().replace(/\s+/g, "_")) + .sort(); + return normalized.length > 0 ? [...new Set(normalized)] : undefined; +} + +function stableEntries(record: Record): [string, unknown][] { + return Object.entries(record) + .filter(([, value]) => value !== undefined) + .sort(([left], [right]) => left.localeCompare(right)); +} + +function normalizeExtra(extra: MetadataScopeInput["extra"]): Record | undefined { + if (!extra) return undefined; + const entries = stableEntries(extra) + .map(([key, value]) => { + if (Array.isArray(value)) { + return [key, value.filter((item) => item !== undefined)] as const; + } + return [key, value] as const; + }) + .filter(([, value]) => { + if (Array.isArray(value)) return value.length > 0; + return value !== undefined; + }); + return entries.length > 0 ? Object.fromEntries(entries) : undefined; +} + +export function metadataScopeParts(input: MetadataScopeInput): Record { + return { + kind: input.kind, + connectionId: normalizeString(input.connectionId), + database: normalizeString(input.database), + schema: normalizeString(input.schema), + nodeKind: normalizeString(input.nodeKind), + tableName: normalizeString(input.tableName), + tableType: normalizeString(input.tableType), + objectTypes: normalizeObjectTypes(input.objectTypes), + searchFilter: normalizeSearchFilter(input.searchFilter), + limit: normalizeNumber(input.limit), + offset: normalizeNumber(input.offset), + sidebarDisplayMode: normalizeString(input.sidebarDisplayMode), + driverProfile: normalizeString(input.driverProfile), + extra: normalizeExtra(input.extra), + }; +} + +export function metadataScopeKey(input: MetadataScopeInput): string { + return JSON.stringify(stableEntries(metadataScopeParts(input))); +} diff --git a/apps/desktop/src/lib/metadata/metadataResultCache.ts b/apps/desktop/src/lib/metadata/metadataResultCache.ts new file mode 100644 index 0000000000..25fb3f36bd --- /dev/null +++ b/apps/desktop/src/lib/metadata/metadataResultCache.ts @@ -0,0 +1,99 @@ +import { metadataScopeKey, metadataScopeParts, type MetadataScopeInput } from "./metadataLoadScope"; + +export interface MetadataCacheHit { + value: T; + cachedAt: number; + ageMs: number; + stale: boolean; +} + +export type MetadataCacheInvalidation = Partial>; + +interface MetadataCacheEntry { + scope: ReturnType; + value: T; + cachedAt: number; +} + +export interface MetadataResultCacheOptions { + ttlMs: number; + maxEntries: number; + now?: () => number; +} + +function normalizeInvalidation(match: MetadataCacheInvalidation): Partial> { + const normalized = metadataScopeParts({ ...match, kind: match.kind ?? "__metadata_cache_match__" }); + const result: Partial> = {}; + for (const key of Object.keys(match) as (keyof MetadataCacheInvalidation)[]) { + const value = normalized[key as keyof typeof normalized]; + if (value !== undefined) { + (result as Record)[key] = value; + } + } + return result; +} + +function valuesEqual(left: unknown, right: unknown): boolean { + if (Array.isArray(left) || Array.isArray(right) || (left && typeof left === "object") || (right && typeof right === "object")) { + return JSON.stringify(left) === JSON.stringify(right); + } + return left === right; +} + +export class MetadataResultCache { + private readonly entries = new Map>(); + private readonly now: () => number; + + constructor(private readonly options: MetadataResultCacheOptions) { + this.now = options.now ?? Date.now; + } + + get(scope: MetadataScopeInput, options?: { allowStale?: boolean }): MetadataCacheHit | undefined { + const key = metadataScopeKey(scope); + const entry = this.entries.get(key); + if (!entry) return undefined; + const ageMs = Math.max(0, this.now() - entry.cachedAt); + const stale = ageMs > this.options.ttlMs; + if (stale && !options?.allowStale) return undefined; + + this.entries.delete(key); + this.entries.set(key, entry); + return { value: entry.value, cachedAt: entry.cachedAt, ageMs, stale }; + } + + set(scope: MetadataScopeInput, value: T): void { + const key = metadataScopeKey(scope); + this.entries.delete(key); + this.entries.set(key, { scope: metadataScopeParts(scope), value, cachedAt: this.now() }); + this.evictOldest(); + } + + invalidate(match: MetadataCacheInvalidation): number { + const normalized = normalizeInvalidation(match); + let removed = 0; + for (const [key, entry] of this.entries) { + const matches = Object.entries(normalized).every(([field, value]) => valuesEqual(entry.scope[field as keyof typeof entry.scope], value)); + if (matches) { + this.entries.delete(key); + removed++; + } + } + return removed; + } + + clear(): void { + this.entries.clear(); + } + + get size(): number { + return this.entries.size; + } + + private evictOldest(): void { + while (this.entries.size > this.options.maxEntries) { + const oldest = this.entries.keys().next().value; + if (oldest === undefined) return; + this.entries.delete(oldest); + } + } +} diff --git a/apps/desktop/src/lib/metadata/tableMetadataCache.ts b/apps/desktop/src/lib/metadata/tableMetadataCache.ts new file mode 100644 index 0000000000..121a60abfb --- /dev/null +++ b/apps/desktop/src/lib/metadata/tableMetadataCache.ts @@ -0,0 +1,126 @@ +import type { ColumnInfo, DatabaseType, IndexInfo, QueryTab } from "@/types/database"; +import * as api from "@/lib/backend/api"; +import { editableRowIdentifierColumns } from "@/lib/table/tableEditing"; +import { createMetadataLoadTrace, logMetadataLoadTrace, MetadataLoadCoordinator, type MetadataLoadCacheStatus, type MetadataLoadTraceLogger } from "./metadataLoadCoordinator"; +import type { MetadataScopeInput } from "./metadataLoadScope"; +import { MetadataResultCache, type MetadataCacheInvalidation } from "./metadataResultCache"; + +export const TABLE_METADATA_CACHE_TTL_MS = 30_000; +const TABLE_METADATA_CACHE_MAX_ENTRIES = 120; + +export interface TableMetadata { + schema?: string; + tableName: string; + tableType?: string; + columns: ColumnInfo[]; + indexes: IndexInfo[]; + primaryKeys: string[]; + cachedAt: number; +} + +export interface TableMetadataRequest { + connectionId: string; + database: string; + schema?: string; + tableName: string; + tableType?: string; + databaseType: DatabaseType | string; + driverProfile?: string; + force?: boolean; + traceLogger?: MetadataLoadTraceLogger; +} + +export interface TableMetadataLoadResult { + metadata: TableMetadata; + cacheStatus: MetadataLoadCacheStatus; + ageMs: number; +} + +const tableMetadataCache = new MetadataResultCache({ + ttlMs: TABLE_METADATA_CACHE_TTL_MS, + maxEntries: TABLE_METADATA_CACHE_MAX_ENTRIES, +}); + +const tableMetadataCoordinator = new MetadataLoadCoordinator((event) => { + console.debug("[DBX][metadata-load:table-coordinator]", event); +}); + +export function tableMetadataScope(request: Pick): MetadataScopeInput { + return { + kind: "table-metadata", + connectionId: request.connectionId, + database: request.database, + schema: request.schema ?? "", + tableName: request.tableName, + tableType: request.tableType, + driverProfile: request.driverProfile || request.databaseType, + }; +} + +export function getCachedTableMetadata(request: Pick): TableMetadataLoadResult | undefined { + const hit = tableMetadataCache.get(tableMetadataScope(request)); + if (!hit) return undefined; + return { metadata: hit.value, cacheStatus: hit.stale ? "stale" : "hit", ageMs: hit.ageMs }; +} + +export function tableMetadataToDataTabMeta(metadata: TableMetadata, schema = metadata.schema): NonNullable { + return { + schema, + tableName: metadata.tableName, + tableType: metadata.tableType, + columns: metadata.columns, + primaryKeys: metadata.primaryKeys, + }; +} + +export async function loadTableMetadata(request: TableMetadataRequest): Promise { + const scope = tableMetadataScope(request); + const trace = createMetadataLoadTrace(scope); + if (!request.force) { + const cached = tableMetadataCache.get(scope); + if (cached) { + logMetadataLoadTrace(request.traceLogger, trace, "cache-hit", { + cacheStatus: cached.stale ? "stale" : "hit", + resultCount: cached.value.columns.length, + stale: cached.stale, + }); + return { metadata: cached.value, cacheStatus: cached.stale ? "stale" : "hit", ageMs: cached.ageMs }; + } + } + + logMetadataLoadTrace(request.traceLogger, trace, "cache-miss", { cacheStatus: request.force ? "refresh" : "miss", force: request.force === true }); + const metadata = await tableMetadataCoordinator.run( + scope, + async () => { + const columns = await api.getColumns(request.connectionId, request.database, request.schema ?? "", request.tableName); + const indexes = await api.listIndexes(request.connectionId, request.database, request.schema ?? "", request.tableName).catch((): IndexInfo[] => []); + const primaryKeys = editableRowIdentifierColumns(request.databaseType as DatabaseType, columns, indexes, request.tableType); + return { + schema: request.schema || undefined, + tableName: request.tableName, + tableType: request.tableType, + columns, + indexes, + primaryKeys, + cachedAt: Date.now(), + }; + }, + { force: request.force, kind: scope.kind }, + ); + + tableMetadataCache.set(scope, metadata); + logMetadataLoadTrace(request.traceLogger, trace, "done", { + cacheStatus: request.force ? "refresh" : "miss", + resultCount: metadata.columns.length, + force: request.force === true, + }); + return { metadata, cacheStatus: request.force ? "refresh" : "miss", ageMs: 0 }; +} + +export function invalidateTableMetadataCache(match: MetadataCacheInvalidation): number { + return tableMetadataCache.invalidate(match); +} + +export function clearTableMetadataCache(): void { + tableMetadataCache.clear(); +} diff --git a/apps/desktop/src/lib/sidebar/sidebarTableSearchControl.ts b/apps/desktop/src/lib/sidebar/sidebarTableSearchControl.ts index c7e6cccffe..445eb44a0b 100644 --- a/apps/desktop/src/lib/sidebar/sidebarTableSearchControl.ts +++ b/apps/desktop/src/lib/sidebar/sidebarTableSearchControl.ts @@ -2,7 +2,7 @@ import type { FlatTreeNode } from "@/composables/useFlatTree"; import type { TreeNode, TreeNodeType } from "@/types/database"; const simpleObjectParentTypes = new Set(["database", "schema", "linked-server-schema"]); -const simpleObjectChildTypes = new Set(["table", "view", "materialized_view", "procedure", "function", "sequence", "package", "package-body", "load-more"]); +const tableSearchableChildTypes = new Set(["table", "view", "materialized_view", "load-more"]); export function isSidebarTableSearchControlNode(node: TreeNode): boolean { return node.type === "table-search-control"; @@ -12,18 +12,19 @@ function tableSearchControlId(parentId: string): string { return `${parentId}:__table_search`; } -function parentHasDirectTableList(node: TreeNode): boolean { - return !!node.children?.some((child) => simpleObjectChildTypes.has(child.type)); +function parentHasSearchableTableList(node: TreeNode): boolean { + // Routine-only schemas should not show a table-specific search box. + return !!node.children?.some((child) => tableSearchableChildTypes.has(child.type)); } function shouldInsertTableSearchControl(item: FlatTreeNode, sidebarObjectDisplay: "simple" | "grouped", activeQueries: Readonly>): boolean { const node = item.node; if (!node.isExpanded) return false; if (sidebarObjectDisplay === "grouped") { - return node.type === "group-tables"; + return node.type === "group-tables" && (parentHasSearchableTableList(node) || !!activeQueries[node.id]?.trim()); } if (!simpleObjectParentTypes.has(node.type)) return false; - return parentHasDirectTableList(node) || !!activeQueries[node.id]?.trim(); + return parentHasSearchableTableList(node) || !!activeQueries[node.id]?.trim(); } function buildTableSearchControlNode(parent: TreeNode): TreeNode { diff --git a/apps/desktop/src/stores/__tests__/connectionStore.metadataLoading.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.metadataLoading.spec.ts new file mode 100644 index 0000000000..d561836305 --- /dev/null +++ b/apps/desktop/src/stores/__tests__/connectionStore.metadataLoading.spec.ts @@ -0,0 +1,99 @@ +import { createPinia, setActivePinia } from "pinia"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { ConnectionConfig, TableInfo, TreeNode } from "@/types/database"; + +function installLocalStorage() { + const data = new Map(); + vi.stubGlobal("localStorage", { + getItem: vi.fn((key: string) => data.get(key) ?? null), + setItem: vi.fn((key: string, value: string) => data.set(key, value)), + removeItem: vi.fn((key: string) => data.delete(key)), + }); +} + +function postgresConnection(): ConnectionConfig { + return { + id: "pg-1", + name: "Postgres", + db_type: "postgres", + host: "127.0.0.1", + port: 5432, + username: "postgres", + password: "", + database: "app", + } as ConnectionConfig; +} + +describe("connectionStore metadata loading", () => { + beforeEach(() => { + vi.resetModules(); + vi.unstubAllGlobals(); + installLocalStorage(); + setActivePinia(createPinia()); + }); + + it("renders simple-mode table children without waiting for supplemental objects", async () => { + const tables: TableInfo[] = [{ name: "users", table_type: "TABLE", comment: null }]; + const listTables = vi.fn().mockResolvedValue(tables); + const listObjects = vi.fn(() => new Promise(() => undefined)); + + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + checkConnectionHealth: vi.fn().mockResolvedValue(undefined), + deleteSchemaCachePrefix: vi.fn().mockResolvedValue(undefined), + listObjects, + listTables, + loadSchemaCache: vi.fn().mockResolvedValue(null), + saveSchemaCache: vi.fn().mockResolvedValue(undefined), + saveConnections: vi.fn().mockResolvedValue(undefined), + saveSidebarLayout: vi.fn().mockResolvedValue(undefined), + })); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const { useSettingsStore } = await import("@/stores/settingsStore"); + const store = useConnectionStore(); + const settingsStore = useSettingsStore(); + settingsStore.editorSettings.sidebarObjectDisplay = "simple"; + + const connection = postgresConnection(); + const schemaNode: TreeNode = { + id: "pg-1:app:public", + label: "public", + type: "schema", + connectionId: connection.id, + database: "app", + schema: "public", + isExpanded: false, + children: [], + }; + store.connections = [connection]; + store.connectedIds.add(connection.id); + store.treeNodes = [ + { + id: connection.id, + label: connection.name, + type: "connection", + connectionId: connection.id, + isExpanded: true, + children: [ + { + id: "pg-1:app", + label: "app", + type: "database", + connectionId: connection.id, + database: "app", + isExpanded: true, + children: [schemaNode], + }, + ], + }, + ]; + + const result = await Promise.race([store.loadTables(connection.id, "app", "public").then(() => "done"), new Promise((resolve) => setTimeout(() => resolve("timeout"), 50))]); + + expect(result).toBe("done"); + expect(listTables).toHaveBeenCalledWith(connection.id, "app", "public", undefined, 1001, 0); + expect(listObjects).toHaveBeenCalled(); + expect(schemaNode.children?.map((node) => node.label)).toEqual(["users"]); + }); +}); diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index f9644a8407..26208e2e79 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -60,6 +60,10 @@ import { completionSchemasFromTree, completionTablesFromTree } from "@/lib/metad import { kvRootNodeLabel } from "@/lib/kv/kvRootPresentation"; import { REDIS_SCAN_PAGE_SIZE_DEFAULT } from "@/lib/redis/redisKeyPattern"; import { appendAgentDriverUpdateHint, hasAgentDriverUpdate, type AgentDriverInstallState } from "@/lib/connection/agentDriverInstallHint"; +import { createMetadataLoadTrace, logMetadataLoadTrace, MetadataLoadCoordinator, type MetadataLoadTraceLogger } from "@/lib/metadata/metadataLoadCoordinator"; +import type { MetadataScopeInput } from "@/lib/metadata/metadataLoadScope"; +import { MetadataResultCache, type MetadataCacheInvalidation } from "@/lib/metadata/metadataResultCache"; +import { invalidateTableMetadataCache } from "@/lib/metadata/tableMetadataCache"; import { MetadataTaskLimiter } from "@/lib/metadata/metadataTaskLimiter"; import i18n from "@/i18n"; import type { MqAdminConfig } from "@/types/mq"; @@ -72,6 +76,8 @@ const METADATA_LOAD_MIN_TIMEOUT_MS = 15_000; const METADATA_LOAD_DISABLED_QUERY_TIMEOUT_MS = 60_000; const DISCONNECT_REQUEST_TIMEOUT_MS = 5_000; const DEFAULT_KEEPALIVE_INTERVAL_SECS = 30; +const METADATA_LIST_PAGE_CACHE_TTL_MS = 30_000; +const METADATA_LIST_PAGE_CACHE_MAX_ENTRIES = 160; export const COMPLETION_METADATA_CONCURRENCY = 2; const MONGO_LEGACY_DRIVER_PROFILE = "mongodb-legacy"; const MONGO_LEGACY_DRIVER_LABEL = "MongoDB (Legacy)"; @@ -175,6 +181,8 @@ interface PersistedTreeChildrenLoadResult { isStale: boolean; } +type MetadataListPageResult = TableInfo[] | ObjectInfo[]; + type BeforeConnectHandler = (config: ConnectionConfig) => Promise; export const CONNECTION_ATTEMPT_CANCELLED_MESSAGE = "Connection attempt was cancelled"; @@ -184,6 +192,10 @@ function redisDbLabel(db: number, _loadedKeyCount?: number, totalKeyCount?: numb return `db${db} (${totalKeyCount})`; } +function metadataDriverProfile(config?: ConnectionConfig): string | undefined { + return config?.driver_profile || config?.db_type; +} + export const useConnectionStore = defineStore("connection", () => { const settingsStore = useSettingsStore(); const connections = ref([]); @@ -280,6 +292,16 @@ export const useConnectionStore = defineStore("connection", () => { const sidebarLayout = ref(emptyLayout()); let layoutPersistTimer: ReturnType | null = null; const staleTreeRefreshIds = new Set(); + const metadataLoadCoordinator = new MetadataLoadCoordinator((event) => { + console.debug("[DBX][metadata-load:coordinator]", event); + }); + const metadataListPageCache = new MetadataResultCache({ + ttlMs: METADATA_LIST_PAGE_CACHE_TTL_MS, + maxEntries: METADATA_LIST_PAGE_CACHE_MAX_ENTRIES, + }); + const metadataTraceLogger: MetadataLoadTraceLogger = (event) => { + console.debug("[DBX][metadata-load:trace]", event); + }; const connectInFlight = new Map>(); const disconnectInFlight = new Map>(); const cancelDisconnectInFlight = new Map>(); @@ -289,6 +311,35 @@ export const useConnectionStore = defineStore("connection", () => { let beforeConnectHandler: BeforeConnectHandler | null = null; let initFromDiskPromise: Promise | null = null; + function runTreeMetadataLoad(scope: MetadataScopeInput, task: () => Promise, options?: LoadTreeOptions): Promise { + return metadataLoadCoordinator.run(scope, task, { force: options?.force, kind: scope.kind }); + } + + async function loadCachedMetadataListPage(scope: MetadataScopeInput, load: () => Promise, options?: { force?: boolean }): Promise { + const trace = createMetadataLoadTrace(scope); + if (!options?.force) { + const cached = metadataListPageCache.get(scope); + if (cached) { + logMetadataLoadTrace(metadataTraceLogger, trace, "cache-hit", { + cacheStatus: cached.stale ? "stale" : "hit", + resultCount: cached.value.length, + stale: cached.stale, + }); + return cached.value as T; + } + } + + logMetadataLoadTrace(metadataTraceLogger, trace, "cache-miss", { cacheStatus: options?.force ? "refresh" : "miss", force: options?.force === true }); + const result = await load(); + metadataListPageCache.set(scope, result); + logMetadataLoadTrace(metadataTraceLogger, trace, "done", { + cacheStatus: options?.force ? "refresh" : "miss", + resultCount: result.length, + force: options?.force === true, + }); + return result; + } + function startEditing(id: string) { editingConnectionId.value = id; } @@ -881,6 +932,75 @@ export const useConnectionStore = defineStore("connection", () => { return schemaCacheKey(node.connectionId || "", node.database || "", node.schema || "", node.type, "objects-v3"); } + function metadataListDriverProfile(connectionId?: string): string | undefined { + return connectionId ? metadataDriverProfile(getConfig(connectionId)) : undefined; + } + + function metadataListCacheScope(options: { + kind: string; + connectionId?: string | null; + database?: string | null; + schema?: string | null; + nodeKind?: string | null; + objectTypes?: readonly string[] | null; + searchFilter?: string | null; + limit?: number | null; + offset?: number | null; + sidebarDisplayMode?: string | null; + extra?: MetadataScopeInput["extra"]; + }): MetadataScopeInput { + return { + kind: options.kind, + connectionId: options.connectionId, + database: options.database, + schema: options.schema, + nodeKind: options.nodeKind, + objectTypes: options.objectTypes, + searchFilter: options.searchFilter, + limit: options.limit, + offset: options.offset, + sidebarDisplayMode: options.sidebarDisplayMode, + driverProfile: metadataListDriverProfile(options.connectionId || undefined), + extra: options.extra, + }; + } + + function invalidateMetadataCaches(match: MetadataCacheInvalidation): number { + return metadataListPageCache.invalidate(match) + invalidateTableMetadataCache(match); + } + + function invalidateMetadataCachesByTreePrefix(prefix: string) { + const [connectionId, database, schema, tableName] = prefix.split(":").map((part) => { + try { + return decodeURIComponent(part); + } catch { + return part; + } + }); + if (!connectionId) return; + invalidateMetadataCaches({ + connectionId, + database: database || undefined, + schema: schema || undefined, + tableName: tableName && !tableName.startsWith("__") ? tableName : undefined, + }); + } + + function invalidateMetadataCachesForNode(node: TreeNode) { + if (!node.connectionId) return; + const tableName = node.tableName || (node.type === "table" || node.type === "view" || node.type === "materialized_view" || node.type === "mongo-collection" ? node.label : undefined); + invalidateMetadataCaches({ + connectionId: node.connectionId, + database: node.database || undefined, + schema: node.schema || undefined, + tableName, + }); + } + + function invalidateMetadataCache(connectionId: string, database?: string, schema?: string, tableName?: string) { + invalidateMetadataCaches({ connectionId, database, schema, tableName }); + } + function buildLoadMoreNode(parent: TreeNode, offset: number, pageSize: number): TreeNode { return { id: `${parent.id}:__load_more:${offset}`, @@ -974,13 +1094,30 @@ export const useConnectionStore = defineStore("connection", () => { offset: number; pageSize: number; searchFilter?: string; + force?: boolean; }): Promise<{ children: TreeNode[]; objectCount: number; hasMore: boolean; nextOffset: number }> { if (!options.node.connectionId || !options.node.database) { return { children: [], objectCount: 0, hasMore: false, nextOffset: options.offset }; } const searchFilter = (options.searchFilter ?? sidebarSearchQuery.value) || undefined; const fetchLimit = searchFilter ? options.pageSize : options.pageSize + 1; - const tables = await api.listTables(options.node.connectionId, options.node.database, options.querySchema, searchFilter, fetchLimit, searchFilter ? undefined : options.offset, options.objectTypes); + const fetchOffset = searchFilter ? undefined : options.offset; + const tables = await loadCachedMetadataListPage( + metadataListCacheScope({ + kind: "table-list-page", + connectionId: options.node.connectionId, + database: options.node.database, + schema: options.querySchema, + nodeKind: options.node.type, + objectTypes: options.objectTypes, + searchFilter, + limit: fetchLimit, + offset: fetchOffset, + sidebarDisplayMode: "grouped", + }), + () => api.listTables(options.node.connectionId!, options.node.database!, options.querySchema, searchFilter, fetchLimit, fetchOffset, options.objectTypes), + { force: options.force }, + ); const hasMore = searchFilter ? false : tables.length > options.pageSize; const pageTables = hasMore ? tables.slice(0, options.pageSize) : tables; indexCompletionTables(options.node.connectionId, options.node.database, options.effectiveSchema, tableInfosToCompletionTables(pageTables, options.effectiveSchema)); @@ -1010,35 +1147,30 @@ export const useConnectionStore = defineStore("connection", () => { offset: number; pageSize: number; searchFilter?: string; + force?: boolean; }): Promise<{ children: TreeNode[]; objectCount: number; hasMore: boolean; nextOffset: number }> { const searchFilter = (options.searchFilter ?? sidebarSearchQuery.value) || undefined; const fetchLimit = searchFilter ? options.pageSize : options.pageSize + 1; - const tables = await api.listTables(options.connectionId, options.database, options.querySchema, searchFilter, fetchLimit, searchFilter ? undefined : options.offset); + const fetchOffset = searchFilter ? undefined : options.offset; + const tables = await loadCachedMetadataListPage( + metadataListCacheScope({ + kind: "table-list-page", + connectionId: options.connectionId, + database: options.database, + schema: options.querySchema, + nodeKind: "simple-tables", + searchFilter, + limit: fetchLimit, + offset: fetchOffset, + sidebarDisplayMode: "simple", + }), + () => api.listTables(options.connectionId, options.database, options.querySchema, searchFilter, fetchLimit, fetchOffset), + { force: options.force }, + ); const hasMore = searchFilter ? false : tables.length > options.pageSize; const pageTables = hasMore ? tables.slice(0, options.pageSize) : tables; indexCompletionTables(options.connectionId, options.database, options.effectiveSchema, tableInfosToCompletionTables(pageTables, options.effectiveSchema)); - if (!searchFilter) { - try { - const objects = options.nonTableObjectTypes.length > 0 ? await api.listObjects(options.connectionId, options.database, options.querySchema, options.nonTableObjectTypes) : []; - const children = buildSimpleObjectTreeNodes({ - nodeId: options.nodeId, - connectionId: options.connectionId, - database: options.database, - schema: options.effectiveSchema, - objects: mergeTableInfosIntoObjects(filterSimpleSidebarSupplementalObjects(objects), pageTables, options.effectiveSchema), - }); - return { - children, - objectCount: children.length, - hasMore, - nextOffset: options.offset + pageTables.length, - }; - } catch { - // Some drivers only expose table metadata; keep the paged table tree usable. - } - } - const children = buildTableTreeNodes({ nodeId: options.nodeId, connectionId: options.connectionId, @@ -1054,6 +1186,55 @@ export const useConnectionStore = defineStore("connection", () => { }; } + async function loadSimpleSupplementalObjectChildren(options: { node: TreeNode; nodeId: string; connectionId: string; database: string; querySchema: string; effectiveSchema?: string; objectTypes: DatabaseObjectTreeKind[]; cacheKey: string; loadOptions?: LoadTreeOptions }) { + if (options.objectTypes.length === 0) return; + const searchFilter = activeTreeLoadSearchFilter(options.loadOptions); + if (searchFilter) return; + + try { + const objects = await loadCachedMetadataListPage( + metadataListCacheScope({ + kind: "object-list-page", + connectionId: options.connectionId, + database: options.database, + schema: options.querySchema, + nodeKind: "simple-supplemental", + objectTypes: options.objectTypes, + sidebarDisplayMode: "simple", + }), + () => api.listObjects(options.connectionId, options.database, options.querySchema, options.objectTypes), + { force: options.loadOptions?.force }, + ); + const supplementalObjects = filterSimpleSidebarSupplementalObjects(objects); + if (supplementalObjects.length === 0) return; + const supplementalChildren = buildSimpleObjectTreeNodes({ + nodeId: options.nodeId, + connectionId: options.connectionId, + database: options.database, + schema: options.effectiveSchema, + objects: supplementalObjects, + }); + if (supplementalChildren.length === 0) return; + if (isTreeLoadSearchChanged(searchFilter, options.loadOptions)) return; + if (!canApplyTreeMetadataResult(options.node)) return; + + const loadMoreNodes = (options.node.children || []).filter((child) => child.type === "load-more"); + const currentChildren = withoutLoadMoreNodes(options.node.children); + const mergedChildren = mergeLocatedTreeChildren(options.node, currentChildren, supplementalChildren, options.connectionId, options.database); + const nextChildren = [...mergedChildren, ...loadMoreNodes]; + setChildren(options.node, nextChildren); + await savePersistedTreeChildren(options.cacheKey, nextChildren); + } catch (error) { + // Some drivers only expose table metadata; keep the already-rendered table tree usable. + console.debug("[DBX][metadata:simple-supplemental:error]", { + connectionId: options.connectionId, + database: options.database, + schema: options.effectiveSchema, + error, + }); + } + } + function refreshStaleTreeNode(node: TreeNode) { const searchFilter = sidebarSearchQuery.value || ""; if (searchFilter) return; @@ -1071,18 +1252,35 @@ export const useConnectionStore = defineStore("connection", () => { } async function loadPersistedTreeChildren(node: TreeNode, cacheKey: string): Promise { + const trace = createMetadataLoadTrace({ + kind: "persisted-tree-cache", + connectionId: node.connectionId, + database: node.database, + schema: node.schema, + nodeKind: node.type, + extra: { cacheKey }, + }); const payload = await api.loadSchemaCache(cacheKey).catch(() => null); const decoded = decodeSchemaTreeCache(payload); - if (!decoded) return { hit: false, isStale: false }; + if (!decoded) { + logMetadataLoadTrace(metadataTraceLogger, trace, "cache-miss", { cacheStatus: "miss" }); + return { hit: false, isStale: false }; + } const config = node.connectionId ? getConfig(node.connectionId) : undefined; const cachedChildren = normalizeCataloglessDatabaseNodes(expandCachedObjectBrowserNodes(decoded.children)); const childrenWithLinkedServers = node.type === "connection" && node.connectionId ? ensureSqlServerLinkedRootNode(node.connectionId, cachedChildren, config) : cachedChildren; if (node.type === "connection" && !hasConnectionMetadataChildren(childrenWithLinkedServers)) { + logMetadataLoadTrace(metadataTraceLogger, trace, "cache-miss", { cacheStatus: "miss", resultCount: 0 }); return { hit: false, isStale: false }; } const normalizedChildren = sortSidebarTreeChildrenForParent(node, childrenWithLinkedServers, config?.db_type); setChildren(node, node.type === "connection" && node.connectionId ? withSavedSqlRoot(node.connectionId, normalizedChildren, node) : normalizedChildren); node.isExpanded = true; + logMetadataLoadTrace(metadataTraceLogger, trace, "cache-hit", { + cacheStatus: decoded.isStale ? "stale" : "hit", + resultCount: normalizedChildren.length, + stale: decoded.isStale, + }); return { hit: true, isStale: decoded.isStale }; } @@ -1132,12 +1330,19 @@ export const useConnectionStore = defineStore("connection", () => { return loadedTreeNodeChildrenIds.value.has(nodeId); } + function canApplyTreeMetadataResult(node: TreeNode): boolean { + if (findNode(treeNodes.value, node.id) !== node) return false; + if (node.connectionId && !connectedIds.value.has(node.connectionId)) return false; + return true; + } + function clearLoadedChildrenCache(prefix: string) { for (const id of loadedTreeNodeChildrenIds.value) { if (id === prefix || id.startsWith(`${prefix}:`)) { loadedTreeNodeChildrenIds.value.delete(id); } } + invalidateMetadataCachesByTreePrefix(prefix); const rawPrefix = `${prefix}:`; const encodedPrefix = `${schemaCacheKey(prefix)}:`; if (rawPrefix === encodedPrefix) { @@ -1235,6 +1440,7 @@ export const useConnectionStore = defineStore("connection", () => { } function invalidateCompletionCache(connectionId: string, database?: string) { + invalidateMetadataCaches({ connectionId, database }); const cachePrefix = database == null ? `${connectionId}:` : `${connectionId}:${database}:`; const exactCacheKey = database == null ? null : `${connectionId}:${database}`; for (const key of Object.keys(completionTablesCache.value)) { @@ -1688,105 +1894,119 @@ export const useConnectionStore = defineStore("connection", () => { } async function loadDatabases(connectionId: string, options?: LoadTreeOptions) { - const node = findNode(treeNodes.value, connectionId); - if (!node) return; - node.isLoading = true; - try { - await ensureConnected(connectionId); - if (useCachedChildren(node, options)) return; + const configForScope = getConfig(connectionId); + return runTreeMetadataLoad( + { + kind: "connection-databases", + connectionId, + driverProfile: metadataDriverProfile(configForScope), + }, + async () => { + const node = findNode(treeNodes.value, connectionId); + if (!node) return; + node.isLoading = true; + try { + await ensureConnected(connectionId); + if (useCachedChildren(node, options)) return; - const config = getConfig(connectionId); - if (config?.db_type === "duckdb") { - const cacheKey = schemaCacheKey(connectionId, "duckdb-root"); - if (!options?.force) { - const cached = await loadPersistedTreeChildren(node, cacheKey); - if (cached.hit) { - if (cached.isStale) refreshStaleTreeNode(node); - return; - } - } - const [databases, schemas] = await Promise.all([withMetadataLoadTimeout(connectionId, api.listDatabases(connectionId), "databases"), withMetadataLoadTimeout(connectionId, api.listSchemas(connectionId, "main"), "schemas")]); - const children = withSavedSqlRoot(connectionId, buildDuckDbConnectionTreeNodes(connectionId, databases, schemas), node); - if (isSidebarSearchQueryChanged(options)) return; - setChildren(node, children); - await savePersistedConnectionTreeChildren(cacheKey, node.children || children); - } else if (config && connectionUsesVisibleSchemaFilter(config)) { - const schemaFilterConfig = config; - const effectiveDb = schemaFilterConfig.database || ""; - const cacheKey = schemaCacheKey(connectionId, effectiveDb, "schemas"); - if (!options?.force) { - const cached = await loadPersistedTreeChildren(node, cacheKey); - if (cached.hit) { - if (cached.isStale) refreshStaleTreeNode(node); - return; - } - } - const schemas = await withMetadataLoadTimeout(connectionId, api.listSchemas(connectionId, effectiveDb, true), "schemas"); - const visibleSchemas = filterSchemaNamesForConnection(schemas, schemaFilterConfig, effectiveDb || ""); - const schemaNodes: TreeNode[] = sortSidebarNames(visibleSchemas).map((s) => ({ - id: `${connectionId}:${s}:${s}`, - label: s, - type: "schema" as const, - connectionId, - database: s, - schema: s, - isExpanded: false, - children: [], - })); - if (isSidebarSearchQueryChanged(options)) return; - setChildren(node, withSavedSqlRoot(connectionId, schemaNodes, node)); - await savePersistedConnectionTreeChildren(cacheKey, node.children || schemaNodes); - } else { - const cacheKey = schemaCacheKey(connectionId, "databases"); - if (!options?.force) { - const cached = await loadPersistedTreeChildren(node, cacheKey); - if (cached.hit) { - if (cached.isStale) refreshStaleTreeNode(node); - return; - } - } - const databases = await withMetadataLoadTimeout(connectionId, api.listDatabases(connectionId), "databases"); - const visibleNames = filterDatabaseNamesForConnection( - databases.map((database) => database.name), - config, - ); - const visibleNameSet = new Set(visibleNames); - const visibleDatabases = databases.filter((database) => visibleNameSet.has(database.name)); - const effectiveDbType = effectiveDatabaseTypeForConnection(config); - const databaseNodes = buildDatabaseTreeNodes(connectionId, visibleDatabases, { - includeDefaultWhenEmpty: usesTreeSchemaMode(effectiveDbType) || shouldIncludeDefaultDatabaseNode(config, visibleDatabases), - }); - if (config?.db_type === "sqlserver") { - const linkedServers = await withMetadataLoadTimeout(connectionId, api.listSqlServerLinkedServers(connectionId), "linked servers").catch(() => []); - const linkedDatabase = sqlServerLinkedRuntimeDatabase(config); - databaseNodes.push({ - ...sqlServerLinkedRootNode(connectionId, linkedDatabase), - children: linkedServers.map((server) => ({ - id: sqlServerLinkedServerId(connectionId, server.name), - label: server.name, - type: "linked-server", + const config = getConfig(connectionId); + if (config?.db_type === "duckdb") { + const cacheKey = schemaCacheKey(connectionId, "duckdb-root"); + if (!options?.force) { + const cached = await loadPersistedTreeChildren(node, cacheKey); + if (cached.hit) { + if (cached.isStale) refreshStaleTreeNode(node); + return; + } + } + const [databases, schemas] = await Promise.all([withMetadataLoadTimeout(connectionId, api.listDatabases(connectionId), "databases"), withMetadataLoadTimeout(connectionId, api.listSchemas(connectionId, "main"), "schemas")]); + const children = withSavedSqlRoot(connectionId, buildDuckDbConnectionTreeNodes(connectionId, databases, schemas), node); + if (isSidebarSearchQueryChanged(options)) return; + if (!canApplyTreeMetadataResult(node)) return; + setChildren(node, children); + await savePersistedConnectionTreeChildren(cacheKey, node.children || children); + } else if (config && connectionUsesVisibleSchemaFilter(config)) { + const schemaFilterConfig = config; + const effectiveDb = schemaFilterConfig.database || ""; + const cacheKey = schemaCacheKey(connectionId, effectiveDb, "schemas"); + if (!options?.force) { + const cached = await loadPersistedTreeChildren(node, cacheKey); + if (cached.hit) { + if (cached.isStale) refreshStaleTreeNode(node); + return; + } + } + const schemas = await withMetadataLoadTimeout(connectionId, api.listSchemas(connectionId, effectiveDb, true), "schemas"); + const visibleSchemas = filterSchemaNamesForConnection(schemas, schemaFilterConfig, effectiveDb || ""); + const schemaNodes: TreeNode[] = sortSidebarNames(visibleSchemas).map((s) => ({ + id: `${connectionId}:${s}:${s}`, + label: s, + type: "schema" as const, connectionId, - database: linkedDatabase, - linkedServer: server.name, - comment: [server.product, server.provider, server.data_source].filter(Boolean).join(" / ") || null, + database: s, + schema: s, isExpanded: false, children: [], - })), - }); - if (linkedServers.length > 0) loadedTreeNodeChildrenIds.value.add(sqlServerLinkedRootId(connectionId)); + })); + if (isSidebarSearchQueryChanged(options)) return; + if (!canApplyTreeMetadataResult(node)) return; + setChildren(node, withSavedSqlRoot(connectionId, schemaNodes, node)); + await savePersistedConnectionTreeChildren(cacheKey, node.children || schemaNodes); + } else { + const cacheKey = schemaCacheKey(connectionId, "databases"); + if (!options?.force) { + const cached = await loadPersistedTreeChildren(node, cacheKey); + if (cached.hit) { + if (cached.isStale) refreshStaleTreeNode(node); + return; + } + } + const databases = await withMetadataLoadTimeout(connectionId, api.listDatabases(connectionId), "databases"); + const visibleNames = filterDatabaseNamesForConnection( + databases.map((database) => database.name), + config, + ); + const visibleNameSet = new Set(visibleNames); + const visibleDatabases = databases.filter((database) => visibleNameSet.has(database.name)); + const effectiveDbType = effectiveDatabaseTypeForConnection(config); + const databaseNodes = buildDatabaseTreeNodes(connectionId, visibleDatabases, { + includeDefaultWhenEmpty: usesTreeSchemaMode(effectiveDbType) || shouldIncludeDefaultDatabaseNode(config, visibleDatabases), + }); + if (config?.db_type === "sqlserver") { + const linkedServers = await withMetadataLoadTimeout(connectionId, api.listSqlServerLinkedServers(connectionId), "linked servers").catch(() => []); + const linkedDatabase = sqlServerLinkedRuntimeDatabase(config); + databaseNodes.push({ + ...sqlServerLinkedRootNode(connectionId, linkedDatabase), + children: linkedServers.map((server) => ({ + id: sqlServerLinkedServerId(connectionId, server.name), + label: server.name, + type: "linked-server", + connectionId, + database: linkedDatabase, + linkedServer: server.name, + comment: [server.product, server.provider, server.data_source].filter(Boolean).join(" / ") || null, + isExpanded: false, + children: [], + })), + }); + if (linkedServers.length > 0) loadedTreeNodeChildrenIds.value.add(sqlServerLinkedRootId(connectionId)); + } + const children = withSavedSqlRoot(connectionId, databaseNodes, node); + if (isSidebarSearchQueryChanged(options)) return; + if (!canApplyTreeMetadataResult(node)) return; + setChildren(node, children); + await savePersistedConnectionTreeChildren(cacheKey, node.children || children); + } + node.isExpanded = true; + } catch (e) { + recordMetadataLoadError(connectionId, e); + throw e; + } finally { + node.isLoading = false; } - const children = withSavedSqlRoot(connectionId, databaseNodes, node); - if (isSidebarSearchQueryChanged(options)) return; - setChildren(node, children); - await savePersistedConnectionTreeChildren(cacheKey, node.children || children); - } - node.isExpanded = true; - } catch (e) { - recordMetadataLoadError(connectionId, e); - throw e; - } finally { - node.isLoading = false; - } + }, + options, + ); } async function loadRedisDatabases(connectionId: string) { @@ -2163,56 +2383,69 @@ export const useConnectionStore = defineStore("connection", () => { } async function loadSchemas(connectionId: string, database: string, options?: LoadTreeOptions) { - const nodeId = `${connectionId}:${database}`; - const node = findNode(treeNodes.value, nodeId); - if (!node) return; - node.isLoading = true; - try { - await ensureConnected(connectionId); - if (useCachedChildren(node, options)) return; - const cacheKey = schemaCacheKey(connectionId, database, "schemas-v2"); - if (!options?.force) { - const cached = await loadPersistedTreeChildren(node, cacheKey); - if (cached.hit) { - if (cached.isStale) refreshStaleTreeNode(node); - return; - } - } + const configForScope = getConfig(connectionId); + return runTreeMetadataLoad( + { + kind: "database-schemas", + connectionId, + database, + driverProfile: metadataDriverProfile(configForScope), + }, + async () => { + const nodeId = `${connectionId}:${database}`; + const node = findNode(treeNodes.value, nodeId); + if (!node) return; + node.isLoading = true; + try { + await ensureConnected(connectionId); + if (useCachedChildren(node, options)) return; + const cacheKey = schemaCacheKey(connectionId, database, "schemas-v2"); + if (!options?.force) { + const cached = await loadPersistedTreeChildren(node, cacheKey); + if (cached.hit) { + if (cached.isStale) refreshStaleTreeNode(node); + return; + } + } - const schemas = sortSidebarSchemaInfos(await api.listSchemaInfos(connectionId, database)); - const visibleSchemaNames = new Set( - filterSchemaNamesForConnection( - schemas.map((schema) => schema.name), - getConfig(connectionId), - database, - ), - ); - const children = schemas - .filter((schema) => visibleSchemaNames.has(schema.name)) - .map((schema) => { - const s = schema.name; - return { - id: `${connectionId}:${database}:${s}`, - label: s, - type: "schema" as const, - connectionId, - database, - schema: s, - comment: schema.comment, - isExpanded: false, - children: [], - }; - }); - if (isSidebarSearchQueryChanged(options)) return; - setChildren(node, children); - await savePersistedTreeChildren(cacheKey, children); - node.isExpanded = true; - } catch (e) { - recordMetadataLoadError(connectionId, e); - throw e; - } finally { - node.isLoading = false; - } + const schemas = sortSidebarSchemaInfos(await api.listSchemaInfos(connectionId, database)); + const visibleSchemaNames = new Set( + filterSchemaNamesForConnection( + schemas.map((schema) => schema.name), + getConfig(connectionId), + database, + ), + ); + const children = schemas + .filter((schema) => visibleSchemaNames.has(schema.name)) + .map((schema) => { + const s = schema.name; + return { + id: `${connectionId}:${database}:${s}`, + label: s, + type: "schema" as const, + connectionId, + database, + schema: s, + comment: schema.comment, + isExpanded: false, + children: [], + }; + }); + if (isSidebarSearchQueryChanged(options)) return; + if (!canApplyTreeMetadataResult(node)) return; + setChildren(node, children); + await savePersistedTreeChildren(cacheKey, children); + node.isExpanded = true; + } catch (e) { + recordMetadataLoadError(connectionId, e); + throw e; + } finally { + node.isLoading = false; + } + }, + options, + ); } async function loadSqlServerDatabaseObjects(connectionId: string, database: string, options?: LoadTreeOptions) { @@ -2359,263 +2592,386 @@ export const useConnectionStore = defineStore("connection", () => { } async function loadTables(connectionId: string, database: string, schema?: string, options?: LoadTreeOptions) { - const nodeId = schema ? `${connectionId}:${database}:${schema}` : `${connectionId}:${database}`; - const node = findNode(treeNodes.value, nodeId); - if (!node) return; - node.isLoading = true; - try { - await ensureConnected(connectionId); - if (useCachedChildren(node, options)) return; - const simpleObjectDisplay = useSettingsStore().editorSettings.sidebarObjectDisplay === "simple"; - const cacheKey = schemaCacheKey(connectionId, database, schema || "", simpleObjectDisplay ? "objects-simple-v3" : "objects-grouped-v3"); - const searchFilter = activeTreeLoadSearchFilter(options); - const isSidebarTableSearch = !!options?.sidebarTableSearchParentId; - if (!options?.force && !searchFilter) { - const cached = await loadPersistedTreeChildren(node, cacheKey); - if (cached.hit) { - if (cached.isStale) refreshStaleTreeNode(node); - return; - } - } + const configForScope = getConfig(connectionId); + const simpleObjectDisplayForScope = useSettingsStore().editorSettings.sidebarObjectDisplay === "simple"; + const objectTypesForScope = simpleObjectDisplayForScope ? supportedSidebarObjectTypes(configForScope) : undefined; + return runTreeMetadataLoad( + { + kind: "schema-tables", + connectionId, + database, + schema, + nodeKind: schema ? "schema" : "database", + objectTypes: objectTypesForScope, + searchFilter: activeTreeLoadSearchFilter(options), + limit: simpleObjectDisplayForScope ? sidebarObjectGroupPageSize() + 1 : undefined, + offset: 0, + sidebarDisplayMode: simpleObjectDisplayForScope ? "simple" : "grouped", + driverProfile: metadataDriverProfile(configForScope), + extra: options?.sidebarTableSearchParentId ? { sidebarTableSearchParentId: options.sidebarTableSearchParentId } : undefined, + }, + async () => { + const nodeId = schema ? `${connectionId}:${database}:${schema}` : `${connectionId}:${database}`; + const node = findNode(treeNodes.value, nodeId); + if (!node) return; + node.isLoading = true; + try { + await ensureConnected(connectionId); + if (useCachedChildren(node, options)) return; + const simpleObjectDisplay = useSettingsStore().editorSettings.sidebarObjectDisplay === "simple"; + const cacheKey = schemaCacheKey(connectionId, database, schema || "", simpleObjectDisplay ? "objects-simple-v3" : "objects-grouped-v3"); + const searchFilter = activeTreeLoadSearchFilter(options); + const isSidebarTableSearch = !!options?.sidebarTableSearchParentId; + if (!options?.force && !searchFilter) { + const cached = await loadPersistedTreeChildren(node, cacheKey); + if (cached.hit) { + if (cached.isStale) refreshStaleTreeNode(node); + return; + } + } - const config = getConfig(connectionId); - const querySchema = connectionObjectTreeQuerySchema(config, database, schema); - const effectiveSchema = connectionObjectTreeNodeSchema(config, database, schema); - const nonTableObjectTypes = simpleObjectDisplay ? supportedSidebarObjectTypes(config).filter((objectType) => objectType !== "TABLE") : []; - let children: TreeNode[]; - if (simpleObjectDisplay) { - const pageSize = sidebarObjectGroupPageSize(); - const page = await loadPagedSimpleTableChildren({ - nodeId, - connectionId, - database, - querySchema, - effectiveSchema, - nonTableObjectTypes, - offset: 0, - pageSize, - searchFilter: searchFilter || undefined, - }); - children = page.hasMore && !searchFilter ? [...page.children, buildLoadMoreNode(node, page.nextOffset, pageSize)] : page.children; - node.objectCount = page.objectCount; - } else { - children = buildObjectGroupPlaceholderNodes({ - nodeId, - connectionId, - database, - schema: effectiveSchema, - objectTypes: supportedSidebarObjectTypes(config), - }); - } - if (isTreeLoadSearchChanged(searchFilter, options)) return; - setChildren(node, children); - if (!searchFilter && !isSidebarTableSearch) { - await savePersistedTreeChildren(cacheKey, children); - } - node.isExpanded = true; - } catch (e) { - recordMetadataLoadError(connectionId, e); - throw e; - } finally { - node.isLoading = false; - } + const config = getConfig(connectionId); + const querySchema = connectionObjectTreeQuerySchema(config, database, schema); + const effectiveSchema = connectionObjectTreeNodeSchema(config, database, schema); + const nonTableObjectTypes = simpleObjectDisplay ? supportedSidebarObjectTypes(config).filter((objectType) => objectType !== "TABLE") : []; + let children: TreeNode[]; + let nextObjectCount: number | undefined; + if (simpleObjectDisplay) { + const pageSize = sidebarObjectGroupPageSize(); + const page = await loadPagedSimpleTableChildren({ + nodeId, + connectionId, + database, + querySchema, + effectiveSchema, + nonTableObjectTypes, + offset: 0, + pageSize, + searchFilter: searchFilter || undefined, + force: options?.force, + }); + children = page.hasMore && !searchFilter ? [...page.children, buildLoadMoreNode(node, page.nextOffset, pageSize)] : page.children; + nextObjectCount = page.objectCount; + } else { + children = buildObjectGroupPlaceholderNodes({ + nodeId, + connectionId, + database, + schema: effectiveSchema, + objectTypes: supportedSidebarObjectTypes(config), + }); + } + if (isTreeLoadSearchChanged(searchFilter, options)) return; + if (!canApplyTreeMetadataResult(node)) return; + if (nextObjectCount !== undefined) node.objectCount = nextObjectCount; + setChildren(node, children); + if (!searchFilter && !isSidebarTableSearch) { + await savePersistedTreeChildren(cacheKey, children); + } + node.isExpanded = true; + if (simpleObjectDisplay && !searchFilter && !isSidebarTableSearch && nonTableObjectTypes.length > 0) { + void loadSimpleSupplementalObjectChildren({ + node, + nodeId, + connectionId, + database, + querySchema, + effectiveSchema, + objectTypes: nonTableObjectTypes, + cacheKey, + loadOptions: options, + }); + } + } catch (e) { + recordMetadataLoadError(connectionId, e); + throw e; + } finally { + node.isLoading = false; + } + }, + options, + ); } async function loadObjectGroupChildren(node: TreeNode, options?: LoadTreeOptions) { - if (!node.connectionId || !hasTreeNodeDatabaseContext(node)) return; - node.isLoading = true; - try { - await ensureConnected(node.connectionId); - if (useCachedChildren(node, options)) return; - const objectTypes = objectTypesForGroupNode(node.type); - const parentNodeId = objectGroupRefreshParentId(node); - if (!objectTypes || !parentNodeId) return; + const configForScope = node.connectionId ? getConfig(node.connectionId) : undefined; + const objectTypesForScope = objectTypesForGroupNode(node.type); + return runTreeMetadataLoad( + { + kind: "object-group", + connectionId: node.connectionId, + database: node.database, + schema: node.schema, + nodeKind: node.type, + objectTypes: objectTypesForScope, + searchFilter: activeTreeLoadSearchFilter(options), + limit: sidebarObjectGroupPageSize() + 1, + offset: 0, + sidebarDisplayMode: useSettingsStore().editorSettings.sidebarObjectDisplay, + driverProfile: metadataDriverProfile(configForScope), + extra: options?.sidebarTableSearchParentId ? { sidebarTableSearchParentId: options.sidebarTableSearchParentId } : undefined, + }, + async () => { + if (!node.connectionId || !hasTreeNodeDatabaseContext(node)) return; + node.isLoading = true; + try { + await ensureConnected(node.connectionId); + if (useCachedChildren(node, options)) return; + const objectTypes = objectTypesForGroupNode(node.type); + const parentNodeId = objectGroupRefreshParentId(node); + if (!objectTypes || !parentNodeId) return; - const config = getConfig(node.connectionId); - const querySchema = connectionObjectTreeQuerySchema(config, node.database, node.schema); - const effectiveSchema = connectionObjectTreeNodeSchema(config, node.database, node.schema); - const cacheKey = objectGroupCacheKey(node); - const searchFilter = activeTreeLoadSearchFilter(options); - const isSidebarTableSearch = !!options?.sidebarTableSearchParentId; - if (!options?.force && !searchFilter) { - const cached = await loadPersistedTreeChildren(node, cacheKey); - if (cached.hit) { - if (cached.isStale) refreshStaleTreeNode(node); - return; + const config = getConfig(node.connectionId); + const querySchema = connectionObjectTreeQuerySchema(config, node.database, node.schema); + const effectiveSchema = connectionObjectTreeNodeSchema(config, node.database, node.schema); + const cacheKey = objectGroupCacheKey(node); + const searchFilter = activeTreeLoadSearchFilter(options); + const isSidebarTableSearch = !!options?.sidebarTableSearchParentId; + if (!options?.force && !searchFilter) { + const cached = await loadPersistedTreeChildren(node, cacheKey); + if (cached.hit) { + if (cached.isStale) refreshStaleTreeNode(node); + return; + } + } + + const wantsOnlyTablesOrViews = objectTypes.every((objectType) => objectType === "TABLE" || objectType === "VIEW" || objectType === "MATERIALIZED_VIEW"); + let children: TreeNode[]; + let nextObjectCount: number; + if (wantsOnlyTablesOrViews) { + const page = await loadPagedTableGroupChildren({ + node, + parentNodeId, + querySchema, + effectiveSchema, + objectTypes, + offset: 0, + pageSize: sidebarObjectGroupPageSize(), + searchFilter: searchFilter || undefined, + force: options?.force, + }); + children = page.hasMore && !searchFilter ? [...page.children, buildLoadMoreNode(node, page.nextOffset, sidebarObjectGroupPageSize())] : page.children; + nextObjectCount = page.objectCount; + } else { + const objects = await loadCachedMetadataListPage( + metadataListCacheScope({ + kind: "object-list-page", + connectionId: node.connectionId, + database: node.database, + schema: querySchema, + nodeKind: node.type, + objectTypes, + searchFilter: searchFilter || undefined, + sidebarDisplayMode: useSettingsStore().editorSettings.sidebarObjectDisplay, + }), + () => api.listObjects(node.connectionId!, node.database!, querySchema, objectTypes, searchFilter || undefined), + { force: options?.force }, + ); + children = objectGroupChildrenFromObjects({ + node, + parentNodeId, + effectiveSchema, + objectTypes, + objects, + }); + nextObjectCount = children.length; + } + if (isTreeLoadSearchChanged(searchFilter, options)) return; + if (!canApplyTreeMetadataResult(node)) return; + node.objectCount = nextObjectCount; + setChildren(node, children); + if (!searchFilter && !isSidebarTableSearch) { + await savePersistedTreeChildren(cacheKey, children); + } + node.isExpanded = true; + } catch (e) { + recordMetadataLoadError(node.connectionId, e); + throw e; + } finally { + node.isLoading = false; } - } - - const wantsOnlyTablesOrViews = objectTypes.every((objectType) => objectType === "TABLE" || objectType === "VIEW" || objectType === "MATERIALIZED_VIEW"); - let children: TreeNode[]; - if (wantsOnlyTablesOrViews) { - const page = await loadPagedTableGroupChildren({ - node, - parentNodeId, - querySchema, - effectiveSchema, - objectTypes, - offset: 0, - pageSize: sidebarObjectGroupPageSize(), - searchFilter: searchFilter || undefined, - }); - children = page.hasMore && !searchFilter ? [...page.children, buildLoadMoreNode(node, page.nextOffset, sidebarObjectGroupPageSize())] : page.children; - node.objectCount = page.objectCount; - } else { - const objects = await api.listObjects(node.connectionId, node.database, querySchema, objectTypes); - children = objectGroupChildrenFromObjects({ - node, - parentNodeId, - effectiveSchema, - objectTypes, - objects, - }); - node.objectCount = children.length; - } - if (isTreeLoadSearchChanged(searchFilter, options)) return; - setChildren(node, children); - if (!searchFilter && !isSidebarTableSearch) { - await savePersistedTreeChildren(cacheKey, children); - } - node.isExpanded = true; - } catch (e) { - recordMetadataLoadError(node.connectionId, e); - throw e; - } finally { - node.isLoading = false; - } + }, + options, + ); } async function loadMoreObjectGroupChildren(node: TreeNode) { if (node.type !== "load-more" || !node.loadMore) return; + const loadMore = node.loadMore; const parent = findNode(treeNodes.value, node.loadMore.parentId); if (!parent?.connectionId || !hasTreeNodeDatabaseContext(parent)) return; - node.isLoading = true; - try { - await ensureConnected(parent.connectionId); - if (parent.type === "database" || parent.type === "schema" || parent.type === "linked-server-schema") { - const parentDatabase = parent.database; - if (!parentDatabase) return; - const config = getConfig(parent.connectionId); - const querySchema = connectionObjectTreeQuerySchema(config, parentDatabase, parent.schema); - const effectiveSchema = connectionObjectTreeNodeSchema(config, parentDatabase, parent.schema); - const page = await loadPagedSimpleTableChildren({ - nodeId: parent.schema ? `${parent.connectionId}:${parentDatabase}:${parent.schema}` : `${parent.connectionId}:${parentDatabase}`, - connectionId: parent.connectionId, - database: parentDatabase, - querySchema, - effectiveSchema, - nonTableObjectTypes: [], - offset: node.loadMore.offset, - pageSize: node.loadMore.pageSize, - }); - const currentChildren = withoutLoadMoreNodes(parent.children); - const mergedChildren = mergeTableTreePageChildren(currentChildren, page.children, parent.connectionId, parentDatabase); - const nextChildren = page.hasMore ? [...mergedChildren, buildLoadMoreNode(parent, page.nextOffset, node.loadMore.pageSize)] : mergedChildren; - parent.objectCount = (parent.objectCount ?? currentChildren.length) + page.objectCount; - setChildren(parent, nextChildren); - await savePersistedTreeChildren(schemaCacheKey(parent.connectionId, parentDatabase, parent.schema || "", "objects-simple-v3"), nextChildren); - parent.isExpanded = true; - return; - } - const objectTypes = objectTypesForGroupNode(parent.type); - const parentNodeId = objectGroupRefreshParentId(parent); - if (!objectTypes || !parentNodeId) return; - if (!objectTypes.every((objectType) => objectType === "TABLE" || objectType === "VIEW" || objectType === "MATERIALIZED_VIEW")) return; + const parentConnectionId = parent.connectionId; + const configForScope = getConfig(parentConnectionId); + const objectTypesForScope = objectTypesForGroupNode(parent.type); + return runTreeMetadataLoad( + { + kind: "object-group-page", + connectionId: parentConnectionId, + database: parent.database, + schema: parent.schema, + nodeKind: parent.type, + objectTypes: objectTypesForScope, + limit: loadMore.pageSize + 1, + offset: loadMore.offset, + sidebarDisplayMode: useSettingsStore().editorSettings.sidebarObjectDisplay, + driverProfile: metadataDriverProfile(configForScope), + }, + async () => { + node.isLoading = true; + try { + await ensureConnected(parentConnectionId); + if (parent.type === "database" || parent.type === "schema" || parent.type === "linked-server-schema") { + const parentDatabase = parent.database; + if (!parentDatabase) return; + const config = getConfig(parentConnectionId); + const querySchema = connectionObjectTreeQuerySchema(config, parentDatabase, parent.schema); + const effectiveSchema = connectionObjectTreeNodeSchema(config, parentDatabase, parent.schema); + const page = await loadPagedSimpleTableChildren({ + nodeId: parent.schema ? `${parentConnectionId}:${parentDatabase}:${parent.schema}` : `${parentConnectionId}:${parentDatabase}`, + connectionId: parentConnectionId, + database: parentDatabase, + querySchema, + effectiveSchema, + nonTableObjectTypes: [], + offset: loadMore.offset, + pageSize: loadMore.pageSize, + force: false, + }); + const currentChildren = withoutLoadMoreNodes(parent.children); + const mergedChildren = mergeTableTreePageChildren(currentChildren, page.children, parentConnectionId, parentDatabase); + const nextChildren = page.hasMore ? [...mergedChildren, buildLoadMoreNode(parent, page.nextOffset, loadMore.pageSize)] : mergedChildren; + if (!canApplyTreeMetadataResult(parent)) return; + parent.objectCount = (parent.objectCount ?? currentChildren.length) + page.objectCount; + setChildren(parent, nextChildren); + await savePersistedTreeChildren(schemaCacheKey(parentConnectionId, parentDatabase, parent.schema || "", "objects-simple-v3"), nextChildren); + parent.isExpanded = true; + return; + } + const objectTypes = objectTypesForGroupNode(parent.type); + const parentNodeId = objectGroupRefreshParentId(parent); + if (!objectTypes || !parentNodeId) return; + if (!objectTypes.every((objectType) => objectType === "TABLE" || objectType === "VIEW" || objectType === "MATERIALIZED_VIEW")) return; - const config = getConfig(parent.connectionId); - const querySchema = connectionObjectTreeQuerySchema(config, parent.database, parent.schema); - const effectiveSchema = connectionObjectTreeNodeSchema(config, parent.database, parent.schema); - const page = await loadPagedTableGroupChildren({ - node: parent, - parentNodeId, - querySchema, - effectiveSchema, - objectTypes, - offset: node.loadMore.offset, - pageSize: node.loadMore.pageSize, - }); - const currentChildren = withoutLoadMoreNodes(parent.children); - const mergedChildren = mergeTableTreePageChildren(currentChildren, page.children, parent.connectionId, parent.database); - const nextChildren = page.hasMore ? [...mergedChildren, buildLoadMoreNode(parent, page.nextOffset, node.loadMore.pageSize)] : mergedChildren; - parent.objectCount = (parent.objectCount ?? currentChildren.length) + page.objectCount; - setChildren(parent, nextChildren); - await savePersistedTreeChildren(objectGroupCacheKey(parent), nextChildren); - parent.isExpanded = true; - } catch (e) { - recordMetadataLoadError(parent.connectionId, e); - throw e; - } finally { - node.isLoading = false; - } + const config = getConfig(parentConnectionId); + const parentDatabase = parent.database; + if (!parentDatabase) return; + const querySchema = connectionObjectTreeQuerySchema(config, parentDatabase, parent.schema); + const effectiveSchema = connectionObjectTreeNodeSchema(config, parentDatabase, parent.schema); + const page = await loadPagedTableGroupChildren({ + node: parent, + parentNodeId, + querySchema, + effectiveSchema, + objectTypes, + offset: loadMore.offset, + pageSize: loadMore.pageSize, + force: false, + }); + const currentChildren = withoutLoadMoreNodes(parent.children); + const mergedChildren = mergeTableTreePageChildren(currentChildren, page.children, parentConnectionId, parentDatabase); + const nextChildren = page.hasMore ? [...mergedChildren, buildLoadMoreNode(parent, page.nextOffset, loadMore.pageSize)] : mergedChildren; + if (!canApplyTreeMetadataResult(parent)) return; + parent.objectCount = (parent.objectCount ?? currentChildren.length) + page.objectCount; + setChildren(parent, nextChildren); + await savePersistedTreeChildren(objectGroupCacheKey(parent), nextChildren); + parent.isExpanded = true; + } catch (e) { + recordMetadataLoadError(parentConnectionId, e); + throw e; + } finally { + node.isLoading = false; + } + }, + ); } async function loadTableForLocate(target: LocateTableTarget): Promise { const config = getConfig(target.connectionId); if (!config) return false; - await ensureConnected(target.connectionId); - - const querySchema = connectionObjectTreeQuerySchema(config, target.database, target.schema); - const effectiveSchema = connectionObjectTreeNodeSchema(config, target.database, target.schema); - const pageSize = sidebarObjectGroupPageSize(); - const simpleObjectDisplay = useSettingsStore().editorSettings.sidebarObjectDisplay === "simple"; - let loaded = false; - - if (simpleObjectDisplay) { - const parentId = target.schema ? `${target.connectionId}:${target.database}:${target.schema}` : `${target.connectionId}:${target.database}`; - const parent = findNode(treeNodes.value, parentId); - if (!parent) return false; - const page = await loadPagedSimpleTableChildren({ - nodeId: parentId, + return runTreeMetadataLoad( + { + kind: "locate-target", connectionId: target.connectionId, database: target.database, - querySchema, - effectiveSchema, - nonTableObjectTypes: [], - offset: 0, - pageSize, + schema: target.schema, + tableName: target.tableName, searchFilter: target.tableName, - }); - if (!page.children.length) return false; - const currentChildren = withoutLoadMoreNodes(parent.children); - const loadMoreNodes = (parent.children || []).filter((child) => child.type === "load-more"); - const mergedChildren = mergeLocatedTreeChildren(parent, currentChildren, page.children, target.connectionId, target.database); - setChildren(parent, [...mergedChildren, ...loadMoreNodes]); - parent.objectCount = Math.max(parent.objectCount ?? currentChildren.length, mergedChildren.length); - parent.isExpanded = true; - return true; - } - - const matchingGroups = findTreeNodes(treeNodes.value, (node) => { - return (node.type === "group-tables" || node.type === "group-views" || node.type === "group-materialized-views") && node.connectionId === target.connectionId && sameSidebarObjectName(node.database, target.database) && (!target.schema || sameSidebarObjectName(node.schema, target.schema)); - }); - - for (const group of matchingGroups) { - const objectTypes = objectTypesForGroupNode(group.type); - const parentNodeId = objectGroupRefreshParentId(group); - if (!objectTypes || !parentNodeId) continue; - - const page = await loadPagedTableGroupChildren({ - node: group, - parentNodeId, - querySchema, - effectiveSchema, - objectTypes, + limit: sidebarObjectGroupPageSize() + 1, offset: 0, - pageSize, - searchFilter: target.tableName, - }); - if (!page.children.length) continue; + sidebarDisplayMode: useSettingsStore().editorSettings.sidebarObjectDisplay, + driverProfile: metadataDriverProfile(config), + }, + async () => { + await ensureConnected(target.connectionId); - const currentChildren = withoutLoadMoreNodes(group.children); - const loadMoreNodes = (group.children || []).filter((child) => child.type === "load-more"); - const mergedChildren = mergeLocatedTreeChildren(group, currentChildren, page.children, target.connectionId, target.database); - setChildren(group, [...mergedChildren, ...loadMoreNodes]); - group.objectCount = Math.max(group.objectCount ?? currentChildren.length, mergedChildren.length); - group.isExpanded = true; - loaded = true; - } + const querySchema = connectionObjectTreeQuerySchema(config, target.database, target.schema); + const effectiveSchema = connectionObjectTreeNodeSchema(config, target.database, target.schema); + const pageSize = sidebarObjectGroupPageSize(); + const simpleObjectDisplay = useSettingsStore().editorSettings.sidebarObjectDisplay === "simple"; + let loaded = false; - return loaded; + if (simpleObjectDisplay) { + const parentId = target.schema ? `${target.connectionId}:${target.database}:${target.schema}` : `${target.connectionId}:${target.database}`; + const parent = findNode(treeNodes.value, parentId); + if (!parent) return false; + const page = await loadPagedSimpleTableChildren({ + nodeId: parentId, + connectionId: target.connectionId, + database: target.database, + querySchema, + effectiveSchema, + nonTableObjectTypes: [], + offset: 0, + pageSize, + searchFilter: target.tableName, + force: false, + }); + if (!page.children.length) return false; + const currentChildren = withoutLoadMoreNodes(parent.children); + const loadMoreNodes = (parent.children || []).filter((child) => child.type === "load-more"); + const mergedChildren = mergeLocatedTreeChildren(parent, currentChildren, page.children, target.connectionId, target.database); + if (!canApplyTreeMetadataResult(parent)) return false; + setChildren(parent, [...mergedChildren, ...loadMoreNodes]); + parent.objectCount = Math.max(parent.objectCount ?? currentChildren.length, mergedChildren.length); + parent.isExpanded = true; + return true; + } + + const matchingGroups = findTreeNodes(treeNodes.value, (node) => { + return (node.type === "group-tables" || node.type === "group-views" || node.type === "group-materialized-views") && node.connectionId === target.connectionId && sameSidebarObjectName(node.database, target.database) && (!target.schema || sameSidebarObjectName(node.schema, target.schema)); + }); + + for (const group of matchingGroups) { + const objectTypes = objectTypesForGroupNode(group.type); + const parentNodeId = objectGroupRefreshParentId(group); + if (!objectTypes || !parentNodeId) continue; + + const page = await loadPagedTableGroupChildren({ + node: group, + parentNodeId, + querySchema, + effectiveSchema, + objectTypes, + offset: 0, + pageSize, + searchFilter: target.tableName, + force: false, + }); + if (!page.children.length) continue; + + const currentChildren = withoutLoadMoreNodes(group.children); + const loadMoreNodes = (group.children || []).filter((child) => child.type === "load-more"); + const mergedChildren = mergeLocatedTreeChildren(group, currentChildren, page.children, target.connectionId, target.database); + if (!canApplyTreeMetadataResult(group)) continue; + setChildren(group, [...mergedChildren, ...loadMoreNodes]); + group.objectCount = Math.max(group.objectCount ?? currentChildren.length, mergedChildren.length); + group.isExpanded = true; + loaded = true; + } + + return loaded; + }, + ); } async function loadAllObjectGroupChildren(parent: TreeNode) { @@ -3006,6 +3362,7 @@ export const useConnectionStore = defineStore("connection", () => { } async function refreshTreeNode(node: TreeNode) { + invalidateMetadataCachesForNode(node); if (objectTypesForGroupNode(node.type)) { clearLoadedChildrenCache(node.id); await loadObjectGroupChildren(node, { force: true }); @@ -3041,6 +3398,7 @@ export const useConnectionStore = defineStore("connection", () => { } async function refreshObjectListTreeNode(connectionId: string, database: string, schema?: string) { + invalidateMetadataCaches({ connectionId, database, schema }); const shouldRefreshSchemaNode = !!schema; const node = shouldRefreshSchemaNode ? findNode(treeNodes.value, `${connectionId}:${database}:${schema}`) : null; if (node) { @@ -4350,6 +4708,7 @@ export const useConnectionStore = defineStore("connection", () => { listMongoCompletionCollections, listMongoCompletionFields, invalidateCompletionCache, + invalidateMetadataCache, exportConnectionsToFile, readImportFile, importConnectionsFromFile, diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 02525b884f..afde47df30 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -27,10 +27,11 @@ import { redisCommandResultToQueryResult } from "@/lib/redis/redisQueryResult"; import { nextRedisCommandDb } from "@/lib/redis/redisCommandSession"; import { isRedisMutatingCommand } from "@/lib/redis/redisCommandTable"; import { usesAgentCursorForQuery } from "@/lib/database/databaseDriverManifest"; -import { canUseKeylessRowPredicate, editableRowIdentifierColumns } from "@/lib/table/tableEditing"; +import { canUseKeylessRowPredicate } from "@/lib/table/tableEditing"; import { TABLE_DATA_EXPORT_PAGE_SIZE } from "@/lib/table/tableDataExport"; import { tableMetaForDataTab } from "@/lib/table/tableDataTabMeta"; import { tableOpenPageLimit } from "@/lib/table/tableOpenPageLimit"; +import { loadTableMetadata } from "@/lib/metadata/tableMetadataCache"; import { quoteTableIdentifier } from "@/lib/table/tableSelectSql"; import { connectionUsesDatabaseObjectTreeMode, connectionUsesSchemaExecutionContext, effectiveDatabaseTypeForConnection, metadataSchemaForConnection } from "@/lib/database/jdbcDialect"; import { frontendQueryTimeoutSecsForSql, queryTimeoutSecsForConnection } from "@/lib/sql/queryTimeout"; @@ -1771,21 +1772,33 @@ export const useQueryStore = defineStore("query", () => { const metadataAnalysis = normalizeOracleLikeQueryAnalysis(dbType, analysis, metadataSchema || undefined, metadataTableName); try { - console.info("[DBX][executeTabSql:metadata:get-columns:start]", { + console.info("[DBX][executeTabSql:metadata:table:start]", { traceId, schema: metadataSchema, table: metadataTableName, elapsed: elapsed?.(), }); - const columns = await api.getColumns(tab.connectionId, tab.database, metadataSchema, metadataTableName); - console.info("[DBX][executeTabSql:metadata:get-columns:done]", { + const loadedMetadata = await loadTableMetadata({ + connectionId: tab.connectionId, + database: tab.database, + schema: metadataSchema, + tableName: metadataTableName, + tableType: tab.tableMeta?.tableType, + databaseType: dbType, + driverProfile: conn?.driver_profile || conn?.db_type, + traceLogger: (event) => console.debug("[DBX][executeTabSql:metadata:table-trace]", { sourceTraceId: traceId, ...event }), + }); + const columns = loadedMetadata.metadata.columns; + const primaryKeys = loadedMetadata.metadata.primaryKeys; + console.info("[DBX][executeTabSql:metadata:table:done]", { traceId, columnCount: columns.length, + primaryKeyCount: primaryKeys.length, + cacheStatus: loadedMetadata.cacheStatus, + ageMs: Math.round(loadedMetadata.ageMs), elapsed: elapsed?.(), }); - const indexes = await api.listIndexes(tab.connectionId, tab.database, metadataSchema, metadataTableName).catch(() => []); - const tableType = tab.tableMeta?.tableType; - const primaryKeys = editableRowIdentifierColumns(dbType as DatabaseType, columns, indexes, tableType); + const tableType = loadedMetadata.metadata.tableType; const tableMeta = { schema: metadataSchema || undefined, tableName: metadataTableName, diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index bc227c0e0d..163210a6fb 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -782,8 +782,17 @@ pub async fn export_database_sql_core( let mut functions: Vec = Vec::new(); if request.include_objects && request.selected_tables.is_empty() { - if let Ok(objects) = - crate::schema::list_objects_core(state, &request.connection_id, &request.database, &request.schema).await + if let Ok(objects) = crate::schema::list_objects_core( + state, + &request.connection_id, + &request.database, + &request.schema, + None, + None, + None, + None, + ) + .await { for obj in &objects { let ot = obj.object_type.to_uppercase(); diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index 9cc2d6af0d..6757da3594 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -628,8 +628,30 @@ impl AgentDriverClient { schema: &str, object_types: Option<&[String]>, timeout_duration: Option, + ) -> Result { + self.list_tables_constrained(database, schema, None, None, None, object_types, timeout_duration).await + } + + pub async fn list_tables_constrained( + &mut self, + database: &str, + schema: &str, + filter: Option<&str>, + limit: Option, + offset: Option, + object_types: Option<&[String]>, + timeout_duration: Option, ) -> Result { let mut params = agent_schema_params(database, schema); + if let Some(filter) = filter.map(str::trim).filter(|filter| !filter.is_empty()) { + params["filter"] = serde_json::json!(filter); + } + if let Some(limit) = limit { + params["limit"] = serde_json::json!(limit); + } + if let Some(offset) = offset { + params["offset"] = serde_json::json!(offset); + } if let Some(object_types) = object_types { params["object_types"] = serde_json::json!(object_types); } @@ -642,8 +664,33 @@ impl AgentDriverClient { schema: &str, timeout_duration: Option, ) -> Result { - self.call_method_with_timeout(AgentMethod::ListObjects, agent_schema_params(database, schema), timeout_duration) - .await + self.list_objects_constrained(database, schema, None, None, None, None, timeout_duration).await + } + + pub async fn list_objects_constrained( + &mut self, + database: &str, + schema: &str, + filter: Option<&str>, + limit: Option, + offset: Option, + object_types: Option<&[String]>, + timeout_duration: Option, + ) -> Result { + let mut params = agent_schema_params(database, schema); + if let Some(filter) = filter.map(str::trim).filter(|filter| !filter.is_empty()) { + params["filter"] = serde_json::json!(filter); + } + if let Some(limit) = limit { + params["limit"] = serde_json::json!(limit); + } + if let Some(offset) = offset { + params["offset"] = serde_json::json!(offset); + } + if let Some(object_types) = object_types { + params["object_types"] = serde_json::json!(object_types); + } + self.call_method_with_timeout(AgentMethod::ListObjects, params, timeout_duration).await } pub async fn list_data_types( diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 9741aef4cb..9d921af39d 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -1453,10 +1453,18 @@ async fn list_tables_once( .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } + let mut params = + serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }); + if let Some(filter) = filter.map(str::trim).filter(|value| !value.is_empty()) { + params["filter"] = serde_json::json!(filter); + } + if let Some(object_types) = object_types { + params["object_types"] = serde_json::json!(object_types); + } return session .invoke_with_timeout::>( "listTables", - serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }), + params, agent_metadata_timeout(Some(config.as_ref())), ) .await @@ -1520,16 +1528,33 @@ 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 use_oracle_agent_paging = db_config.as_ref().is_some_and(is_default_oracle_agent_config); 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_limit = if use_oracle_agent_paging { limit } else { None }; + let agent_offset = if use_oracle_agent_paging { offset } else { None }; match client - .list_tables_filtered::>(database, schema, object_types, timeout_duration) + .list_tables_constrained::>( + database, + schema, + filter, + agent_limit, + agent_offset, + object_types, + timeout_duration, + ) .await { Ok(tables) if !tables.is_empty() => { - let mut tables = filter_table_infos(tables, filter, limit, offset, object_types); + let final_offset = + if oracle_agent_paging_likely_applied(use_oracle_agent_paging, limit, tables.len()) { + Some(0) + } else { + offset + }; + let mut tables = filter_table_infos(tables, filter, limit, final_offset, object_types); if is_oracle { load_oracle_table_comments_for_tables( &mut client, @@ -1685,6 +1710,50 @@ fn filter_table_infos( .collect() } +fn filter_object_infos( + objects: Vec, + filter: Option<&str>, + limit: Option, + offset: Option, + object_types: Option<&[String]>, +) -> Vec { + let filter = filter.unwrap_or(""); + let limit = limit.unwrap_or(usize::MAX); + let offset = offset.unwrap_or(0); + objects + .into_iter() + .filter(|object| crate::sql::contains_or_fuzzy_match(&object.name, filter)) + .filter(|object| object_info_matches_object_types(object, object_types)) + .skip(offset) + .take(limit) + .collect() +} + +fn object_info_matches_object_types(object: &db::ObjectInfo, object_types: Option<&[String]>) -> bool { + let Some(object_types) = object_types else { + return true; + }; + if object_types.is_empty() { + return true; + } + let object_type = normalize_object_info_object_type(&object.object_type); + object_types.iter().any(|expected| normalize_object_info_object_type(expected) == object_type) +} + +fn normalize_object_info_object_type(value: &str) -> String { + let upper = value.to_ascii_uppercase().replace(' ', "_"); + if upper.contains("MATERIALIZED") && upper.contains("VIEW") { + return "MATERIALIZED_VIEW".to_string(); + } + if upper == "BASE_TABLE" || upper.contains("TABLE") { + return "TABLE".to_string(); + } + if upper.contains("VIEW") { + return "VIEW".to_string(); + } + upper +} + fn table_info_matches_object_types(table: &db::TableInfo, object_types: Option<&[String]>) -> bool { let Some(object_types) = object_types else { return true; @@ -1750,8 +1819,12 @@ async fn external_driver_presto_like_objects( config: &ConnectionConfig, database: &str, schema: &str, + filter: Option<&str>, + object_types: Option<&[String]>, ) -> Result, String> { - let tables = external_driver_presto_like_tables(session, config, database, schema, None, None, None).await?; + let tables = external_driver_presto_like_tables(session, config, database, schema, filter, None, None) + .await + .map(|tables| filter_table_infos(tables, filter, None, None, object_types))?; Ok(tables .into_iter() .map(|table| db::ObjectInfo { @@ -1953,10 +2026,10 @@ mod tests { use super::db; use super::{ clickhouse_metadata_database, deduplicate_column_infos, filter_mysql_system_databases_for_config, - filter_table_infos, filter_visible_schema_names, is_agent_postgres_metadata_fallback_config, - is_retryable_metadata_error, 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, + filter_object_infos, filter_table_infos, filter_visible_schema_names, + is_agent_postgres_metadata_fallback_config, is_retryable_metadata_error, 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, @@ -2083,6 +2156,30 @@ mod tests { assert_eq!(visible_schema_filter(Some(&config), "OTHER", true), None); } + #[test] + fn default_oracle_agent_config_excludes_legacy_profiles() { + let mut config = test_connection_config(DatabaseType::Oracle); + assert!(super::is_default_oracle_agent_config(&config)); + + config.driver_profile = Some("oracle".to_string()); + assert!(super::is_default_oracle_agent_config(&config)); + + config.driver_profile = Some("oracle-legacy".to_string()); + assert!(!super::is_default_oracle_agent_config(&config)); + + config.driver_profile = Some("oracle-10g".to_string()); + assert!(!super::is_default_oracle_agent_config(&config)); + } + + #[test] + fn oracle_agent_paging_detection_avoids_double_offset_only_when_page_sized() { + assert!(super::oracle_agent_paging_likely_applied(true, Some(500), 500)); + assert!(super::oracle_agent_paging_likely_applied(true, Some(500), 120)); + assert!(!super::oracle_agent_paging_likely_applied(true, Some(500), 501)); + assert!(!super::oracle_agent_paging_likely_applied(false, Some(500), 120)); + assert!(!super::oracle_agent_paging_likely_applied(true, None, 120)); + } + #[test] fn filter_visible_schema_names_preserves_database_order() { let schemas = vec!["APP".to_string(), "SYS".to_string(), "REPORTING".to_string()]; @@ -2101,6 +2198,20 @@ mod tests { } } + fn test_object_info(name: &str, object_type: &str) -> super::db::ObjectInfo { + super::db::ObjectInfo { + name: name.to_string(), + object_type: object_type.to_string(), + schema: Some("app".to_string()), + signature: None, + comment: None, + created_at: None, + updated_at: None, + parent_schema: None, + parent_name: None, + } + } + fn test_database_info(name: &str) -> super::db::DatabaseInfo { super::db::DatabaseInfo { name: name.to_string() } } @@ -2220,6 +2331,22 @@ mod tests { assert_eq!(filtered[0].name, "active_users"); } + #[test] + fn filter_object_infos_filters_object_type_before_offset_and_limit() { + let objects = vec![ + test_object_info("sync_user", "PROCEDURE"), + test_object_info("find_user", "FUNCTION"), + test_object_info("fetch_name", "FUNCTION"), + test_object_info("orders", "TABLE"), + ]; + let object_types = vec!["FUNCTION".to_string()]; + + let filtered = filter_object_infos(objects, Some("fn"), Some(1), Some(1), Some(&object_types)); + + assert_eq!(filtered.len(), 1); + assert_eq!(filtered[0].name, "fetch_name"); + } + #[test] fn presto_like_information_schema_sql_uses_catalog_and_schema_without_system_jdbc() { let sql = presto_like_information_schema_tables_sql("hive", "sales_analytics", None, None); @@ -2793,9 +2920,26 @@ pub async fn list_objects_core( connection_id: &str, database: &str, schema: &str, + filter: Option<&str>, + limit: Option, + offset: Option, + object_types: Option<&[String]>, ) -> Result, String> { - retry_metadata_connection(state, connection_id, Some(database), || { - list_objects_once(state, connection_id, database, schema) + let use_oracle_agent_paging = + connection_config(state, connection_id).await.as_ref().is_some_and(is_default_oracle_agent_config); + retry_metadata_connection(state, connection_id, Some(database), || async { + let objects = list_objects_once(state, connection_id, database, schema, filter, limit, offset, object_types) + .await + .map(|objects| { + let final_offset = if oracle_agent_paging_likely_applied(use_oracle_agent_paging, limit, objects.len()) + { + Some(0) + } else { + offset + }; + filter_object_infos(objects, filter, limit, final_offset, object_types) + })?; + Ok(objects) }) .await } @@ -3109,6 +3253,10 @@ async fn list_objects_once( connection_id: &str, database: &str, schema: &str, + filter: Option<&str>, + limit: Option, + offset: Option, + object_types: Option<&[String]>, ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; @@ -3144,12 +3292,28 @@ async fn list_objects_once( let session = session.clone(); drop(connections); if uses_presto_like_information_schema_tables(&config.db_type) { - return external_driver_presto_like_objects(session, config.as_ref(), database, schema).await; + return external_driver_presto_like_objects( + session, + config.as_ref(), + database, + schema, + filter, + object_types, + ) + .await; + } + let mut params = + serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }); + if let Some(filter) = filter.map(str::trim).filter(|value| !value.is_empty()) { + params["filter"] = serde_json::json!(filter); + } + if let Some(object_types) = object_types { + params["object_types"] = serde_json::json!(object_types); } return session .invoke_with_timeout::>( "listObjects", - serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }), + params, agent_metadata_timeout(Some(config.as_ref())), ) .await; @@ -3157,18 +3321,41 @@ async fn list_objects_once( try_sqlserver!(connections, &pool_key, list_objects, schema); 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 use_oracle_agent_paging = db_config.as_ref().is_some_and(is_default_oracle_agent_config); + let timeout_duration = agent_metadata_timeout(db_config.as_ref()); let fallback_config = db_config.clone(); drop(connections); - if is_oracle { - return oracle_agent_list_objects(client, database, schema, agent_metadata_timeout(db_config.as_ref())) - .await; + if is_oracle && !use_oracle_agent_paging { + return oracle_agent_list_objects(client, database, schema, timeout_duration).await; } let mut client = client.lock().await; + let agent_limit = if use_oracle_agent_paging { limit } else { None }; + let agent_offset = if use_oracle_agent_paging { offset } else { None }; match client - .list_objects::>(database, schema, agent_metadata_timeout(db_config.as_ref())) + .list_objects_constrained::>( + database, + schema, + filter, + agent_limit, + agent_offset, + object_types, + timeout_duration, + ) .await { - Ok(objects) if !objects.is_empty() => return Ok(objects), + Ok(mut objects) if !objects.is_empty() => { + if is_oracle { + load_oracle_table_comments_for_objects( + &mut client, + database, + schema, + &mut objects, + timeout_duration, + ) + .await?; + } + return Ok(objects); + } Ok(objects) => { if let Some(config) = fallback_config.as_ref() { match native_postgres_metadata_pool(state, connection_id, database, config).await { @@ -3338,7 +3525,7 @@ async fn list_completion_objects_once( PoolKind::Postgres(p) => db::postgres::list_objects(p, schema).await.map(filter_completion_objects), PoolKind::SqlServer(_) => { drop(connections); - let objects = list_objects_once(state, connection_id, database, schema).await?; + let objects = list_objects_once(state, connection_id, database, schema, None, None, None, None).await?; Ok(filter_completion_objects(objects)) } _ => Ok(Vec::new()), @@ -4070,6 +4257,16 @@ fn is_opengauss_family_config(config: &ConnectionConfig) -> bool { || matches!(config.driver_profile.as_deref(), Some("opengauss" | "gaussdb")) } +fn is_default_oracle_agent_config(config: &ConnectionConfig) -> bool { + // Only the default go-oracle agent handles filtered/paged metadata; legacy profiles keep Rust fallback paging. + matches!(config.db_type, DatabaseType::Oracle) + && !matches!(config.driver_profile.as_deref(), Some("oracle-legacy" | "oracle-10g")) +} + +fn oracle_agent_paging_likely_applied(enabled: bool, limit: Option, returned_len: usize) -> bool { + enabled && limit.is_some_and(|limit| returned_len <= limit) +} + fn is_doris_family_config(config: &ConnectionConfig) -> bool { matches!(config.db_type, DatabaseType::Doris | DatabaseType::StarRocks | DatabaseType::ManticoreSearch) || matches!(config.driver_profile.as_deref(), Some("doris" | "selectdb" | "starrocks" | "manticoresearch")) diff --git a/packages/app-tests/metadataLoadCoordinator.test.ts b/packages/app-tests/metadataLoadCoordinator.test.ts new file mode 100644 index 0000000000..06aa127351 --- /dev/null +++ b/packages/app-tests/metadataLoadCoordinator.test.ts @@ -0,0 +1,169 @@ +import assert from "node:assert/strict"; +import { test } from "vitest"; +import { createMetadataLoadTrace, MetadataLoadCoordinator } from "../../apps/desktop/src/lib/metadata/metadataLoadCoordinator.ts"; +import { metadataScopeKey } from "../../apps/desktop/src/lib/metadata/metadataLoadScope.ts"; + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +} + +test("metadata scope key treats equivalent object type sets as equal", () => { + const left = metadataScopeKey({ + kind: "object-group", + connectionId: "conn", + database: "app", + schema: "public", + nodeKind: "group-tables", + objectTypes: ["view", "TABLE", "TABLE"], + searchFilter: " user ", + limit: 501, + offset: 0, + sidebarDisplayMode: "grouped", + driverProfile: "postgres", + }); + const right = metadataScopeKey({ + kind: "object-group", + connectionId: "conn", + database: "app", + schema: "public", + nodeKind: "group-tables", + objectTypes: ["TABLE", "VIEW"], + searchFilter: "user", + limit: 501, + offset: 0, + sidebarDisplayMode: "grouped", + driverProfile: "postgres", + }); + + assert.equal(left, right); +}); + +test("metadata scope key separates result-affecting fields", () => { + const base = { + kind: "sidebar-table-search", + connectionId: "conn", + database: "app", + schema: "public", + objectTypes: ["TABLE"], + limit: 500, + offset: 0, + sidebarDisplayMode: "simple", + driverProfile: "postgres", + }; + + assert.notEqual(metadataScopeKey({ ...base, searchFilter: "user" }), metadataScopeKey({ ...base, searchFilter: "order" })); + assert.notEqual(metadataScopeKey({ ...base, offset: 0 }), metadataScopeKey({ ...base, offset: 500 })); + assert.notEqual(metadataScopeKey({ ...base, objectTypes: ["TABLE"] }), metadataScopeKey({ ...base, objectTypes: ["VIEW"] })); + assert.notEqual(metadataScopeKey({ ...base, sidebarDisplayMode: "simple" }), metadataScopeKey({ ...base, sidebarDisplayMode: "grouped" })); +}); + +test("metadata load coordinator deduplicates equivalent in-flight loads", async () => { + const events: string[] = []; + const coordinator = new MetadataLoadCoordinator((event) => events.push(event.event)); + const gate = deferred(); + let calls = 0; + const scope = { kind: "database-schemas", connectionId: "conn", database: "app", driverProfile: "postgres" }; + const first = coordinator.run(scope, async () => { + calls++; + return gate.promise; + }); + const second = coordinator.run(scope, async () => { + calls++; + return "unexpected"; + }); + + assert.equal(calls, 0); + await Promise.resolve(); + assert.equal(calls, 1); + assert.equal(first, second); + assert.deepEqual(events.slice(0, 2), ["start", "dedupe"]); + + gate.resolve("schemas"); + assert.equal(await first, "schemas"); + assert.equal(await second, "schemas"); + assert.equal(coordinator.has(scope), false); +}); + +test("metadata load coordinator keeps different filters independent", async () => { + const coordinator = new MetadataLoadCoordinator(); + const first = deferred(); + const second = deferred(); + let calls = 0; + + const user = coordinator.run({ kind: "sidebar-table-search", connectionId: "conn", database: "app", schema: "public", searchFilter: "user" }, async () => { + calls++; + return first.promise; + }); + const order = coordinator.run({ kind: "sidebar-table-search", connectionId: "conn", database: "app", schema: "public", searchFilter: "order" }, async () => { + calls++; + return second.promise; + }); + + await Promise.resolve(); + assert.equal(calls, 2); + first.resolve("users"); + second.resolve("orders"); + assert.deepEqual(await Promise.all([user, order]), ["users", "orders"]); +}); + +test("metadata load coordinator force bypasses in-flight reuse", async () => { + const coordinator = new MetadataLoadCoordinator(); + const slow = deferred(); + let calls = 0; + const scope = { kind: "connection-databases", connectionId: "conn" }; + const normal = coordinator.run(scope, async () => { + calls++; + return slow.promise; + }); + const forced = coordinator.run( + scope, + async () => { + calls++; + return "fresh"; + }, + { force: true }, + ); + + assert.equal(await forced, "fresh"); + assert.equal(calls, 2); + slow.resolve("cached"); + assert.equal(await normal, "cached"); +}); + +test("metadata load trace reports scope, elapsed time, cache and request state", () => { + let now = 100; + const trace = createMetadataLoadTrace( + { + kind: "table-metadata", + connectionId: "conn", + database: "app", + schema: "public", + tableName: "users", + }, + { traceId: "trace-1", now: () => now }, + ); + + now = 137; + assert.deepEqual(trace.event("done", { cacheStatus: "stale", resultCount: 12, deduped: true, superseded: false }), { + event: "done", + traceId: "trace-1", + scopeKey: metadataScopeKey({ + kind: "table-metadata", + connectionId: "conn", + database: "app", + schema: "public", + tableName: "users", + }), + elapsedMs: 37, + cacheStatus: "stale", + resultCount: 12, + deduped: true, + superseded: false, + }); +}); diff --git a/packages/app-tests/metadataResultCache.test.ts b/packages/app-tests/metadataResultCache.test.ts new file mode 100644 index 0000000000..a478cfc6e0 --- /dev/null +++ b/packages/app-tests/metadataResultCache.test.ts @@ -0,0 +1,54 @@ +import assert from "node:assert/strict"; +import { test } from "vitest"; +import { MetadataResultCache } from "../../apps/desktop/src/lib/metadata/metadataResultCache.ts"; + +test("metadata result cache returns fresh hits and exposes stale hits when requested", () => { + let now = 1_000; + const cache = new MetadataResultCache({ ttlMs: 50, maxEntries: 10, now: () => now }); + const scope = { kind: "table-list-page", connectionId: "conn", database: "app", schema: "public", limit: 501, offset: 0 }; + + cache.set(scope, "tables"); + assert.equal(cache.get(scope)?.value, "tables"); + + now = 1_100; + assert.equal(cache.get(scope), undefined); + const stale = cache.get(scope, { allowStale: true }); + assert.equal(stale?.value, "tables"); + assert.equal(stale?.stale, true); + assert.equal(stale?.ageMs, 100); +}); + +test("metadata result cache invalidates affected scopes without clearing unrelated databases", () => { + const cache = new MetadataResultCache({ ttlMs: 30_000, maxEntries: 10, now: () => 1 }); + const publicTables = { kind: "table-list-page", connectionId: "conn", database: "app", schema: "public", objectTypes: ["VIEW", "TABLE"], limit: 501, offset: 0 }; + const auditTables = { kind: "table-list-page", connectionId: "conn", database: "app", schema: "audit", objectTypes: ["TABLE"], limit: 501, offset: 0 }; + const otherDatabase = { kind: "table-list-page", connectionId: "conn", database: "other", schema: "public", objectTypes: ["TABLE"], limit: 501, offset: 0 }; + const otherConnection = { kind: "table-list-page", connectionId: "other", database: "app", schema: "public", objectTypes: ["TABLE"], limit: 501, offset: 0 }; + + cache.set(publicTables, "public"); + cache.set(auditTables, "audit"); + cache.set(otherDatabase, "other-db"); + cache.set(otherConnection, "other-conn"); + + assert.equal(cache.invalidate({ connectionId: "conn", database: "app", schema: "public" }), 1); + assert.equal(cache.get(publicTables), undefined); + assert.equal(cache.get(auditTables)?.value, "audit"); + assert.equal(cache.get(otherDatabase)?.value, "other-db"); + assert.equal(cache.get(otherConnection)?.value, "other-conn"); +}); + +test("metadata result cache evicts least recently used entries", () => { + const cache = new MetadataResultCache({ ttlMs: 30_000, maxEntries: 2, now: () => 1 }); + const first = { kind: "object-list-page", connectionId: "conn", database: "app", schema: "public", nodeKind: "group-functions" }; + const second = { kind: "object-list-page", connectionId: "conn", database: "app", schema: "audit", nodeKind: "group-functions" }; + const third = { kind: "object-list-page", connectionId: "conn", database: "other", schema: "public", nodeKind: "group-functions" }; + + cache.set(first, "first"); + cache.set(second, "second"); + assert.equal(cache.get(first)?.value, "first"); + cache.set(third, "third"); + + assert.equal(cache.get(first)?.value, "first"); + assert.equal(cache.get(second), undefined); + assert.equal(cache.get(third)?.value, "third"); +}); diff --git a/packages/app-tests/sidebarTableSearchControl.test.ts b/packages/app-tests/sidebarTableSearchControl.test.ts index 7adbb3df6f..c4059ef9b7 100644 --- a/packages/app-tests/sidebarTableSearchControl.test.ts +++ b/packages/app-tests/sidebarTableSearchControl.test.ts @@ -71,7 +71,7 @@ test("inserts a grouped local table search control only for expanded table group connectionId: "conn", database: "app", isExpanded: true, - children: [], + children: [{ id: "conn:app:orders", label: "orders", type: "table", connectionId: "conn", database: "app" }], }; const viewGroup: TreeNode = { id: "conn:app:__views", @@ -96,6 +96,65 @@ test("inserts a grouped local table search control only for expanded table group assert.equal(nodes[1].node.tableSearchParentId, "conn:app:__tables"); }); +test("does not insert a local table search control for empty table scopes", () => { + const tableGroup: TreeNode = { + id: "conn:app:__tables", + label: "tree.tables", + type: "group-tables", + connectionId: "conn", + database: "app", + isExpanded: true, + children: [], + }; + const routineOnlySchema: TreeNode = { + id: "conn:app:public", + label: "public", + type: "schema", + connectionId: "conn", + database: "app", + schema: "public", + isExpanded: true, + children: [{ id: "conn:app:public:sync_orders", label: "sync_orders", type: "procedure", connectionId: "conn", database: "app", schema: "public" }], + }; + + const groupedNodes = insertSidebarTableSearchControls([flat(tableGroup, 1)], { + enabled: true, + sidebarObjectDisplay: "grouped", + activeQueries: {}, + }); + const simpleNodes = insertSidebarTableSearchControls([flat(routineOnlySchema, 1)], { + enabled: true, + sidebarObjectDisplay: "simple", + activeQueries: {}, + }); + + assert.deepEqual(groupedNodes.map((item) => item.node.type), ["group-tables"]); + assert.deepEqual(simpleNodes.map((item) => item.node.type), ["schema"]); +}); + +test("keeps a grouped local table search control visible when the current search has no results", () => { + const tableGroup: TreeNode = { + id: "conn:app:__tables", + label: "tree.tables", + type: "group-tables", + connectionId: "conn", + database: "app", + isExpanded: true, + children: [], + }; + + const nodes = insertSidebarTableSearchControls([flat(tableGroup, 1)], { + enabled: true, + sidebarObjectDisplay: "grouped", + activeQueries: { "conn:app:__tables": "invoice" }, + }); + + assert.deepEqual( + nodes.map((item) => item.node.type), + ["group-tables", "table-search-control"], + ); +}); + test("uses isolated virtual scroller pools for each local table search input", () => { const first: TreeNode = { id: "conn:first:__tables", @@ -104,7 +163,7 @@ test("uses isolated virtual scroller pools for each local table search input", ( connectionId: "conn", database: "first", isExpanded: true, - children: [], + children: [{ id: "conn:first:orders", label: "orders", type: "table", connectionId: "conn", database: "first" }], }; const second: TreeNode = { id: "conn:second:__tables", @@ -113,7 +172,7 @@ test("uses isolated virtual scroller pools for each local table search input", ( connectionId: "conn", database: "second", isExpanded: true, - children: [], + children: [{ id: "conn:second:orders", label: "orders", type: "table", connectionId: "conn", database: "second" }], }; const nodes = insertSidebarTableSearchControls([flat(first, 1), flat(second, 1)], { diff --git a/packages/app-tests/tableMetadataCache.test.ts b/packages/app-tests/tableMetadataCache.test.ts new file mode 100644 index 0000000000..d555f598e0 --- /dev/null +++ b/packages/app-tests/tableMetadataCache.test.ts @@ -0,0 +1,101 @@ +import assert from "node:assert/strict"; +import { beforeEach, test, vi } from "vitest"; +import type { ColumnInfo, IndexInfo } from "../../apps/desktop/src/types/database.ts"; + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +} + +beforeEach(() => { + vi.resetModules(); +}); + +test("table metadata loader deduplicates equivalent in-flight requests and caches the result", async () => { + const columnsGate = deferred(); + const indexes: IndexInfo[] = [{ name: "users_pkey", columns: ["id"], is_unique: true, is_primary: true }]; + const getColumns = vi.fn(() => columnsGate.promise); + const listIndexes = vi.fn().mockResolvedValue(indexes); + vi.doMock("@/lib/backend/api", () => ({ getColumns, listIndexes })); + + const { clearTableMetadataCache, loadTableMetadata } = await import("../../apps/desktop/src/lib/metadata/tableMetadataCache.ts"); + clearTableMetadataCache(); + + const request = { + connectionId: "conn", + database: "app", + schema: "public", + tableName: "users", + tableType: "TABLE", + databaseType: "postgres", + driverProfile: "postgres", + }; + const first = loadTableMetadata(request); + const second = loadTableMetadata(request); + await Promise.resolve(); + assert.equal(getColumns.mock.calls.length, 1); + + columnsGate.resolve([ + { + name: "id", + data_type: "integer", + is_nullable: false, + column_default: null, + is_primary_key: true, + extra: null, + }, + ]); + + const [firstResult, secondResult] = await Promise.all([first, second]); + assert.deepEqual(firstResult.metadata.primaryKeys, ["id"]); + assert.equal(secondResult.metadata.columns.length, 1); + assert.equal(listIndexes.mock.calls.length, 1); + + const cached = await loadTableMetadata(request); + assert.equal(cached.cacheStatus, "hit"); + assert.equal(getColumns.mock.calls.length, 1); +}); + +test("table metadata invalidation keeps unrelated schemas cached", async () => { + const getColumns = vi.fn( + async (_connectionId: string, _database: string, _schema: string, table: string): Promise => [ + { + name: `${table}_id`, + data_type: "integer", + is_nullable: false, + column_default: null, + is_primary_key: true, + extra: null, + }, + ], + ); + const listIndexes = vi.fn(async (): Promise => []); + vi.doMock("@/lib/backend/api", () => ({ getColumns, listIndexes })); + + const { clearTableMetadataCache, getCachedTableMetadata, invalidateTableMetadataCache, loadTableMetadata } = await import("../../apps/desktop/src/lib/metadata/tableMetadataCache.ts"); + clearTableMetadataCache(); + + const publicRequest = { + connectionId: "conn", + database: "app", + schema: "public", + tableName: "users", + tableType: "TABLE", + databaseType: "postgres", + driverProfile: "postgres", + }; + const auditRequest = { ...publicRequest, schema: "audit", tableName: "audit_log" }; + + await loadTableMetadata(publicRequest); + await loadTableMetadata(auditRequest); + assert.equal(getColumns.mock.calls.length, 2); + + assert.equal(invalidateTableMetadataCache({ connectionId: "conn", database: "app", schema: "public" }), 1); + assert.equal(getCachedTableMetadata(publicRequest), undefined); + assert.equal(getCachedTableMetadata(auditRequest)?.metadata.tableName, "audit_log"); +}); diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index fc13e099c3..53316816b8 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -263,12 +263,18 @@ public final class DbxJdbcPlugin { optionalText(params, "database"), optionalText(params, "schema"), optionalText(params, "filter"), - nonNegativeInt(params, "limit", 0) + nonNegativeInt(params, "limit", 0), + nonNegativeInt(params, "offset", 0), + optionalStringList(params, "object_types") ); case "listObjects", "list_objects" -> listObjects( connection, optionalText(params, "database"), - optionalText(params, "schema") + optionalText(params, "schema"), + optionalText(params, "filter"), + nonNegativeInt(params, "limit", 0), + nonNegativeInt(params, "offset", 0), + optionalStringList(params, "object_types") ); case "listDataTypes", "list_data_types" -> listDataTypes(connection, optionalText(params, "database")); case "getObjectSource", "get_object_source" -> getObjectSource( @@ -1121,45 +1127,82 @@ public final class DbxJdbcPlugin { return result; } - private static JsonNode listTables(JsonNode connection, String database, String schema, String filter, int limit) throws SQLException { + private static JsonNode listTables( + JsonNode connection, + String database, + String schema, + String filter, + int limit, + int offset, + List objectTypes + ) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); Connection conn = openConnection(connection); JdbcDriverQuirks quirks = driverQuirks(connection); if (quirks.useOracleMetadata()) { - return oracleListTables(conn, oracleEffectiveSchema(conn, schema)); + return filterMetadataNodes( + (ArrayNode) oracleListTables(conn, oracleEffectiveSchema(conn, schema)), + filter, + limit, + offset, + objectTypes, + "table_type", + true + ); } if (usePrestoInformationSchemaTables(connection)) { - return prestoListTables(conn, database, schema, filter, limit); + return prestoListTables(conn, database, schema, filter, limit, offset, objectTypes); } DatabaseMetaData meta = conn.getMetaData(); - String[] types = jdbcTableTypes(meta); + String[] types = constrainedJdbcTableTypes(jdbcTableTypes(meta), objectTypes); + if (types.length == 0) { + return result; + } String catalog = metadataCatalog(database, quirks); String schemaPattern = resolveSchemaPattern(meta, database, schema, quirks); appendTables(result, meta, catalog, schemaPattern, types); if (result.isEmpty() && catalog != null) { appendTables(result, meta, null, schemaPattern, types); } - return result; + return filterMetadataNodes(result, filter, limit, offset, objectTypes, "table_type", true); } - private static JsonNode listObjects(JsonNode connection, String database, String schema) throws SQLException { + private static JsonNode listObjects( + JsonNode connection, + String database, + String schema, + String filter, + int limit, + int offset, + List objectTypes + ) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); Connection conn = openConnection(connection); if (driverQuirks(connection).useOracleMetadata()) { - return oracleListObjects(conn, oracleEffectiveSchema(conn, schema), schema); + return filterMetadataNodes( + (ArrayNode) oracleListObjects(conn, oracleEffectiveSchema(conn, schema), schema), + filter, + limit, + offset, + objectTypes, + "object_type", + false + ); } if (usePrestoInformationSchemaTables(connection)) { - return prestoListObjects(conn, database, schema); + return prestoListObjects(conn, database, schema, filter, limit, offset, objectTypes); } DatabaseMetaData meta = conn.getMetaData(); JdbcDriverQuirks quirks = driverQuirks(connection); String catalog = metadataCatalog(database, quirks); String schemaPattern = resolveSchemaPattern(meta, database, schema, quirks); - String[] tableTypes = jdbcTableTypes(meta); - appendTableObjects(result, meta, catalog, schemaPattern, schema, tableTypes); - if (result.isEmpty() && catalog != null) { - appendTableObjects(result, meta, null, schemaPattern, schema, tableTypes); + String[] tableTypes = constrainedJdbcTableTypes(jdbcTableTypes(meta), objectTypes); + if (tableTypes.length > 0) { + appendTableObjects(result, meta, catalog, schemaPattern, schema, tableTypes); + if (result.isEmpty() && catalog != null) { + appendTableObjects(result, meta, null, schemaPattern, schema, tableTypes); + } } try (ResultSet rs = meta.getProcedures(catalog, schemaPattern, "%")) { @@ -1195,7 +1238,7 @@ public final class DbxJdbcPlugin { } catch (SQLException ignored) { } - return result; + return filterMetadataNodes(result, filter, limit, offset, objectTypes, "object_type", false); } private static JsonNode listDataTypes(JsonNode connection, String database) throws SQLException { @@ -1378,6 +1421,111 @@ public final class DbxJdbcPlugin { return DEFAULT_TABLE_TYPES; } + private static String[] constrainedJdbcTableTypes(String[] tableTypes, List objectTypes) { + Set allowed = normalizedObjectTypes(objectTypes); + if (allowed.isEmpty()) { + return tableTypes; + } + List result = new ArrayList<>(); + for (String tableType : tableTypes) { + if (allowed.contains(normalizeTableObjectType(tableType))) { + result.add(tableType); + } + } + return result.toArray(new String[0]); + } + + private static ArrayNode filterMetadataNodes( + ArrayNode source, + String filter, + int limit, + int offset, + List objectTypes, + String typeField, + boolean defaultBlankTypeToTable + ) { + ArrayNode result = MAPPER.createArrayNode(); + Set allowedTypes = normalizedObjectTypes(objectTypes); + String normalizedFilter = filter == null ? "" : filter.trim().toLowerCase(Locale.ROOT); + int start = Math.max(0, offset); + int max = limit <= 0 ? Integer.MAX_VALUE : limit; + int skipped = 0; + for (JsonNode item : source) { + if (!metadataNameMatches(item.path("name").asText(""), normalizedFilter)) { + continue; + } + String type = item.path(typeField).asText(""); + String normalizedType = defaultBlankTypeToTable ? normalizeTableObjectType(type) : normalizeObjectType(type); + if (!allowedTypes.isEmpty() && (normalizedType.isEmpty() || !allowedTypes.contains(normalizedType))) { + continue; + } + if (skipped++ < start) { + continue; + } + if (result.size() >= max) { + break; + } + result.add(item); + } + return result; + } + + private static boolean metadataNameMatches(String name, String filter) { + if (filter == null || filter.isEmpty()) { + return true; + } + String candidate = name == null ? "" : name.toLowerCase(Locale.ROOT); + return candidate.contains(filter) || (filter.length() >= 2 && fuzzySubsequenceMatches(candidate, filter)); + } + + private static boolean fuzzySubsequenceMatches(String candidate, String expected) { + int cursor = 0; + for (int i = 0; i < expected.length(); i++) { + cursor = candidate.indexOf(expected.charAt(i), cursor); + if (cursor < 0) { + return false; + } + cursor++; + } + return true; + } + + private static Set normalizedObjectTypes(List objectTypes) { + Set result = new HashSet<>(); + if (objectTypes == null) { + return result; + } + for (String objectType : objectTypes) { + String normalized = normalizeObjectType(objectType); + if (!normalized.isEmpty()) { + result.add(normalized); + } + } + return result; + } + + private static String normalizeTableObjectType(String value) { + String normalized = normalizeObjectType(value); + return normalized.isEmpty() ? "TABLE" : normalized; + } + + private static String normalizeObjectType(String value) { + if (value == null || value.isBlank()) { + return ""; + } + String upper = value.trim().toUpperCase(Locale.ROOT).replace(' ', '_'); + if (upper.contains("MATERIALIZED") && upper.contains("VIEW")) { + return "MATERIALIZED_VIEW"; + } + if ("BASE_TABLE".equals(upper) || upper.contains("TABLE")) { + return "TABLE"; + } + if (upper.contains("VIEW")) { + return "VIEW"; + } + return upper; + } + private static void appendTableObjects( ArrayNode result, DatabaseMetaData meta, @@ -1403,9 +1551,18 @@ public final class DbxJdbcPlugin { return urlMatchesPrefix(url, "jdbc:presto:") || urlMatchesPrefix(url, "jdbc:trino:"); } - private static JsonNode prestoListTables(Connection conn, String database, String schema, String filter, int limit) throws SQLException { + private static JsonNode prestoListTables( + Connection conn, + String database, + String schema, + String filter, + int limit, + int offset, + List objectTypes + ) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); - try (PreparedStatement ps = conn.prepareStatement(prestoInformationSchemaTablesSql(database, filter, limit))) { + int queryLimit = limit > 0 ? Math.max(1, limit + Math.max(0, offset)) : 0; + try (PreparedStatement ps = conn.prepareStatement(prestoInformationSchemaTablesSql(database, filter, queryLimit))) { ps.setString(1, schema); if (emptyToNull(filter) != null) { ps.setString(2, escapeLikePattern(filter.toLowerCase(Locale.ROOT)) + "%"); @@ -1420,13 +1577,25 @@ public final class DbxJdbcPlugin { } } } - return result; + return filterMetadataNodes(result, filter, limit, offset, objectTypes, "table_type", true); } - private static JsonNode prestoListObjects(Connection conn, String database, String schema) throws SQLException { + private static JsonNode prestoListObjects( + Connection conn, + String database, + String schema, + String filter, + int limit, + int offset, + List objectTypes + ) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); - try (PreparedStatement ps = conn.prepareStatement(prestoInformationSchemaTablesSql(database, null, 0))) { + int queryLimit = limit > 0 ? Math.max(1, limit + Math.max(0, offset)) : 0; + try (PreparedStatement ps = conn.prepareStatement(prestoInformationSchemaTablesSql(database, filter, queryLimit))) { ps.setString(1, schema); + if (emptyToNull(filter) != null) { + ps.setString(2, escapeLikePattern(filter.toLowerCase(Locale.ROOT)) + "%"); + } try (ResultSet rs = ps.executeQuery()) { while (rs.next()) { ObjectNode item = MAPPER.createObjectNode(); @@ -1438,7 +1607,7 @@ public final class DbxJdbcPlugin { } } } - return result; + return filterMetadataNodes(result, filter, limit, offset, objectTypes, "object_type", false); } private static JsonNode prestoGetColumns(Connection conn, String database, String schema, String table) throws SQLException { @@ -2310,6 +2479,34 @@ public final class DbxJdbcPlugin { return text.isEmpty() ? null : text; } + private static List optionalStringList(JsonNode node, String field) { + JsonNode value = node.path(field); + if (value.isMissingNode() || value.isNull()) { + return null; + } + List result = new ArrayList<>(); + if (value.isArray()) { + for (JsonNode item : value) { + String text = item.asText("").trim(); + if (!text.isEmpty()) { + result.add(text); + } + } + return result; + } + String text = value.asText("").trim(); + if (text.isEmpty()) { + return null; + } + for (String part : text.split(",")) { + String item = part.trim(); + if (!item.isEmpty()) { + result.add(item); + } + } + return result; + } + private static int positiveInt(JsonNode node, String field, int defaultValue) { return Math.max(1, nonNegativeInt(node, field, defaultValue)); } diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java index a53764709e..f08284d39e 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -731,6 +731,43 @@ final class DbxJdbcPluginTest { assertEquals("PEOPLE", response.path("result").path(0).path("name").asText()); } + @Test + void listTablesAppliesMetadataConstraints() throws Exception { + request("executeQuery", """ + { + "connection": %s, + "sql": "CREATE SCHEMA IF NOT EXISTS app" + } + """.formatted(CONNECTION)); + request("executeQuery", """ + { + "connection": %s, + "sql": "CREATE TABLE IF NOT EXISTS app.people (id INT PRIMARY KEY)" + } + """.formatted(CONNECTION)); + request("executeQuery", """ + { + "connection": %s, + "sql": "CREATE TABLE IF NOT EXISTS app.people_archive (id INT PRIMARY KEY)" + } + """.formatted(CONNECTION)); + + JsonNode response = request("listTables", """ + { + "connection": %s, + "schema": "APP", + "filter": "people", + "limit": 1, + "offset": 1, + "object_types": ["TABLE"] + } + """.formatted(CONNECTION)); + + assertFalse(response.has("error"), response.toString()); + assertEquals(1, response.path("result").size()); + assertEquals("PEOPLE_ARCHIVE", response.path("result").path(0).path("name").asText()); + } + @Test void listDatabasesIncludesConfiguredDatabaseWhenDriverDoesNotReturnIt() throws Exception { String connection = """ @@ -794,6 +831,43 @@ final class DbxJdbcPluginTest { assertEquals("PEOPLE", response.path("result").path(0).path("name").asText()); } + @Test + void listObjectsAppliesMetadataConstraints() throws Exception { + request("executeQuery", """ + { + "connection": %s, + "sql": "CREATE SCHEMA IF NOT EXISTS app" + } + """.formatted(CONNECTION)); + request("executeQuery", """ + { + "connection": %s, + "sql": "CREATE TABLE IF NOT EXISTS app.people (id INT PRIMARY KEY)" + } + """.formatted(CONNECTION)); + request("executeQuery", """ + { + "connection": %s, + "sql": "CREATE TABLE IF NOT EXISTS app.people_archive (id INT PRIMARY KEY)" + } + """.formatted(CONNECTION)); + + JsonNode response = request("listObjects", """ + { + "connection": %s, + "schema": "APP", + "filter": "people", + "limit": 1, + "offset": 1, + "object_types": ["TABLE"] + } + """.formatted(CONNECTION)); + + assertFalse(response.has("error"), response.toString()); + assertEquals(1, response.path("result").size()); + assertEquals("PEOPLE_ARCHIVE", response.path("result").path(0).path("name").asText()); + } + @Test void getColumnsFallsBackWhenCatalogFiltersEverything() throws Exception { createPeopleTable(); diff --git a/src-tauri/src/commands/schema.rs b/src-tauri/src/commands/schema.rs index cdb7b372df..907b03bb29 100644 --- a/src-tauri/src/commands/schema.rs +++ b/src-tauri/src/commands/schema.rs @@ -138,8 +138,22 @@ pub async fn list_objects( connection_id: String, database: String, schema: String, + filter: Option, + limit: Option, + offset: Option, + object_types: Option>, ) -> Result, String> { - dbx_core::schema::list_objects_core(&state, &connection_id, &database, &schema).await + dbx_core::schema::list_objects_core( + &state, + &connection_id, + &database, + &schema, + filter.as_deref(), + limit, + offset, + object_types.as_deref(), + ) + .await } #[tauri::command]