1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 package org.apache.shardingsphere.infra.database.h2.metadata.data.loader;
19
20 import org.apache.shardingsphere.infra.database.core.metadata.data.loader.DialectMetaDataLoader;
21 import org.apache.shardingsphere.infra.database.core.metadata.data.loader.MetaDataLoaderMaterial;
22 import org.apache.shardingsphere.infra.database.core.metadata.data.model.ColumnMetaData;
23 import org.apache.shardingsphere.infra.database.core.metadata.data.model.IndexMetaData;
24 import org.apache.shardingsphere.infra.database.core.metadata.data.model.SchemaMetaData;
25 import org.apache.shardingsphere.infra.database.core.metadata.data.model.TableMetaData;
26 import org.apache.shardingsphere.infra.database.core.metadata.database.datatype.DataTypeLoader;
27
28 import java.sql.Connection;
29 import java.sql.PreparedStatement;
30 import java.sql.ResultSet;
31 import java.sql.SQLException;
32 import java.util.Collection;
33 import java.util.Collections;
34 import java.util.HashMap;
35 import java.util.LinkedList;
36 import java.util.Map;
37 import java.util.Map.Entry;
38 import java.util.stream.Collectors;
39
40
41
42
43 public final class H2MetaDataLoader implements DialectMetaDataLoader {
44
45 private static final String TABLE_META_DATA_NO_ORDER = "SELECT TABLE_CATALOG, TABLE_NAME, COLUMN_NAME, DATA_TYPE, ORDINAL_POSITION, COALESCE(IS_VISIBLE, FALSE) IS_VISIBLE, IS_NULLABLE"
46 + " FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_CATALOG=? AND TABLE_SCHEMA=?";
47
48 private static final String ORDER_BY_ORDINAL_POSITION = " ORDER BY ORDINAL_POSITION";
49
50 private static final String TABLE_META_DATA_SQL = TABLE_META_DATA_NO_ORDER + ORDER_BY_ORDINAL_POSITION;
51
52 private static final String TABLE_META_DATA_SQL_IN_TABLES = TABLE_META_DATA_NO_ORDER + " AND UPPER(TABLE_NAME) IN (%s)" + ORDER_BY_ORDINAL_POSITION;
53
54 private static final String INDEX_META_DATA_SQL = "SELECT TABLE_CATALOG, TABLE_NAME, INDEX_NAME, INDEX_TYPE_NAME FROM INFORMATION_SCHEMA.INDEXES"
55 + " WHERE TABLE_CATALOG=? AND TABLE_SCHEMA=? AND UPPER(TABLE_NAME) IN (%s)";
56
57 private static final String PRIMARY_KEY_META_DATA_SQL = "SELECT TABLE_NAME, INDEX_NAME FROM INFORMATION_SCHEMA.INDEXES WHERE TABLE_CATALOG=? AND TABLE_SCHEMA=?"
58 + " AND INDEX_TYPE_NAME = 'PRIMARY KEY'";
59
60 private static final String PRIMARY_KEY_META_DATA_SQL_IN_TABLES = PRIMARY_KEY_META_DATA_SQL + " AND UPPER(TABLE_NAME) IN (%s)";
61
62 private static final String GENERATED_INFO_SQL = "SELECT C.TABLE_NAME TABLE_NAME, C.COLUMN_NAME COLUMN_NAME, COALESCE(I.IS_GENERATED, FALSE) IS_GENERATED FROM INFORMATION_SCHEMA.COLUMNS C"
63 + " RIGHT JOIN INFORMATION_SCHEMA.INDEXES I ON C.TABLE_NAME=I.TABLE_NAME WHERE C.TABLE_CATALOG=? AND C.TABLE_SCHEMA=?";
64
65 private static final String GENERATED_INFO_SQL_IN_TABLES = GENERATED_INFO_SQL + " AND UPPER(C.TABLE_NAME) IN (%s)";
66
67 @Override
68 public Collection<SchemaMetaData> load(final MetaDataLoaderMaterial material) throws SQLException {
69 Collection<TableMetaData> tableMetaDataList = new LinkedList<>();
70 try (Connection connection = material.getDataSource().getConnection()) {
71 Map<String, Collection<ColumnMetaData>> columnMetaDataMap = loadColumnMetaDataMap(connection, material.getActualTableNames());
72 Map<String, Collection<IndexMetaData>> indexMetaDataMap = columnMetaDataMap.isEmpty() ? Collections.emptyMap() : loadIndexMetaData(connection, columnMetaDataMap.keySet());
73 for (Entry<String, Collection<ColumnMetaData>> entry : columnMetaDataMap.entrySet()) {
74 Collection<IndexMetaData> indexMetaDataList = indexMetaDataMap.getOrDefault(entry.getKey(), Collections.emptyList());
75 tableMetaDataList.add(new TableMetaData(entry.getKey(), entry.getValue(), indexMetaDataList, Collections.emptyList()));
76 }
77 }
78 return Collections.singleton(new SchemaMetaData(material.getDefaultSchemaName(), tableMetaDataList));
79 }
80
81 private Map<String, Collection<ColumnMetaData>> loadColumnMetaDataMap(final Connection connection, final Collection<String> tables) throws SQLException {
82 Map<String, Collection<ColumnMetaData>> result = new HashMap<>();
83 try (PreparedStatement preparedStatement = connection.prepareStatement(getTableMetaDataSQL(tables))) {
84 Map<String, Integer> dataTypes = new DataTypeLoader().load(connection.getMetaData(), getType());
85 Map<String, Collection<String>> tablePrimaryKeys = loadTablePrimaryKeys(connection, tables);
86 Map<String, Map<String, Boolean>> tableGenerated = loadTableGenerated(connection, tables);
87 preparedStatement.setString(1, connection.getCatalog());
88 preparedStatement.setString(2, "PUBLIC");
89 try (ResultSet resultSet = preparedStatement.executeQuery()) {
90 while (resultSet.next()) {
91 String tableName = resultSet.getString("TABLE_NAME");
92 ColumnMetaData columnMetaData = loadColumnMetaData(dataTypes, resultSet, tablePrimaryKeys.getOrDefault(tableName, Collections.emptyList()),
93 tableGenerated.getOrDefault(tableName, new HashMap<>()));
94 if (!result.containsKey(tableName)) {
95 result.put(tableName, new LinkedList<>());
96 }
97 result.get(tableName).add(columnMetaData);
98 }
99 }
100 }
101 return result;
102 }
103
104 private ColumnMetaData loadColumnMetaData(final Map<String, Integer> dataTypeMap, final ResultSet resultSet, final Collection<String> primaryKeys,
105 final Map<String, Boolean> tableGenerated) throws SQLException {
106 String columnName = resultSet.getString("COLUMN_NAME");
107 String dataType = resultSet.getString("DATA_TYPE");
108 boolean primaryKey = primaryKeys.contains(columnName);
109 boolean generated = tableGenerated.getOrDefault(columnName, Boolean.FALSE);
110 boolean isVisible = resultSet.getBoolean("IS_VISIBLE");
111 boolean isNullable = "YES".equals(resultSet.getString("IS_NULLABLE"));
112 return new ColumnMetaData(columnName, dataTypeMap.get(dataType), primaryKey, generated, false, isVisible, false, isNullable);
113 }
114
115 private String getTableMetaDataSQL(final Collection<String> tables) {
116 return tables.isEmpty() ? TABLE_META_DATA_SQL
117 : String.format(TABLE_META_DATA_SQL_IN_TABLES, tables.stream().map(each -> String.format("'%s'", each).toUpperCase()).collect(Collectors.joining(",")));
118 }
119
120 private Map<String, Collection<IndexMetaData>> loadIndexMetaData(final Connection connection, final Collection<String> tableNames) throws SQLException {
121 Map<String, Collection<IndexMetaData>> result = new HashMap<>();
122 try (PreparedStatement preparedStatement = connection.prepareStatement(getIndexMetaDataSQL(tableNames))) {
123 preparedStatement.setString(1, connection.getCatalog());
124 preparedStatement.setString(2, "PUBLIC");
125 try (ResultSet resultSet = preparedStatement.executeQuery()) {
126 while (resultSet.next()) {
127 String indexName = resultSet.getString("INDEX_NAME");
128 String tableName = resultSet.getString("TABLE_NAME");
129 boolean uniqueIndex = "UNIQUE INDEX".equals(resultSet.getString("INDEX_TYPE_NAME"));
130 if (!result.containsKey(tableName)) {
131 result.put(tableName, new LinkedList<>());
132 }
133 IndexMetaData indexMetaData = new IndexMetaData(indexName);
134 indexMetaData.setUnique(uniqueIndex);
135 result.get(tableName).add(indexMetaData);
136
137 }
138 }
139 }
140 return result;
141 }
142
143 private String getIndexMetaDataSQL(final Collection<String> tableNames) {
144 return String.format(INDEX_META_DATA_SQL, tableNames.stream().map(each -> String.format("'%s'", each).toUpperCase()).collect(Collectors.joining(",")));
145 }
146
147 private String getPrimaryKeyMetaDataSQL(final Collection<String> tables) {
148 return tables.isEmpty() ? PRIMARY_KEY_META_DATA_SQL
149 : String.format(PRIMARY_KEY_META_DATA_SQL_IN_TABLES, tables.stream().map(each -> String.format("'%s'", each).toUpperCase()).collect(Collectors.joining(",")));
150 }
151
152 private Map<String, Collection<String>> loadTablePrimaryKeys(final Connection connection, final Collection<String> tableNames) throws SQLException {
153 Map<String, Collection<String>> result = new HashMap<>();
154 try (PreparedStatement preparedStatement = connection.prepareStatement(getPrimaryKeyMetaDataSQL(tableNames))) {
155 preparedStatement.setString(1, connection.getCatalog());
156 preparedStatement.setString(2, "PUBLIC");
157 try (ResultSet resultSet = preparedStatement.executeQuery()) {
158 while (resultSet.next()) {
159 String indexName = resultSet.getString("INDEX_NAME");
160 String tableName = resultSet.getString("TABLE_NAME");
161 result.computeIfAbsent(tableName, k -> new LinkedList<>()).add(indexName);
162 }
163 }
164 }
165 return result;
166 }
167
168 private String getGeneratedInfoSQL(final Collection<String> tables) {
169 return tables.isEmpty() ? GENERATED_INFO_SQL
170 : String.format(GENERATED_INFO_SQL_IN_TABLES, tables.stream().map(each -> String.format("'%s'", each).toUpperCase()).collect(Collectors.joining(",")));
171 }
172
173 private Map<String, Map<String, Boolean>> loadTableGenerated(final Connection connection, final Collection<String> tableNames) throws SQLException {
174 Map<String, Map<String, Boolean>> result = new HashMap<>();
175 try (PreparedStatement preparedStatement = connection.prepareStatement(getGeneratedInfoSQL(tableNames))) {
176 preparedStatement.setString(1, connection.getCatalog());
177 preparedStatement.setString(2, "PUBLIC");
178 try (ResultSet resultSet = preparedStatement.executeQuery()) {
179 while (resultSet.next()) {
180 String columnName = resultSet.getString("COLUMN_NAME");
181 String tableName = resultSet.getString("TABLE_NAME");
182 boolean generated = resultSet.getBoolean("IS_GENERATED");
183 result.computeIfAbsent(tableName, k -> new HashMap<>()).put(columnName, generated);
184 }
185 }
186 }
187 return result;
188 }
189
190 @Override
191 public String getDatabaseType() {
192 return "H2";
193 }
194 }