1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 package org.apache.shardingsphere.data.pipeline.postgresql.ddlgenerator;
19
20 import java.sql.Connection;
21 import java.util.Collection;
22 import java.util.HashMap;
23 import java.util.HashSet;
24 import java.util.LinkedHashMap;
25 import java.util.LinkedList;
26 import java.util.Map;
27 import java.util.Optional;
28 import java.util.Set;
29 import java.util.stream.Collectors;
30
31
32
33
34 public final class PostgreSQLConstraintsPropertiesAppender extends AbstractPostgreSQLDDLAdapter {
35
36 private static final Integer PG_CONSTRAINTS_INCLUDE_VERSION = 11;
37
38 public PostgreSQLConstraintsPropertiesAppender(final Connection connection, final int majorVersion, final int minorVersion) {
39 super(connection, majorVersion, minorVersion);
40 }
41
42
43
44
45
46
47 public void append(final Map<String, Object> context) {
48 loadPrimaryOrUniqueConstraint(context, "primary_key", "p");
49 loadPrimaryOrUniqueConstraint(context, "unique_constraint", "u");
50 context.put("foreign_key", fetchForeignKeys(context));
51 context.put("check_constraint", fetchCheckConstraints(context));
52 context.put("exclude_constraint", getExclusionConstraints(context));
53 }
54
55 private Collection<Map<String, Object>> fetchCheckConstraints(final Map<String, Object> context) {
56 Collection<Map<String, Object>> result = new LinkedList<>();
57 for (Map<String, Object> each : getCheckConstraints((Long) context.get("tid"))) {
58 if (!isPartitionAndConstraintInherited(each, context)) {
59 result.add(each);
60 }
61 }
62 return result;
63 }
64
65 private Collection<Map<String, Object>> fetchForeignKeys(final Map<String, Object> context) {
66 Collection<Map<String, Object>> result = new LinkedList<>();
67 for (Map<String, Object> each : getForeignKeys((Long) context.get("tid"))) {
68 if (!isPartitionAndConstraintInherited(each, context)) {
69 result.add(each);
70 }
71 }
72 return result;
73 }
74
75 private void loadPrimaryOrUniqueConstraint(final Map<String, Object> context, final String name, final String type) {
76 Collection<Map<String, Object>> constraintsProps = fetchConstraintsProperties(context, type);
77 fetchConstraintsColumns(constraintsProps);
78 context.put(name, constraintsProps.stream().filter(each -> !isPartitionAndConstraintInherited(each, context)).collect(Collectors.toList()));
79 }
80
81 private void fetchConstraintsColumns(final Collection<Map<String, Object>> constraintsProps) {
82 for (Map<String, Object> each : constraintsProps) {
83 Collection<Map<String, Object>> columns = new LinkedList<>();
84 for (Map<String, Object> col : fetchConstraintsCols(each)) {
85 Map<String, Object> column = new HashMap<>();
86 column.put("column", stripQuote((String) col.get("column")));
87 columns.add(column);
88 }
89 each.put("columns", columns);
90 appendConstraintsInclude(each);
91 }
92 }
93
94 private void appendConstraintsInclude(final Map<String, Object> constraintsProp) {
95 Map<String, Object> params = new LinkedHashMap<>();
96 params.put("cid", constraintsProp.get("oid"));
97 Collection<Object> includes = new LinkedList<>();
98 if (getMajorVersion() >= PG_CONSTRAINTS_INCLUDE_VERSION) {
99 for (Map<String, Object> each : executeByTemplate(params, "component/index_constraint/%s/get_constraint_include.ftl")) {
100 includes.add(each.get("colname"));
101 }
102 }
103 constraintsProp.put("include", includes);
104 }
105
106 private String stripQuote(final String column) {
107 String result = column;
108 if (column.startsWith("\"")) {
109 result = result.substring(1);
110 }
111 if (column.endsWith("\"")) {
112 result = result.substring(0, result.length() - 1);
113 }
114 return result;
115 }
116
117 private Collection<Map<String, Object>> fetchConstraintsCols(final Map<String, Object> constraintColProps) {
118 Map<String, Object> params = new HashMap<>();
119 params.put("cid", constraintColProps.get("oid"));
120 params.put("colcnt", constraintColProps.get("col_count"));
121 return executeByTemplate(params, "component/index_constraint/%s/get_costraint_cols.ftl");
122 }
123
124 private Collection<Map<String, Object>> fetchConstraintsProperties(final Map<String, Object> context, final String constraintType) {
125 Map<String, Object> params = new HashMap<>();
126 params.put("did", context.get("did"));
127 params.put("tid", context.get("tid"));
128 params.put("cid", context.get("cid"));
129 params.put("constraint_type", constraintType);
130 return executeByTemplate(params, "component/index_constraint/%s/properties.ftl");
131 }
132
133 private Collection<Map<String, Object>> getExclusionConstraints(final Map<String, Object> context) {
134 Map<String, Object> params = new HashMap<>();
135 params.put("tid", context.get("tid"));
136 params.put("did", context.get("did"));
137 Collection<Map<String, Object>> result = executeByTemplate(params, "component/exclusion_constraint/%s/properties.ftl");
138 for (Map<String, Object> each : result) {
139 getExclusionConstraintsColumns(each);
140 }
141 return result;
142 }
143
144 private void getExclusionConstraintsColumns(final Map<String, Object> exclusionConstraintsProps) {
145 Map<String, Object> params = new HashMap<>();
146 params.put("cid", exclusionConstraintsProps.get("oid"));
147 params.put("col_count", exclusionConstraintsProps.get("col_count"));
148 Collection<Map<String, Object>> columns = new LinkedList<>();
149 for (Map<String, Object> each : executeByTemplate(params, "component/exclusion_constraint/%s/get_constraint_cols.ftl")) {
150 boolean order = 0 == (((int) each.get("options")) & 1);
151 boolean nullsOrder = 0 != (((int) each.get("options")) & 2);
152 Map<String, Object> col = new HashMap<>();
153 col.put("column", strip((String) each.get("coldef")));
154 col.put("oper_class", each.get("opcname"));
155 col.put("order", order);
156 col.put("nulls_order", nullsOrder);
157 col.put("operator", each.get("oprname"));
158 col.put("col_type", each.get("datatype"));
159 col.put("is_exp", each.get("is_exp"));
160 columns.add(col);
161 }
162 exclusionConstraintsProps.put("columns", columns);
163 Map<String, Object> map = new HashMap<>();
164 map.put("cid", exclusionConstraintsProps.get("oid"));
165 Collection<String> include = new LinkedList<>();
166 if (getMajorVersion() >= PG_CONSTRAINTS_INCLUDE_VERSION) {
167 for (Map<String, Object> each : executeByTemplate(map, "exclusion_constraint/%s/get_constraint_include.ftl")) {
168 include.add(each.get("colname").toString());
169 }
170 }
171 exclusionConstraintsProps.put("include", include);
172 }
173
174 private Collection<Map<String, Object>> getForeignKeys(final Long tid) {
175 Map<String, Object> params = new HashMap<>();
176 params.put("tid", tid);
177 Collection<Map<String, Object>> result = executeByTemplate(params, "component/foreign_key/%s/properties.ftl");
178 for (Map<String, Object> each : result) {
179 Collection<Map<String, Object>> columns = new LinkedList<>();
180 Set<String> cols = new HashSet<>();
181 for (Map<String, Object> col : getForeignKeysCols(tid, each)) {
182 Map<String, Object> foreignKeysRef = new HashMap<>();
183 foreignKeysRef.put("local_column", col.get("conattname"));
184 foreignKeysRef.put("references", each.get("confrelid"));
185 foreignKeysRef.put("referenced", col.get("confattname"));
186 foreignKeysRef.put("references_table_name", each.get("refnsp") + "." + each.get("reftab"));
187 columns.add(foreignKeysRef);
188 cols.add((String) col.get("conattname"));
189 }
190 setRemoteName(each, columns);
191 Optional<String> coveringindex = searchCoveringIndex(tid, cols);
192 each.put("coveringindex", coveringindex.orElse(null));
193 each.put("autoindex", !coveringindex.isPresent());
194 each.put("hasindex", coveringindex.isPresent());
195 each.put("columns", columns);
196 }
197 return result;
198 }
199
200 private void setRemoteName(final Map<String, Object> foreignKey, final Collection<Map<String, Object>> columns) {
201 Map<String, Object> params = new HashMap<>();
202 params.put("tid", columns.iterator().next().get("references"));
203 Map<String, Object> parents = executeByTemplateForSingleRow(params, "component/foreign_key/%s/get_parent.ftl");
204 foreignKey.put("remote_schema", parents.get("schema"));
205 foreignKey.put("remote_table", parents.get("table"));
206 }
207
208 private Collection<Map<String, Object>> getForeignKeysCols(final Long tid, final Map<String, Object> foreignKeyProps) {
209 Map<String, Object> params = new HashMap<>();
210 params.put("tid", tid);
211 Collection<Map<String, Object>> keys = new LinkedList<>();
212 Map<String, Object> key = new HashMap<>();
213 key.put("confkey", foreignKeyProps.get("confkey"));
214 key.put("conkey", foreignKeyProps.get("conkey"));
215 keys.add(key);
216 params.put("keys", keys);
217 return executeByTemplate(params, "component/foreign_key/%s/get_constraint_cols.ftl");
218 }
219
220 private boolean isPartitionAndConstraintInherited(final Map<String, Object> constraint, final Map<String, Object> context) {
221 return context.containsKey("relispartition") && (boolean) context.get("relispartition") && constraint.containsKey("conislocal") && (boolean) constraint.get("conislocal");
222 }
223
224 private Optional<String> searchCoveringIndex(final Long tid, final Set<String> cols) {
225 Map<String, Object> params = new HashMap<>();
226 params.put("tid", tid);
227 for (Map<String, Object> each : executeByTemplate(params, "component/foreign_key/%s/get_constraints.ftl")) {
228 Map<String, Object> map = new HashMap<>();
229 map.put("cid", each.get("oid"));
230 map.put("colcnt", each.get("col_count"));
231 Collection<Map<String, Object>> rows = executeByTemplate(map, "component/foreign_key/%s/get_cols.ftl");
232 Set<String> indexCols = new HashSet<>();
233 for (Map<String, Object> row : rows) {
234 indexCols.add(strip(row.get("column").toString()));
235 }
236 if (isSame(indexCols, cols)) {
237 return Optional.of((String) each.get("idxname"));
238 }
239 }
240 return Optional.empty();
241 }
242
243 private boolean isSame(final Set<String> indexCols, final Set<String> cols) {
244 Set<String> copyIndexCols = new HashSet<>(indexCols);
245 Set<String> copyCols = new HashSet<>(cols);
246 copyIndexCols.removeAll(copyCols);
247 if (copyIndexCols.isEmpty()) {
248 cols.removeAll(indexCols);
249 return cols.isEmpty();
250 }
251 return false;
252 }
253
254 private String strip(final String column) {
255 if (column.startsWith("\"")) {
256 return column.substring(1);
257 }
258 if (column.endsWith("\"")) {
259 return column.substring(0, column.length() - 1);
260 }
261 return column;
262 }
263
264 private Collection<Map<String, Object>> getCheckConstraints(final Long tid) {
265 Map<String, Object> params = new HashMap<>();
266 params.put("tid", tid);
267 return executeByTemplate(params, "component/check_constraint/%s/properties.ftl");
268 }
269 }