Skip to content

Commit

Permalink
Collect table type in PostgreSQLMetaDataLoader (#32634)
Browse files Browse the repository at this point in the history 
* Collect table type in PostgreSQLMetaDataLoader

* Fix checkstyle in try block

* Change LinkedList to HashSet
  • Loading branch information
@RaigorJiang
RaigorJiang authored Aug 22, 2024
1 parent 103afe6 commit 0ecc7d9
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 20 deletions.
4 changes: 3 additions & 1 deletion ...he/shardingsphere/infra/database/core/metadata/data/loader/type/ColumnMetaDataLoader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,9 @@ public static Collection<ColumnMetaData> load(final Connection connection, final
}
}
}
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(generateEmptyResultSQL(tableNamePattern, columnNames, databaseType))) {
try (
Statement statement = connection.createStatement();
ResultSet resultSet = statement.executeQuery(generateEmptyResultSQL(tableNamePattern, columnNames, databaseType))) {
for (int i = 0; i < columnNames.size(); i++) {
boolean generated = resultSet.getMetaData().isAutoIncrement(i + 1);
caseSensitiveFlags.add(resultSet.getMetaData().isCaseSensitive(resultSet.findColumn(columnNames.get(i))));
Expand Down
6 changes: 3 additions & 3 deletions .../apache/shardingsphere/infra/database/mysql/metadata/data/loader/MySQLMetaDataLoader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ private String getViewMetaDataSQL(final Collection<String> tableNames) {
}

private Map<String, Collection<ConstraintMetaData>> loadConstraintMetaDataMap(final DataSource dataSource, final Collection<String> tables) throws SQLException {
Map<String, Collection<ConstraintMetaData>> result = new LinkedHashMap<>();
Map<String, Collection<ConstraintMetaData>> result = new LinkedHashMap<>(tables.size(), 1F);
try (
Connection connection = dataSource.getConnection();
PreparedStatement preparedStatement = connection.prepareStatement(getConstraintMetaDataSQL(tables))) {
Expand All @@ -127,7 +127,7 @@ private String getConstraintMetaDataSQL(final Collection<String> tableNames) {
}

private Map<String, Collection<ColumnMetaData>> loadColumnMetaDataMap(final DataSource dataSource, final Collection<String> tables) throws SQLException {
Map<String, Collection<ColumnMetaData>> result = new HashMap<>();
Map<String, Collection<ColumnMetaData>> result = new HashMap<>(tables.size(), 1F);
try (
Connection connection = dataSource.getConnection();
PreparedStatement preparedStatement = connection.prepareStatement(getTableMetaDataSQL(tables))) {
Expand Down Expand Up @@ -167,7 +167,7 @@ private String getTableMetaDataSQL(final Collection<String> tables) {
}

private Map<String, Collection<IndexMetaData>> loadIndexMetaData(final DataSource dataSource, final Collection<String> tableNames) throws SQLException {
Map<String, Map<String, IndexMetaData>> tableToIndex = new HashMap<>();
Map<String, Map<String, IndexMetaData>> tableToIndex = new HashMap<>(tableNames.size(), 1F);
try (
Connection connection = dataSource.getConnection();
PreparedStatement preparedStatement = connection.prepareStatement(getIndexMetaDataSQL(tableNames))) {
Expand Down
16 changes: 12 additions & 4 deletions ...shardingsphere/infra/database/opengauss/metadata/data/loader/OpenGaussMetaDataLoader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,9 @@ public Collection<SchemaMetaData> load(final MetaDataLoaderMaterial material) th

private Map<String, Multimap<String, IndexMetaData>> loadIndexMetaDataMap(final Connection connection, final Collection<String> schemaNames) throws SQLException {
Map<String, Multimap<String, IndexMetaData>> result = new LinkedHashMap<>();
try (PreparedStatement preparedStatement = connection.prepareStatement(getIndexMetaDataSQL(schemaNames)); ResultSet resultSet = preparedStatement.executeQuery()) {
try (
PreparedStatement preparedStatement = connection.prepareStatement(getIndexMetaDataSQL(schemaNames));
ResultSet resultSet = preparedStatement.executeQuery()) {
while (resultSet.next()) {
String schemaName = resultSet.getString("schemaname");
String tableName = resultSet.getString("tablename");
Expand All @@ -93,7 +95,9 @@ private Map<String, Multimap<String, IndexMetaData>> loadIndexMetaDataMap(final
indexMetaDataMap.put(tableName, new IndexMetaData(indexName));
}
}
try (PreparedStatement preparedStatement = connection.prepareStatement(getAdvanceIndexMetaDataSQL(schemaNames)); ResultSet resultSet = preparedStatement.executeQuery()) {
try (
PreparedStatement preparedStatement = connection.prepareStatement(getAdvanceIndexMetaDataSQL(schemaNames));
ResultSet resultSet = preparedStatement.executeQuery()) {
while (resultSet.next()) {
String schemaName = resultSet.getString("index_schema");
String tableName = resultSet.getString("table_name");
Expand Down Expand Up @@ -125,7 +129,9 @@ private String getAdvanceIndexMetaDataSQL(final Collection<String> schemaNames)
private Map<String, Multimap<String, ColumnMetaData>> loadColumnMetaDataMap(final Connection connection, final Collection<String> tables,
final Collection<String> schemaNames) throws SQLException {
Map<String, Multimap<String, ColumnMetaData>> result = new LinkedHashMap<>();
try (PreparedStatement preparedStatement = connection.prepareStatement(getColumnMetaDataSQL(schemaNames, tables)); ResultSet resultSet = preparedStatement.executeQuery()) {
try (
PreparedStatement preparedStatement = connection.prepareStatement(getColumnMetaDataSQL(schemaNames, tables));
ResultSet resultSet = preparedStatement.executeQuery()) {
Map<String, Integer> dataTypes = new DataTypeLoader().load(connection.getMetaData(), getType());
Collection<String> primaryKeys = loadPrimaryKeys(connection, schemaNames);
while (resultSet.next()) {
Expand All @@ -146,7 +152,9 @@ private String getColumnMetaDataSQL(final Collection<String> schemaNames, final

private Collection<String> loadPrimaryKeys(final Connection connection, final Collection<String> schemaNames) throws SQLException {
Collection<String> result = new HashSet<>();
try (PreparedStatement preparedStatement = connection.prepareStatement(getPrimaryKeyMetaDataSQL(schemaNames)); ResultSet resultSet = preparedStatement.executeQuery()) {
try (
PreparedStatement preparedStatement = connection.prepareStatement(getPrimaryKeyMetaDataSQL(schemaNames));
ResultSet resultSet = preparedStatement.executeQuery()) {
while (resultSet.next()) {
String schemaName = resultSet.getString("table_schema");
String tableName = resultSet.getString("table_name");
Expand Down
62 changes: 50 additions & 12 deletions ...ardingsphere/infra/database/postgresql/metadata/data/loader/PostgreSQLMetaDataLoader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import org.apache.shardingsphere.infra.database.core.metadata.data.model.SchemaMetaData;
import org.apache.shardingsphere.infra.database.core.metadata.data.model.TableMetaData;
import org.apache.shardingsphere.infra.database.core.metadata.database.datatype.DataTypeLoader;
import org.apache.shardingsphere.infra.database.core.metadata.database.enums.TableType;
import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader;

Expand All @@ -36,6 +37,7 @@
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedList;
Expand Down Expand Up @@ -76,27 +78,33 @@ public final class PostgreSQLMetaDataLoader implements DialectMetaDataLoader {

private static final String LOAD_FILTERED_ROLE_TABLE_GRANTS_SQL = LOAD_ALL_ROLE_TABLE_GRANTS_SQL + " WHERE table_name IN (%s)";

private static final String VIEW_META_DATA_SQL = "SELECT table_schema, table_name FROM information_schema.views WHERE table_schema IN (%s) and table_name IN (%s)";

@Override
public Collection<SchemaMetaData> load(final MetaDataLoaderMaterial material) throws SQLException {
try (Connection connection = material.getDataSource().getConnection()) {
Collection<String> schemaNames = SchemaMetaDataLoader.loadSchemaNames(connection, TypedSPILoader.getService(DatabaseType.class, "PostgreSQL"));
Map<String, Multimap<String, IndexMetaData>> schemaIndexMetaDataMap = loadIndexMetaDataMap(connection, schemaNames);
Map<String, Multimap<String, ColumnMetaData>> schemaColumnMetaDataMap = loadColumnMetaDataMap(connection, material.getActualTableNames(), schemaNames);
Map<String, Multimap<String, ConstraintMetaData>> schemaConstraintMetaDataMap = loadConstraintMetaDataMap(connection, schemaNames);
Map<String, Collection<String>> schemaViewNames = loadViewNames(connection, schemaNames, material.getActualTableNames());
Collection<SchemaMetaData> result = new LinkedList<>();
for (String each : schemaNames) {
Multimap<String, IndexMetaData> tableIndexMetaDataMap = schemaIndexMetaDataMap.getOrDefault(each, LinkedHashMultimap.create());
Multimap<String, ColumnMetaData> tableColumnMetaDataMap = schemaColumnMetaDataMap.getOrDefault(each, LinkedHashMultimap.create());
Multimap<String, ConstraintMetaData> tableConstraintMetaDataMap = schemaConstraintMetaDataMap.getOrDefault(each, LinkedHashMultimap.create());
result.add(new SchemaMetaData(each, createTableMetaDataList(tableIndexMetaDataMap, tableColumnMetaDataMap, tableConstraintMetaDataMap)));
Collection<String> viewNames = schemaViewNames.getOrDefault(each, Collections.emptySet());
result.add(new SchemaMetaData(each, createTableMetaDataList(tableIndexMetaDataMap, tableColumnMetaDataMap, tableConstraintMetaDataMap, viewNames)));
}
return result;
}
}

private Map<String, Multimap<String, IndexMetaData>> loadIndexMetaDataMap(final Connection connection, final Collection<String> schemaNames) throws SQLException {
Map<String, Multimap<String, IndexMetaData>> result = new LinkedHashMap<>();
try (PreparedStatement preparedStatement = connection.prepareStatement(getIndexMetaDataSQL(schemaNames)); ResultSet resultSet = preparedStatement.executeQuery()) {
Map<String, Multimap<String, IndexMetaData>> result = new LinkedHashMap<>(schemaNames.size(), 1F);
try (
PreparedStatement preparedStatement = connection.prepareStatement(getIndexMetaDataSQL(schemaNames));
ResultSet resultSet = preparedStatement.executeQuery()) {
while (resultSet.next()) {
String schemaName = resultSet.getString("schemaname");
String tableName = resultSet.getString("tablename");
Expand All @@ -105,7 +113,9 @@ private Map<String, Multimap<String, IndexMetaData>> loadIndexMetaDataMap(final
indexMetaDataMap.put(tableName, new IndexMetaData(indexName));
}
}
try (PreparedStatement preparedStatement = connection.prepareStatement(getAdvanceIndexMetaDataSQL(schemaNames)); ResultSet resultSet = preparedStatement.executeQuery()) {
try (
PreparedStatement preparedStatement = connection.prepareStatement(getAdvanceIndexMetaDataSQL(schemaNames));
ResultSet resultSet = preparedStatement.executeQuery()) {
while (resultSet.next()) {
String schemaName = resultSet.getString("index_schema");
String tableName = resultSet.getString("table_name");
Expand Down Expand Up @@ -136,9 +146,11 @@ private String getAdvanceIndexMetaDataSQL(final Collection<String> schemaNames)

private Map<String, Multimap<String, ColumnMetaData>> loadColumnMetaDataMap(final Connection connection, final Collection<String> tables,
final Collection<String> schemaNames) throws SQLException {
Map<String, Multimap<String, ColumnMetaData>> result = new LinkedHashMap<>();
Map<String, Multimap<String, ColumnMetaData>> result = new LinkedHashMap<>(schemaNames.size(), 1F);
Collection<String> roleTableGrants = loadRoleTableGrants(connection, tables);
try (PreparedStatement preparedStatement = connection.prepareStatement(getColumnMetaDataSQL(schemaNames, tables)); ResultSet resultSet = preparedStatement.executeQuery()) {
try (
PreparedStatement preparedStatement = connection.prepareStatement(getColumnMetaDataSQL(schemaNames, tables));
ResultSet resultSet = preparedStatement.executeQuery()) {
Map<String, Integer> dataTypes = new DataTypeLoader().load(connection.getMetaData(), getType());
Collection<String> primaryKeys = loadPrimaryKeys(connection, schemaNames);
while (resultSet.next()) {
Expand All @@ -156,7 +168,9 @@ private Map<String, Multimap<String, ColumnMetaData>> loadColumnMetaDataMap(fina

private Collection<String> loadRoleTableGrants(final Connection connection, final Collection<String> tables) throws SQLException {
Collection<String> result = new HashSet<>(tables.size(), 1F);
try (PreparedStatement preparedStatement = connection.prepareStatement(getLoadRoleTableGrantsSQL(tables)); ResultSet resultSet = preparedStatement.executeQuery()) {
try (
PreparedStatement preparedStatement = connection.prepareStatement(getLoadRoleTableGrantsSQL(tables));
ResultSet resultSet = preparedStatement.executeQuery()) {
while (resultSet.next()) {
result.add(resultSet.getString("table_name"));
}
Expand All @@ -177,7 +191,9 @@ private String getColumnMetaDataSQL(final Collection<String> schemaNames, final

private Set<String> loadPrimaryKeys(final Connection connection, final Collection<String> schemaNames) throws SQLException {
Set<String> result = new HashSet<>();
try (PreparedStatement preparedStatement = connection.prepareStatement(getPrimaryKeyMetaDataSQL(schemaNames)); ResultSet resultSet = preparedStatement.executeQuery()) {
try (
PreparedStatement preparedStatement = connection.prepareStatement(getPrimaryKeyMetaDataSQL(schemaNames));
ResultSet resultSet = preparedStatement.executeQuery()) {
while (resultSet.next()) {
String schemaName = resultSet.getString("table_schema");
String tableName = resultSet.getString("table_name");
Expand Down Expand Up @@ -207,8 +223,10 @@ private ColumnMetaData loadColumnMetaData(final Map<String, Integer> dataTypeMap
}

private Map<String, Multimap<String, ConstraintMetaData>> loadConstraintMetaDataMap(final Connection connection, final Collection<String> schemaNames) throws SQLException {
Map<String, Multimap<String, ConstraintMetaData>> result = new LinkedHashMap<>();
try (PreparedStatement preparedStatement = connection.prepareStatement(getConstraintKeyMetaDataSQL(schemaNames)); ResultSet resultSet = preparedStatement.executeQuery()) {
Map<String, Multimap<String, ConstraintMetaData>> result = new LinkedHashMap<>(schemaNames.size(), 1F);
try (
PreparedStatement preparedStatement = connection.prepareStatement(getConstraintKeyMetaDataSQL(schemaNames));
ResultSet resultSet = preparedStatement.executeQuery()) {
while (resultSet.next()) {
String schemaName = resultSet.getString("table_schema");
Multimap<String, ConstraintMetaData> constraintMetaData = result.computeIfAbsent(schemaName, key -> LinkedHashMultimap.create());
Expand All @@ -225,14 +243,34 @@ private String getConstraintKeyMetaDataSQL(final Collection<String> schemaNames)
return String.format(FOREIGN_KEY_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
}

private Map<String, Collection<String>> loadViewNames(final Connection connection, final Collection<String> schemaNames, final Collection<String> tables) throws SQLException {
Map<String, Collection<String>> result = new LinkedHashMap<>(schemaNames.size(), 1F);
try (
PreparedStatement preparedStatement = connection.prepareStatement(getViewMetaDataSQL(schemaNames, tables));
ResultSet resultSet = preparedStatement.executeQuery()) {
while (resultSet.next()) {
String schemaName = resultSet.getString("table_schema");
Collection<String> viewMetaData = result.computeIfAbsent(schemaName, key -> new HashSet<>());
String tableName = resultSet.getString("table_name");
viewMetaData.add(tableName);
}
}
return result;
}

private String getViewMetaDataSQL(final Collection<String> schemaNames, final Collection<String> tables) {
return String.format(VIEW_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")),
tables.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
}

private Collection<TableMetaData> createTableMetaDataList(final Multimap<String, IndexMetaData> tableIndexMetaDataMap, final Multimap<String, ColumnMetaData> tableColumnMetaDataMap,
final Multimap<String, ConstraintMetaData> tableConstraintMetaDataMap) {
final Multimap<String, ConstraintMetaData> tableConstraintMetaDataMap, final Collection<String> viewNames) {
Collection<TableMetaData> result = new LinkedList<>();
for (String each : tableColumnMetaDataMap.keySet()) {
Collection<ColumnMetaData> columnMetaDataList = tableColumnMetaDataMap.get(each);
Collection<IndexMetaData> indexMetaDataList = tableIndexMetaDataMap.get(each);
Collection<ConstraintMetaData> constraintMetaDataList = tableConstraintMetaDataMap.get(each);
result.add(new TableMetaData(each, columnMetaDataList, indexMetaDataList, constraintMetaDataList));
result.add(new TableMetaData(each, columnMetaDataList, indexMetaDataList, constraintMetaDataList, viewNames.contains(each) ? TableType.VIEW : TableType.TABLE));
}
return result;
}
Expand Down

0 comments on commit 0ecc7d9

Please sign in to comment.