Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Collect table type in PostgreSQLMetaDataLoader #32634

Merged
merged 3 commits into from
Aug 22, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
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
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
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