1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 package org.apache.shardingsphere.infra.database.opengauss.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.IndexMetaData;
27 import org.apache.shardingsphere.infra.database.core.metadata.data.model.SchemaMetaData;
28 import org.apache.shardingsphere.infra.database.core.metadata.data.model.TableMetaData;
29 import org.apache.shardingsphere.infra.database.core.metadata.database.datatype.DataTypeRegistry;
30
31 import java.sql.Connection;
32 import java.sql.PreparedStatement;
33 import java.sql.ResultSet;
34 import java.sql.SQLException;
35 import java.sql.Types;
36 import java.util.Collection;
37 import java.util.Collections;
38 import java.util.HashSet;
39 import java.util.LinkedHashMap;
40 import java.util.LinkedList;
41 import java.util.Map;
42 import java.util.Optional;
43 import java.util.stream.Collectors;
44
45
46
47
48 public final class OpenGaussMetaDataLoader implements DialectMetaDataLoader {
49
50 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"
51 + " FROM information_schema.columns WHERE table_schema IN (%s)";
52
53 private static final String TABLE_META_DATA_SQL_WITHOUT_TABLES = BASIC_TABLE_META_DATA_SQL + " ORDER BY ordinal_position";
54
55 private static final String TABLE_META_DATA_SQL_WITH_TABLES = BASIC_TABLE_META_DATA_SQL + " AND table_name IN (%s) ORDER BY ordinal_position";
56
57 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"
58 + " 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"
59 + " WHERE tc.constraint_type = 'PRIMARY KEY' AND kc.ordinal_position IS NOT NULL AND kc.table_schema IN (%s)";
60
61 private static final String BASIC_INDEX_META_DATA_SQL = "SELECT tablename, indexname, schemaname FROM pg_indexes WHERE schemaname IN (%s)";
62
63 private static final String ADVANCE_INDEX_META_DATA_SQL =
64 "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"
65 + " 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"
66 + " 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)";
67
68 @Override
69 public Collection<SchemaMetaData> load(final MetaDataLoaderMaterial material) throws SQLException {
70 try (Connection connection = material.getDataSource().getConnection()) {
71 Collection<String> schemaNames = new SchemaMetaDataLoader(getType()).loadSchemaNames(connection);
72 Map<String, Multimap<String, IndexMetaData>> schemaIndexMetaDataMap = loadIndexMetaDataMap(connection, schemaNames);
73 Map<String, Multimap<String, ColumnMetaData>> schemaColumnMetaDataMap = loadColumnMetaDataMap(connection, material.getActualTableNames(), schemaNames);
74 Collection<SchemaMetaData> result = new LinkedList<>();
75 for (String each : schemaNames) {
76 Multimap<String, IndexMetaData> tableIndexMetaDataMap = schemaIndexMetaDataMap.getOrDefault(each, LinkedHashMultimap.create());
77 Multimap<String, ColumnMetaData> tableColumnMetaDataMap = schemaColumnMetaDataMap.getOrDefault(each, LinkedHashMultimap.create());
78 result.add(new SchemaMetaData(each, createTableMetaDataList(tableIndexMetaDataMap, tableColumnMetaDataMap)));
79 }
80 return result;
81 }
82 }
83
84 private Map<String, Multimap<String, IndexMetaData>> loadIndexMetaDataMap(final Connection connection, final Collection<String> schemaNames) throws SQLException {
85 Map<String, Multimap<String, IndexMetaData>> result = new LinkedHashMap<>();
86 try (
87 PreparedStatement preparedStatement = connection.prepareStatement(getIndexMetaDataSQL(schemaNames));
88 ResultSet resultSet = preparedStatement.executeQuery()) {
89 while (resultSet.next()) {
90 String schemaName = resultSet.getString("schemaname");
91 String tableName = resultSet.getString("tablename");
92 String indexName = resultSet.getString("indexname");
93 Multimap<String, IndexMetaData> indexMetaDataMap = result.computeIfAbsent(schemaName, key -> LinkedHashMultimap.create());
94 indexMetaDataMap.put(tableName, new IndexMetaData(indexName));
95 }
96 }
97 try (
98 PreparedStatement preparedStatement = connection.prepareStatement(getAdvanceIndexMetaDataSQL(schemaNames));
99 ResultSet resultSet = preparedStatement.executeQuery()) {
100 while (resultSet.next()) {
101 String schemaName = resultSet.getString("index_schema");
102 String tableName = resultSet.getString("table_name");
103 String columnName = resultSet.getString("column_name");
104 String indexName = resultSet.getString("index_name");
105 boolean isUnique = resultSet.getBoolean("is_unique");
106 Collection<IndexMetaData> indexMetaDatas = result.getOrDefault(schemaName, LinkedHashMultimap.create()).get(tableName);
107 if (indexMetaDatas.isEmpty()) {
108 continue;
109 }
110 Optional<IndexMetaData> indexMetaData = indexMetaDatas.stream().filter(each -> each.getName().equals(indexName)).findFirst();
111 if (indexMetaData.isPresent()) {
112 indexMetaData.get().setUnique(isUnique);
113 indexMetaData.get().getColumns().add(columnName);
114 }
115 }
116 }
117 return result;
118 }
119
120 private String getIndexMetaDataSQL(final Collection<String> schemaNames) {
121 return String.format(BASIC_INDEX_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
122 }
123
124 private String getAdvanceIndexMetaDataSQL(final Collection<String> schemaNames) {
125 return String.format(ADVANCE_INDEX_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
126 }
127
128 private Map<String, Multimap<String, ColumnMetaData>> loadColumnMetaDataMap(final Connection connection, final Collection<String> tables,
129 final Collection<String> schemaNames) throws SQLException {
130 Map<String, Multimap<String, ColumnMetaData>> result = new LinkedHashMap<>();
131 try (
132 PreparedStatement preparedStatement = connection.prepareStatement(getColumnMetaDataSQL(schemaNames, tables));
133 ResultSet resultSet = preparedStatement.executeQuery()) {
134 Collection<String> primaryKeys = loadPrimaryKeys(connection, schemaNames);
135 while (resultSet.next()) {
136 String tableName = resultSet.getString("table_name");
137 String schemaName = resultSet.getString("table_schema");
138 Multimap<String, ColumnMetaData> columnMetaDataMap = result.computeIfAbsent(schemaName, key -> LinkedHashMultimap.create());
139 columnMetaDataMap.put(tableName, loadColumnMetaData(primaryKeys, resultSet));
140 }
141 }
142 return result;
143 }
144
145 private String getColumnMetaDataSQL(final Collection<String> schemaNames, final Collection<String> tables) {
146 String schemaNameParam = schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(","));
147 return tables.isEmpty() ? String.format(TABLE_META_DATA_SQL_WITHOUT_TABLES, schemaNameParam)
148 : String.format(TABLE_META_DATA_SQL_WITH_TABLES, schemaNameParam, tables.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
149 }
150
151 private Collection<String> loadPrimaryKeys(final Connection connection, final Collection<String> schemaNames) throws SQLException {
152 Collection<String> result = new HashSet<>();
153 try (
154 PreparedStatement preparedStatement = connection.prepareStatement(getPrimaryKeyMetaDataSQL(schemaNames));
155 ResultSet resultSet = preparedStatement.executeQuery()) {
156 while (resultSet.next()) {
157 String schemaName = resultSet.getString("table_schema");
158 String tableName = resultSet.getString("table_name");
159 String columnName = resultSet.getString("column_name");
160 result.add(schemaName + "," + tableName + "," + columnName);
161 }
162 }
163 return result;
164 }
165
166 private String getPrimaryKeyMetaDataSQL(final Collection<String> schemaNames) {
167 return String.format(PRIMARY_KEY_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(",")));
168 }
169
170 private ColumnMetaData loadColumnMetaData(final Collection<String> primaryKeys, final ResultSet resultSet) throws SQLException {
171 String schemaName = resultSet.getString("table_schema");
172 String tableName = resultSet.getString("table_name");
173 String columnName = resultSet.getString("column_name");
174 String dataType = resultSet.getString("udt_name");
175 boolean isPrimaryKey = primaryKeys.contains(schemaName + "," + tableName + "," + columnName);
176 String columnDefault = resultSet.getString("column_default");
177 boolean generated = null != columnDefault && columnDefault.startsWith("nextval(");
178
179 boolean caseSensitive = true;
180 boolean isNullable = "YES".equals(resultSet.getString("is_nullable"));
181 return new ColumnMetaData(columnName, DataTypeRegistry.getDataType(getDatabaseType(), dataType).orElse(Types.OTHER), isPrimaryKey, generated, caseSensitive, true, false, isNullable);
182 }
183
184 private Collection<TableMetaData> createTableMetaDataList(final Multimap<String, IndexMetaData> tableIndexMetaDataMap, final Multimap<String, ColumnMetaData> tableColumnMetaDataMap) {
185 Collection<TableMetaData> result = new LinkedList<>();
186 for (String each : tableColumnMetaDataMap.keySet()) {
187 Collection<ColumnMetaData> columnMetaDataList = tableColumnMetaDataMap.get(each);
188 Collection<IndexMetaData> indexMetaDataList = tableIndexMetaDataMap.get(each);
189 result.add(new TableMetaData(each, columnMetaDataList, indexMetaDataList, Collections.emptyList()));
190 }
191 return result;
192 }
193
194 @Override
195 public String getDatabaseType() {
196 return "openGauss";
197 }
198 }