From 5242e252f58000ea4829124fc1dc7dfac10ae2d0 Mon Sep 17 00:00:00 2001 From: Qi Yu Date: Tue, 9 Apr 2024 14:39:24 +0800 Subject: [PATCH] [#2085] improvement(jdbc-mysql): Remove depends on JDBC system table. (#2616) ### What changes were proposed in this pull request? Use the JDBC connection meta to replace it with the JDBC system table to obtain schema information. ### Why are the changes needed? We shouldn't assume that we have access to system tables, which are frequently protected by ACLs. Fix: #2085 ### Does this PR introduce _any_ user-facing change? N/A. ### How was this patch tested? Current tests can cover this changs. --- .../operation/MysqlDatabaseOperations.java | 51 ++++---------- .../operation/PostgreSqlSchemaOperations.java | 69 +++++++++++-------- 2 files changed, 56 insertions(+), 64 deletions(-) diff --git a/catalogs/catalog-jdbc-mysql/src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlDatabaseOperations.java b/catalogs/catalog-jdbc-mysql/src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlDatabaseOperations.java index ca6a6eb74fc..9971b9ee490 100644 --- a/catalogs/catalog-jdbc-mysql/src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlDatabaseOperations.java +++ b/catalogs/catalog-jdbc-mysql/src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlDatabaseOperations.java @@ -9,14 +9,14 @@ import com.datastrato.gravitino.catalog.jdbc.operation.JdbcDatabaseOperations; import com.datastrato.gravitino.exceptions.NoSuchSchemaException; import com.datastrato.gravitino.meta.AuditInfo; +import com.google.common.collect.ImmutableMap; import java.sql.Connection; -import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; import java.util.Collections; -import java.util.HashMap; import java.util.HashSet; +import java.util.List; import java.util.Locale; import java.util.Map; import java.util.Set; @@ -53,11 +53,6 @@ public String generateCreateDatabaseSql( if (MapUtils.isNotEmpty(properties)) { // TODO #804 Properties will be unified in the future. throw new UnsupportedOperationException("Properties are not supported yet"); - // sqlBuilder.append("\n"); - // sqlBuilder.append( - // properties.entrySet().stream() - // .map(entry -> entry.getKey() + " " + entry.getValue()) - // .collect(Collectors.joining("\n"))); } String result = sqlBuilder.toString(); LOG.info("Generated create database:{} sql: {}", databaseName, result); @@ -92,37 +87,19 @@ public String generateDropDatabaseSql(String databaseName, boolean cascade) { @Override public JdbcSchema load(String databaseName) throws NoSuchSchemaException { - try (final Connection connection = this.dataSource.getConnection()) { - String query = "SELECT * FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = ?"; - try (PreparedStatement preparedStatement = connection.prepareStatement(query)) { - preparedStatement.setString(1, databaseName); - - // Execute the query - try (ResultSet resultSet = preparedStatement.executeQuery()) { - if (!resultSet.next()) { - throw new NoSuchSchemaException( - "Database %s could not be found in information_schema.SCHEMATA", databaseName); - } - String schemaName = resultSet.getString("SCHEMA_NAME"); - // Mysql currently only supports these two attributes - String characterSetName = resultSet.getString("DEFAULT_CHARACTER_SET_NAME"); - String collationName = resultSet.getString("DEFAULT_COLLATION_NAME"); - Map properties = new HashMap<>(); - properties.put("CHARACTER SET", characterSetName); - properties.put("COLLATE", collationName); - - JdbcSchema.Builder builder = - JdbcSchema.builder() - .withName(schemaName) - .withProperties(properties) - .withAuditInfo(AuditInfo.EMPTY); + List allDatabases = listDatabases(); + String dbName = + allDatabases.stream() + .filter(db -> db.equals(databaseName)) + .findFirst() + .orElseThrow( + () -> new NoSuchSchemaException("Database %s could not be found", databaseName)); - return builder.build(); - } - } - } catch (final SQLException se) { - throw this.exceptionMapper.toGravitinoException(se); - } + return JdbcSchema.builder() + .withName(dbName) + .withProperties(ImmutableMap.of()) + .withAuditInfo(AuditInfo.EMPTY) + .build(); } @Override diff --git a/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlSchemaOperations.java b/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlSchemaOperations.java index b0a9e106a6a..978121434b7 100644 --- a/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlSchemaOperations.java +++ b/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlSchemaOperations.java @@ -11,6 +11,7 @@ import com.datastrato.gravitino.exceptions.NoSuchSchemaException; import com.datastrato.gravitino.meta.AuditInfo; import java.sql.Connection; +import java.sql.DatabaseMetaData; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; @@ -20,6 +21,7 @@ import java.util.List; import java.util.Locale; import java.util.Map; +import java.util.Objects; import java.util.Set; import javax.sql.DataSource; import org.apache.commons.collections4.MapUtils; @@ -50,25 +52,27 @@ public void initialize( @Override public JdbcSchema load(String schema) throws NoSuchSchemaException { try (Connection connection = getConnection()) { - String sql = - "SELECT schema_name FROM information_schema.schemata WHERE schema_name = ? AND catalog_name = ?"; - try (PreparedStatement statement = connection.prepareStatement(sql)) { - statement.setString(1, schema); - statement.setString(2, database); - try (ResultSet resultSet = statement.executeQuery()) { - if (!resultSet.next()) { - throw new NoSuchSchemaException("No such schema: %s", schema); - } - String schemaName = resultSet.getString(1); - String comment = getSchemaComment(schema, connection); - return JdbcSchema.builder() - .withName(schemaName) - .withComment(comment) - .withAuditInfo(AuditInfo.EMPTY) - .withProperties(Collections.emptyMap()) - .build(); + ResultSet resultSet = getSchema(connection, schema); + + boolean found = false; + while (resultSet.next()) { + if (Objects.equals(resultSet.getString(1), schema)) { + found = true; + break; } } + + if (!found) { + throw new NoSuchSchemaException("No such schema: %s", schema); + } + + String comment = getSchemaComment(schema, connection); + return JdbcSchema.builder() + .withName(schema) + .withComment(comment) + .withAuditInfo(AuditInfo.EMPTY) + .withProperties(Collections.emptyMap()) + .build(); } catch (SQLException e) { throw exceptionMapper.toGravitinoException(e); } @@ -78,16 +82,11 @@ public JdbcSchema load(String schema) throws NoSuchSchemaException { public List listDatabases() { List result = new ArrayList<>(); try (Connection connection = getConnection()) { - try (PreparedStatement statement = - connection.prepareStatement( - "SELECT schema_name FROM information_schema.schemata WHERE catalog_name = ?")) { - statement.setString(1, database); - ResultSet resultSet = statement.executeQuery(); - while (resultSet.next()) { - String databaseName = resultSet.getString(1); - if (!isSystemDatabase(databaseName)) { - result.add(databaseName); - } + ResultSet resultSet = getSchema(connection, null); + while (resultSet.next()) { + String schemaName = resultSet.getString(1); + if (!isSystemDatabase(schemaName)) { + result.add(resultSet.getString(1)); } } } catch (final SQLException se) { @@ -116,6 +115,22 @@ public String generateCreateDatabaseSql( return sqlBuilder.toString(); } + /** + * Get the schema with the given name. + * + *

Note: This method will return a result set that may contain multiple rows as the schemaName + * in `getSchemas` is a pattern. The result set will contain all schemas that match the pattern. + * + *

Database in PG corresponds to Catalog in JDBC. Schema in PG corresponds to Schema in JDBC. + * + * @param connection the connection to the database + * @param schemaName the name of the schema + */ + private ResultSet getSchema(Connection connection, String schemaName) throws SQLException { + final DatabaseMetaData metaData = connection.getMetaData(); + return metaData.getSchemas(database, schemaName); + } + @Override public String generateDropDatabaseSql(String schema, boolean cascade) { StringBuilder sqlBuilder = new StringBuilder(String.format("DROP SCHEMA %s", schema));