diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index d40a83b8ab..9c45fd6dfa 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -6800,6 +6800,21 @@ async fn list_indexes_core_for_session( drop(connections); return external_driver_gaussdb_m_indexes(session, config.as_ref(), database, schema, table).await; } + let config = config.clone(); + let session = session.clone(); + drop(connections); + return session + .invoke_with_timeout::>( + "listIndexes", + serde_json::json!({ + "connection": config.as_ref(), + "database": database, + "schema": schema, + "table": table, + }), + agent_metadata_timeout(Some(config.as_ref())), + ) + .await; } if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { drop(connections); 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 5b5dd9c6e6..a1197036b1 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -44,12 +44,14 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.Properties; import java.util.ServiceLoader; import java.util.Set; +import java.util.TreeMap; import java.util.UUID; import java.util.logging.Logger; import java.util.regex.Pattern; @@ -406,6 +408,12 @@ public final class DbxJdbcPlugin { nonNegativeInt(params, "offset", 0), optionalStringList(params, "object_types") ); + case "listIndexes", "list_indexes" -> listIndexes( + connection, + optionalText(params, "database"), + optionalText(params, "schema"), + requireText(params, "table") + ); case "listDataTypes", "list_data_types" -> listDataTypes(connection, optionalText(params, "database")); case "getObjectSource", "get_object_source" -> getObjectSource( connection, @@ -2086,6 +2094,131 @@ public final class DbxJdbcPlugin { return result; } + private static JsonNode listIndexes(JsonNode connection, String database, String schema, String table) + throws SQLException { + Connection conn = openConnection(connection); + JdbcDriverQuirks quirks = driverQuirks(connection); + if (quirks.useOracleMetadata()) { + String owner = oracleEffectiveSchema(conn, schema); + String resolvedTable = oracleResolveTable(conn, owner, table); + return oracleListIndexes(conn, owner, resolvedTable == null ? table : resolvedTable); + } + + DatabaseMetaData meta = conn.getMetaData(); + String catalog = metadataCatalog(database, quirks); + String schemaPattern = resolveSchemaPattern(meta, database, schema, quirks); + Set primaryIndexNames = new HashSet<>(); + Map primaryColumnsBySequence = new TreeMap<>(); + try (ResultSet rs = meta.getPrimaryKeys(catalog, schemaPattern, table)) { + while (rs != null && rs.next()) { + String name = rs.getString("PK_NAME"); + if (name != null && !name.isBlank()) { + primaryIndexNames.add(name); + } + String column = rs.getString("COLUMN_NAME"); + if (column != null && !column.isBlank()) { + primaryColumnsBySequence.put((int) rs.getShort("KEY_SEQ"), column); + } + } + } catch (SQLException | AbstractMethodError | UnsupportedOperationException ignored) { + } + + // Presto/Trino JDBC throws SQLFeatureNotSupportedException from getIndexInfo, and + // drivers compiled before JDBC 4 surface unimplemented DatabaseMetaData methods as + // AbstractMethodError. Both must degrade to an empty list, matching the metadata + // error tolerance used by the other DatabaseMetaData readers in this plugin. + LinkedHashMap indexes = new LinkedHashMap<>(); + try (ResultSet rs = meta.getIndexInfo(catalog, schemaPattern, table, false, false)) { + appendJdbcIndexes(indexes, primaryIndexNames, rs); + } catch (SQLException | AbstractMethodError | UnsupportedOperationException ignored) { + } + if (indexes.isEmpty() && catalog != null) { + try (ResultSet rs = meta.getIndexInfo(null, schemaPattern, table, false, false)) { + appendJdbcIndexes(indexes, primaryIndexNames, rs); + } catch (SQLException | AbstractMethodError | UnsupportedOperationException ignored) { + } + } + markPrimaryIndexByColumns(indexes.values(), new ArrayList<>(primaryColumnsBySequence.values())); + ArrayNode result = MAPPER.createArrayNode(); + indexes.values().forEach(result::add); + return result; + } + + private static void markPrimaryIndexByColumns(Iterable indexes, List primaryColumns) { + if (primaryColumns.isEmpty()) { + return; + } + for (ObjectNode index : indexes) { + if (index.path("is_primary").asBoolean()) { + return; + } + } + for (ObjectNode index : indexes) { + JsonNode columns = index.path("columns"); + if (columns.size() != primaryColumns.size()) { + continue; + } + boolean matches = true; + for (int i = 0; i < primaryColumns.size(); i++) { + if (!primaryColumns.get(i).equalsIgnoreCase(columns.path(i).asText())) { + matches = false; + break; + } + } + if (matches) { + index.put("is_primary", true); + return; + } + } + } + + private static void appendJdbcIndexes( + Map indexes, + Set primaryIndexNames, + ResultSet rs + ) throws SQLException { + while (rs != null && rs.next()) { + String name = rs.getString("INDEX_NAME"); + String column = rs.getString("COLUMN_NAME"); + if (name == null || name.isBlank() || column == null || column.isBlank()) { + continue; + } + ObjectNode item = indexes.get(name); + if (item == null) { + item = indexNode( + name, + !rs.getBoolean("NON_UNIQUE"), + primaryIndexNames.contains(name), + jdbcIndexType(rs.getShort("TYPE")) + ); + indexes.put(name, item); + } + ((ArrayNode) item.path("columns")).add(column); + } + } + + private static String jdbcIndexType(short type) { + return switch (type) { + case DatabaseMetaData.tableIndexClustered -> "CLUSTERED"; + case DatabaseMetaData.tableIndexHashed -> "HASHED"; + case DatabaseMetaData.tableIndexOther -> "OTHER"; + default -> null; + }; + } + + private static ObjectNode indexNode(String name, boolean unique, boolean primary, String indexType) { + ObjectNode item = MAPPER.createObjectNode(); + item.put("name", name); + item.set("columns", MAPPER.createArrayNode()); + item.put("is_unique", unique); + item.put("is_primary", primary); + item.putNull("filter"); + putNullable(item, "index_type", indexType); + item.putNull("included_columns"); + item.putNull("comment"); + return item; + } + private static JsonNode getColumns(JsonNode connection, String database, String schema, String table) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); Connection conn = openConnection(connection); @@ -3361,6 +3494,74 @@ public final class DbxJdbcPlugin { } } } + String sequenceSql = "SELECT sequence_name AS name FROM all_sequences WHERE sequence_owner = ? ORDER BY sequence_name"; + try (PreparedStatement ps = conn.prepareStatement(sequenceSql)) { + ps.setString(1, owner); + try (ResultSet rs = ps.executeQuery()) { + while (rs.next()) { + ObjectNode item = MAPPER.createObjectNode(); + item.put("name", rs.getString("name")); + item.put("object_type", "SEQUENCE"); + putNullable(item, "schema", schemaLabel); + item.putNull("comment"); + result.add(item); + } + } + } + String synonymSql = "SELECT synonym_name AS name FROM all_synonyms WHERE owner = ? ORDER BY synonym_name"; + try (PreparedStatement ps = conn.prepareStatement(synonymSql)) { + ps.setString(1, owner); + try (ResultSet rs = ps.executeQuery()) { + while (rs.next()) { + ObjectNode item = MAPPER.createObjectNode(); + item.put("name", rs.getString("name")); + item.put("object_type", "SYNONYM"); + putNullable(item, "schema", schemaLabel); + item.putNull("comment"); + result.add(item); + } + } + } + return result; + } + + private static JsonNode oracleListIndexes(Connection conn, String owner, String table) throws SQLException { + String sql = + "SELECT i.index_name, i.uniqueness, i.index_type, ic.column_name, " + + "CASE WHEN pk.index_name IS NULL THEN 0 ELSE 1 END AS is_primary " + + "FROM all_indexes i " + + "JOIN all_ind_columns ic ON ic.index_owner = i.owner AND ic.index_name = i.index_name " + + "AND ic.table_owner = i.table_owner AND ic.table_name = i.table_name " + + "LEFT JOIN (SELECT owner, index_name, table_name FROM all_constraints WHERE constraint_type = 'P') pk " + + "ON pk.owner = i.owner AND pk.index_name = i.index_name AND pk.table_name = i.table_name " + + "WHERE i.table_owner = ? AND i.table_name = ? " + + "ORDER BY i.index_name, ic.column_position"; + LinkedHashMap indexes = new LinkedHashMap<>(); + try (PreparedStatement ps = conn.prepareStatement(sql)) { + ps.setString(1, owner); + ps.setString(2, table); + try (ResultSet rs = ps.executeQuery()) { + while (rs.next()) { + String name = rs.getString("index_name"); + ObjectNode item = indexes.get(name); + if (item == null) { + item = indexNode( + name, + "UNIQUE".equalsIgnoreCase(rs.getString("uniqueness")), + rs.getInt("is_primary") != 0, + rs.getString("index_type") + ); + indexes.put(name, item); + } + String column = rs.getString("column_name"); + if (column != null && !column.isBlank()) { + ((ArrayNode) item.path("columns")).add(column); + } + } + } + } + ArrayNode result = MAPPER.createArrayNode(); + indexes.values().forEach(result::add); return result; } 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 50d5a6cbcb..db45c822e4 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -2582,6 +2582,105 @@ final class DbxJdbcPluginTest { assertEquals("PACKAGE", method.invoke(null, "PACKAGE")); } + @Test + void listIndexesReturnsPrimaryAndUniqueIndexes() throws Exception { + request("executeQuery", """ + { + "connection": %s, + "sql": "DROP TABLE IF EXISTS codex_jdbc_indexes" + } + """.formatted(CONNECTION)); + request("executeQuery", """ + { + "connection": %s, + "sql": "CREATE TABLE codex_jdbc_indexes (id INT CONSTRAINT codex_jdbc_indexes_pk PRIMARY KEY, code VARCHAR(32), CONSTRAINT codex_jdbc_indexes_code_uq UNIQUE(code))" + } + """.formatted(CONNECTION)); + + JsonNode response = request("listIndexes", """ + { + "connection": %s, + "schema": "PUBLIC", + "table": "CODEX_JDBC_INDEXES" + } + """.formatted(CONNECTION)); + + assertFalse(response.has("error"), response.toString()); + JsonNode indexes = response.path("result"); + JsonNode primary = findByName(indexes, "CODEX_JDBC_INDEXES_PK_INDEX_4"); + if (primary == null) { + primary = findIndexByColumn(indexes, "ID"); + } + JsonNode unique = findByName(indexes, "CODEX_JDBC_INDEXES_CODE_UQ_INDEX_4"); + if (unique == null) { + unique = findIndexByColumn(indexes, "CODE"); + } + assertEquals(true, primary.path("is_primary").asBoolean(), indexes.toString()); + assertEquals(true, primary.path("is_unique").asBoolean(), indexes.toString()); + assertEquals(true, unique.path("is_unique").asBoolean(), indexes.toString()); + assertEquals(false, unique.path("is_primary").asBoolean(), indexes.toString()); + } + + @Test + void oracleListIndexesGroupsColumnsAndMarksPrimaryKey() throws Exception { + ResultSet rows = rowsResultSet( + new String[] { "index_name", "uniqueness", "index_type", "column_name", "is_primary" }, + new Object[][] { + { "CODEX_7046_CODE_IDX", "NONUNIQUE", "NORMAL", "CODE", 0 }, + { "CODEX_7046_PK", "UNIQUE", "NORMAL", "ID", 1 }, + { "CODEX_7046_PK", "UNIQUE", "NORMAL", "TENANT_ID", 1 } + } + ); + Connection conn = preparedStatementConnection(rows); + Method method = DbxJdbcPlugin.class.getDeclaredMethod( + "oracleListIndexes", + Connection.class, + String.class, + String.class + ); + method.setAccessible(true); + + JsonNode indexes = (JsonNode) method.invoke(null, conn, "DBX_TEST", "CODEX_7046_META"); + + JsonNode primary = findByName(indexes, "CODEX_7046_PK"); + assertEquals(true, primary.path("is_primary").asBoolean()); + assertEquals(true, primary.path("is_unique").asBoolean()); + assertEquals("ID", primary.path("columns").path(0).asText()); + assertEquals("TENANT_ID", primary.path("columns").path(1).asText()); + JsonNode secondary = findByName(indexes, "CODEX_7046_CODE_IDX"); + assertEquals(false, secondary.path("is_primary").asBoolean()); + assertEquals(false, secondary.path("is_unique").asBoolean()); + } + + @Test + void listIndexesDegradesToEmptyWhenDriverRejectsIndexMetadata() throws Exception { + for (Throwable failure : List.of( + new SQLFeatureNotSupportedException("indexes not supported"), + new UnsupportedOperationException("unsupported"), + new AbstractMethodError("unsupported") + )) { + String connection = """ + { "connection_string": "jdbc:dbx-index-metadata:%s" } + """.formatted(failure.getClass().getSimpleName()); + Driver driver = testDriver("jdbc:dbx-index-metadata:", indexMetadataConnection(failure)); + DriverManager.registerDriver(driver); + try { + JsonNode response = request("listIndexes", """ + { + "connection": %s, + "schema": "PUBLIC", + "table": "SOME_TABLE" + } + """.formatted(connection)); + + assertFalse(response.has("error"), failure.getClass().getSimpleName() + ": " + response); + assertEquals(0, response.path("result").size(), failure.getClass().getSimpleName() + ": " + response); + } finally { + closeAndDeregister(connection, driver); + } + } + } + @Test void oracleEffectiveSchemaUsesExactOwnerBeforeUppercaseFallback() throws Exception { Method method = DbxJdbcPlugin.class.getDeclaredMethod("oracleEffectiveSchema", Connection.class, String.class); @@ -3027,6 +3126,39 @@ final class DbxJdbcPluginTest { ); } + private static Connection preparedStatementConnection(ResultSet rs) { + return (Connection) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Connection.class }, + (proxy, method, args) -> switch (method.getName()) { + case "prepareStatement" -> preparedStatement(rs); + case "isClosed" -> false; + case "close" -> null; + default -> defaultValue(method.getReturnType()); + } + ); + } + + private static JsonNode findByName(JsonNode items, String name) { + for (JsonNode item : items) { + if (name.equalsIgnoreCase(item.path("name").asText())) { + return item; + } + } + return null; + } + + private static JsonNode findIndexByColumn(JsonNode indexes, String column) { + for (JsonNode index : indexes) { + for (JsonNode value : index.path("columns")) { + if (column.equalsIgnoreCase(value.asText())) { + return index; + } + } + } + return null; + } + private static ResultSet rowsResultSet(String[] columns, Object[][] rows) { return (ResultSet) Proxy.newProxyInstance( DbxJdbcPluginTest.class.getClassLoader(), @@ -3291,6 +3423,29 @@ final class DbxJdbcPluginTest { ); } + private static Connection indexMetadataConnection(Throwable failure) { + DatabaseMetaData metadata = (DatabaseMetaData) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { DatabaseMetaData.class }, + (proxy, method, args) -> { + if ("getPrimaryKeys".equals(method.getName()) || "getIndexInfo".equals(method.getName())) { + throw failure; + } + return defaultValue(method.getReturnType()); + } + ); + return (Connection) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Connection.class }, + (proxy, method, args) -> switch (method.getName()) { + case "getMetaData" -> metadata; + case "isClosed" -> false; + case "close" -> null; + default -> defaultValue(method.getReturnType()); + } + ); + } + private static Connection transactionMetadataConnection(Throwable metadataFailure, boolean supportsTransactions) { boolean[] autoCommit = { true }; DatabaseMetaData metadata = (DatabaseMetaData) Proxy.newProxyInstance(