1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 package org.apache.shardingsphere.infra.database.postgresql.metadata.data.loader;
19
20 import com.google.common.collect.LinkedHashMultimap;
21 import com.google.common.collect.Multimap;
22 import org.apache.shardingsphere.infra.database.core.metadata.data.loader.DialectMetaDataLoader;
23 import org.apache.shardingsphere.infra.database.core.metadata.data.loader.MetaDataLoaderMaterial;
24 import org.apache.shardingsphere.infra.database.core.metadata.data.loader.type.SchemaMetaDataLoader;
25 import org.apache.shardingsphere.infra.database.core.metadata.data.model.ColumnMetaData;
26 import org.apache.shardingsphere.infra.database.core.metadata.data.model.ConstraintMetaData;
27 import org.apache.shardingsphere.infra.database.core.metadata.data.model.IndexMetaData;
28 import org.apache.shardingsphere.infra.database.core.metadata.data.model.SchemaMetaData;
29 import org.apache.shardingsphere.infra.database.core.metadata.data.model.TableMetaData;
30 import org.apache.shardingsphere.infra.database.core.metadata.database.datatype.DataTypeRegistry;
31 import org.apache.shardingsphere.infra.database.core.metadata.database.enums.TableType;
32
33 import java.sql.Connection;
34 import java.sql.PreparedStatement;
35 import java.sql.ResultSet;
36 import java.sql.SQLException;
37 import java.sql.Types;
38 import java.util.Collection;
39 import java.util.Collections;
40 import java.util.HashSet;
41 import java.util.LinkedHashMap;
42 import java.util.LinkedList;
43 import java.util.Map;
44 import java.util.Optional;
45 import java.util.Set;
46 import java.util.stream.Collectors;
47
48
49
50
51 public final class PostgreSQLMetaDataLoader implements DialectMetaDataLoader {
52
53 private static final String BASIC_TABLE_META_DATA_SQL = "SELECT table_name, column_name, ordinal_position, data_type, udt_name, column_default, table_schema, is_nullable"
54 + " FROM information_schema.columns WHERE table_schema IN (%s)";
55
56 private static final String TABLE_META_DATA_SQL_WITHOUT_TABLES = BASIC_TABLE_META_DATA_SQL + " ORDER BY ordinal_position";
57
58 private static final String TABLE_META_DATA_SQL_WITH_TABLES = BASIC_TABLE_META_DATA_SQL + " AND table_name IN (%s) ORDER BY ordinal_position";
59
60 private static final String FOREIGN_KEY_META_DATA_SQL = "SELECT tc.table_schema,tc.table_name,tc.constraint_name,pgo.relname refer_table_name "
61 + "FROM information_schema.table_constraints tc "
62 + "JOIN pg_constraint pgc ON tc.constraint_name = pgc.conname AND contype='f' "
63 + "JOIN pg_class pgo ON pgc.confrelid = pgo.oid WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_schema IN (%s)";
64
65 private static final String PRIMARY_KEY_META_DATA_SQL = "SELECT tc.table_name, kc.column_name, kc.table_schema FROM information_schema.table_constraints tc"
66 + " JOIN information_schema.key_column_usage kc ON kc.table_schema = tc.table_schema AND kc.table_name = tc.table_name AND kc.constraint_name = tc.constraint_name"
67 + " WHERE tc.constraint_type = 'PRIMARY KEY' AND kc.ordinal_position IS NOT NULL AND kc.table_schema IN (%s)";
68
69 private static final String BASIC_INDEX_META_DATA_SQL = "SELECT tablename, indexname, schemaname FROM pg_indexes WHERE schemaname IN (%s)";
70
71 private static final String ADVANCE_INDEX_META_DATA_SQL =
72 "SELECT idx.relname as index_name, insp.nspname as index_schema, tbl.relname as table_name, att.attname AS column_name, pgi.indisunique as is_unique"
73 + " FROM pg_index pgi JOIN pg_class idx ON idx.oid = pgi.indexrelid JOIN pg_namespace insp ON insp.oid = idx.relnamespace JOIN pg_class tbl ON tbl.oid = pgi.indrelid"
74 + " JOIN pg_namespace tnsp ON tnsp.oid = tbl.relnamespace JOIN pg_attribute att ON att.attrelid = tbl.oid AND att.attnum = ANY(pgi.indkey) WHERE tnsp.nspname IN (%s)";
75
76 private static final String LOAD_ALL_ROLE_TABLE_GRANTS_SQL = "SELECT table_name FROM information_schema.role_table_grants";
77
78 private static final String LOAD_FILTERED_ROLE_TABLE_GRANTS_SQL = LOAD_ALL_ROLE_TABLE_GRANTS_SQL + " WHERE table_name IN (%s)";
79
80 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)";
81
82 @Override
83 public Collection<SchemaMetaData> load(final MetaDataLoaderMaterial material) throws SQLException {
84 try (Connection connection = material.getDataSource().getConnection()) {
85 Collection<String> schemaNames = new SchemaMetaDataLoader(getType()).loadSchemaNames(connection);
86 Map<String, Multimap<String, IndexMetaData>> schemaIndexMetaDataMap = loadIndexMetaDataMap(connection, schemaNames);
87 Map<String, Multimap<String, ColumnMetaData>> schemaColumnMetaDataMap = loadColumnMetaDataMap(connection, material.getActualTableNames(), schemaNames);
88 Map<String, Multimap<String, ConstraintMetaData>> schemaConstraintMetaDataMap = loadConstraintMetaDataMap(connection, schemaNames);
89 Map<String, Collection<String>> schemaViewNames = loadViewNames(connection, schemaNames, material.getActualTableNames());
90 Collection<SchemaMetaData> result = new LinkedList<>();
91 for (String each : schemaNames) {
92 Multimap<String, IndexMetaData> tableIndexMetaDataMap = schemaIndexMetaDataMap.getOrDefault(each, LinkedHashMultimap.create());
93 Multimap<String, ColumnMetaData> tableColumnMetaDataMap = schemaColumnMetaDataMap.getOrDefault(each, LinkedHashMultimap.create());
94 Multimap<String, ConstraintMetaData> tableConstraintMetaDataMap = schemaConstraintMetaDataMap.getOrDefault(each, LinkedHashMultimap.create());
95 Collection<String> viewNames = schemaViewNames.getOrDefault(each, Collections.emptySet());
96 result.add(new SchemaMetaData(each, createTableMetaDataList(tableIndexMetaDataMap, tableColumnMetaDataMap, tableConstraintMetaDataMap, viewNames)));
97 }
98 return result;
99 }
100 }
101
102 private Map<String, Multimap<String, IndexMetaData>> loadIndexMetaDataMap(final Connection connection, final Collection<String> schemaNames) throws SQLException {
103 Map<String, Multimap<String, IndexMetaData>> result = new LinkedHashMap<>(schemaNames.size(), 1F);
104 try (
105 PreparedStatement preparedStatement = connection.prepareStatement(getIndexMetaDataSQL(schemaNames));
106 ResultSet resultSet = preparedStatement.executeQuery()) {
107 while (resultSet.next()) {
108 String schemaName = resultSet.getString("schemaname");
109 String tableName = resultSet.getString("tablename");
110 String indexName = resultSet.getString("indexname");
111 Multimap<String, IndexMetaData> indexMetaDataMap = result.computeIfAbsent(schemaName, key -> LinkedHashMultimap.create());
112 indexMetaDataMap.put(tableName, new IndexMetaData(indexName));
113 }
114 }
115 try (
116 PreparedStatement preparedStatement = connection.prepareStatement(getAdvanceIndexMetaDataSQL(schemaNames));
117 ResultSet resultSet = preparedStatement.executeQuery()) {
118 while (resultSet.next()) {
119 String schemaName = resultSet.getString("index_schema");
120 String tableName = resultSet.getString("table_name");
121 String columnName = resultSet.getString("column_name");
122 String indexName = resultSet.getString("index_name");
123 boolean isUnique = resultSet.getBoolean("is_unique");
124 Collection<IndexMetaData> indexMetaDatas = result.getOrDefault(schemaName, LinkedHashMultimap.create()).get(tableName);
125 if (indexMetaDatas.isEmpty()) {
126 continue;
127 }
128 Optional<IndexMetaData> indexMetaData = indexMetaDatas.stream().filter(each -> each.getName().equals(indexName)).findFirst();
129 if (indexMetaData.isPresent()) {
130 indexMetaData.get().setUnique(isUnique);
131 indexMetaData.get().getColumns().add(columnName);
132 }
133 }
134 }
135 return result;
136 }
137
138 private String getIndexMetaDataSQL(final Collection<String> schemaNames) {
139 return String.format(BASIC_INDEX_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
140 }
141
142 private String getAdvanceIndexMetaDataSQL(final Collection<String> schemaNames) {
143 return String.format(ADVANCE_INDEX_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
144 }
145
146 private Map<String, Multimap<String, ColumnMetaData>> loadColumnMetaDataMap(final Connection connection, final Collection<String> tables,
147 final Collection<String> schemaNames) throws SQLException {
148 Map<String, Multimap<String, ColumnMetaData>> result = new LinkedHashMap<>(schemaNames.size(), 1F);
149 Collection<String> roleTableGrants = loadRoleTableGrants(connection, tables);
150 try (
151 PreparedStatement preparedStatement = connection.prepareStatement(getColumnMetaDataSQL(schemaNames, tables));
152 ResultSet resultSet = preparedStatement.executeQuery()) {
153 Collection<String> primaryKeys = loadPrimaryKeys(connection, schemaNames);
154 while (resultSet.next()) {
155 String tableName = resultSet.getString("table_name");
156 if (!roleTableGrants.contains(tableName)) {
157 continue;
158 }
159 String schemaName = resultSet.getString("table_schema");
160 Multimap<String, ColumnMetaData> columnMetaDataMap = result.computeIfAbsent(schemaName, key -> LinkedHashMultimap.create());
161 columnMetaDataMap.put(tableName, loadColumnMetaData(primaryKeys, resultSet));
162 }
163 }
164 return result;
165 }
166
167 private Collection<String> loadRoleTableGrants(final Connection connection, final Collection<String> tables) throws SQLException {
168 Collection<String> result = new HashSet<>(tables.size(), 1F);
169 try (
170 PreparedStatement preparedStatement = connection.prepareStatement(getLoadRoleTableGrantsSQL(tables));
171 ResultSet resultSet = preparedStatement.executeQuery()) {
172 while (resultSet.next()) {
173 result.add(resultSet.getString("table_name"));
174 }
175 }
176 return result;
177 }
178
179 private String getLoadRoleTableGrantsSQL(final Collection<String> tables) {
180 return tables.isEmpty() ? LOAD_ALL_ROLE_TABLE_GRANTS_SQL
181 : String.format(LOAD_FILTERED_ROLE_TABLE_GRANTS_SQL, tables.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
182 }
183
184 private String getColumnMetaDataSQL(final Collection<String> schemaNames, final Collection<String> tables) {
185 String schemaNameParam = schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(","));
186 return tables.isEmpty() ? String.format(TABLE_META_DATA_SQL_WITHOUT_TABLES, schemaNameParam)
187 : String.format(TABLE_META_DATA_SQL_WITH_TABLES, schemaNameParam, tables.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
188 }
189
190 private Set<String> loadPrimaryKeys(final Connection connection, final Collection<String> schemaNames) throws SQLException {
191 Set<String> result = new HashSet<>();
192 try (
193 PreparedStatement preparedStatement = connection.prepareStatement(getPrimaryKeyMetaDataSQL(schemaNames));
194 ResultSet resultSet = preparedStatement.executeQuery()) {
195 while (resultSet.next()) {
196 String schemaName = resultSet.getString("table_schema");
197 String tableName = resultSet.getString("table_name");
198 String columnName = resultSet.getString("column_name");
199 result.add(schemaName + "," + tableName + "," + columnName);
200 }
201 }
202 return result;
203 }
204
205 private String getPrimaryKeyMetaDataSQL(final Collection<String> schemaNames) {
206 return String.format(PRIMARY_KEY_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
207 }
208
209 private ColumnMetaData loadColumnMetaData(final Collection<String> primaryKeys, final ResultSet resultSet) throws SQLException {
210 String schemaName = resultSet.getString("table_schema");
211 String tableName = resultSet.getString("table_name");
212 String columnName = resultSet.getString("column_name");
213 String dataType = resultSet.getString("udt_name");
214 boolean isPrimaryKey = primaryKeys.contains(schemaName + "," + tableName + "," + columnName);
215 String columnDefault = resultSet.getString("column_default");
216 boolean generated = null != columnDefault && columnDefault.startsWith("nextval(");
217
218 boolean caseSensitive = true;
219 boolean isNullable = "YES".equals(resultSet.getString("is_nullable"));
220 return new ColumnMetaData(columnName, DataTypeRegistry.getDataType(getDatabaseType(), dataType).orElse(Types.OTHER), isPrimaryKey, generated, caseSensitive, true, false, isNullable);
221 }
222
223 private Map<String, Multimap<String, ConstraintMetaData>> loadConstraintMetaDataMap(final Connection connection, final Collection<String> schemaNames) throws SQLException {
224 Map<String, Multimap<String, ConstraintMetaData>> result = new LinkedHashMap<>(schemaNames.size(), 1F);
225 try (
226 PreparedStatement preparedStatement = connection.prepareStatement(getConstraintKeyMetaDataSQL(schemaNames));
227 ResultSet resultSet = preparedStatement.executeQuery()) {
228 while (resultSet.next()) {
229 String schemaName = resultSet.getString("table_schema");
230 Multimap<String, ConstraintMetaData> constraintMetaData = result.computeIfAbsent(schemaName, key -> LinkedHashMultimap.create());
231 String tableName = resultSet.getString("table_name");
232 String constraintName = resultSet.getString("constraint_name");
233 String referencedTableName = resultSet.getString("refer_table_name");
234 constraintMetaData.put(tableName, new ConstraintMetaData(constraintName, referencedTableName));
235 }
236 }
237 return result;
238 }
239
240 private String getConstraintKeyMetaDataSQL(final Collection<String> schemaNames) {
241 return String.format(FOREIGN_KEY_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
242 }
243
244 private Map<String, Collection<String>> loadViewNames(final Connection connection, final Collection<String> schemaNames, final Collection<String> tables) throws SQLException {
245 Map<String, Collection<String>> result = new LinkedHashMap<>(schemaNames.size(), 1F);
246 try (
247 PreparedStatement preparedStatement = connection.prepareStatement(getViewMetaDataSQL(schemaNames, tables));
248 ResultSet resultSet = preparedStatement.executeQuery()) {
249 while (resultSet.next()) {
250 String schemaName = resultSet.getString("table_schema");
251 Collection<String> viewMetaData = result.computeIfAbsent(schemaName, key -> new HashSet<>());
252 String tableName = resultSet.getString("table_name");
253 viewMetaData.add(tableName);
254 }
255 }
256 return result;
257 }
258
259 private String getViewMetaDataSQL(final Collection<String> schemaNames, final Collection<String> tables) {
260 return String.format(VIEW_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")),
261 tables.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
262 }
263
264 private Collection<TableMetaData> createTableMetaDataList(final Multimap<String, IndexMetaData> tableIndexMetaDataMap, final Multimap<String, ColumnMetaData> tableColumnMetaDataMap,
265 final Multimap<String, ConstraintMetaData> tableConstraintMetaDataMap, final Collection<String> viewNames) {
266 Collection<TableMetaData> result = new LinkedList<>();
267 for (String each : tableColumnMetaDataMap.keySet()) {
268 Collection<ColumnMetaData> columnMetaDataList = tableColumnMetaDataMap.get(each);
269 Collection<IndexMetaData> indexMetaDataList = tableIndexMetaDataMap.get(each);
270 Collection<ConstraintMetaData> constraintMetaDataList = tableConstraintMetaDataMap.get(each);
271 result.add(new TableMetaData(each, columnMetaDataList, indexMetaDataList, constraintMetaDataList, viewNames.contains(each) ? TableType.VIEW : TableType.TABLE));
272 }
273 return result;
274 }
275
276 @Override
277 public String getDatabaseType() {
278 return "PostgreSQL";
279 }
280 }