1 | /* |
2 | * Copyright 2004-2014 H2 Group. Multiple-Licensed under the MPL 2.0, |
3 | * and the EPL 1.0 (http://h2database.com/html/license.html). |
4 | * Initial Developer: H2 Group |
5 | */ |
6 | package org.h2.command.dml; |
7 | |
8 | import java.util.ArrayList; |
9 | import java.util.HashMap; |
10 | |
11 | import org.h2.api.ErrorCode; |
12 | import org.h2.api.Trigger; |
13 | import org.h2.command.Command; |
14 | import org.h2.command.CommandInterface; |
15 | import org.h2.command.Prepared; |
16 | import org.h2.engine.Right; |
17 | import org.h2.engine.Session; |
18 | import org.h2.engine.UndoLogRecord; |
19 | import org.h2.expression.Comparison; |
20 | import org.h2.expression.ConditionAndOr; |
21 | import org.h2.expression.Expression; |
22 | import org.h2.expression.ExpressionColumn; |
23 | import org.h2.expression.Parameter; |
24 | import org.h2.index.Index; |
25 | import org.h2.message.DbException; |
26 | import org.h2.result.ResultInterface; |
27 | import org.h2.result.ResultTarget; |
28 | import org.h2.result.Row; |
29 | import org.h2.table.Column; |
30 | import org.h2.table.Table; |
31 | import org.h2.util.New; |
32 | import org.h2.util.StatementBuilder; |
33 | import org.h2.value.Value; |
34 | import org.h2.value.ValueNull; |
35 | |
36 | /** |
37 | * This class represents the statement |
38 | * INSERT |
39 | */ |
40 | public class Insert extends Prepared implements ResultTarget { |
41 | |
42 | private Table table; |
43 | private Column[] columns; |
44 | private final ArrayList<Expression[]> list = New.arrayList(); |
45 | private Query query; |
46 | private boolean sortedInsertMode; |
47 | private int rowNumber; |
48 | private boolean insertFromSelect; |
49 | |
50 | /** |
51 | * For MySQL-style INSERT ... ON DUPLICATE KEY UPDATE .... |
52 | */ |
53 | private HashMap<Column, Expression> duplicateKeyAssignmentMap; |
54 | |
55 | public Insert(Session session) { |
56 | super(session); |
57 | } |
58 | |
59 | @Override |
60 | public void setCommand(Command command) { |
61 | super.setCommand(command); |
62 | if (query != null) { |
63 | query.setCommand(command); |
64 | } |
65 | } |
66 | |
67 | public void setTable(Table table) { |
68 | this.table = table; |
69 | } |
70 | |
71 | public void setColumns(Column[] columns) { |
72 | this.columns = columns; |
73 | } |
74 | |
75 | public void setQuery(Query query) { |
76 | this.query = query; |
77 | } |
78 | |
79 | /** |
80 | * Keep a collection of the columns to pass to update if a duplicate key |
81 | * happens, for MySQL-style INSERT ... ON DUPLICATE KEY UPDATE .... |
82 | * |
83 | * @param column the column |
84 | * @param expression the expression |
85 | */ |
86 | public void addAssignmentForDuplicate(Column column, Expression expression) { |
87 | if (duplicateKeyAssignmentMap == null) { |
88 | duplicateKeyAssignmentMap = New.hashMap(); |
89 | } |
90 | if (duplicateKeyAssignmentMap.containsKey(column)) { |
91 | throw DbException.get(ErrorCode.DUPLICATE_COLUMN_NAME_1, |
92 | column.getName()); |
93 | } |
94 | duplicateKeyAssignmentMap.put(column, expression); |
95 | } |
96 | |
97 | /** |
98 | * Add a row to this merge statement. |
99 | * |
100 | * @param expr the list of values |
101 | */ |
102 | public void addRow(Expression[] expr) { |
103 | list.add(expr); |
104 | } |
105 | |
106 | @Override |
107 | public int update() { |
108 | Index index = null; |
109 | if (sortedInsertMode) { |
110 | index = table.getScanIndex(session); |
111 | index.setSortedInsertMode(true); |
112 | } |
113 | try { |
114 | return insertRows(); |
115 | } finally { |
116 | if (index != null) { |
117 | index.setSortedInsertMode(false); |
118 | } |
119 | } |
120 | } |
121 | |
122 | private int insertRows() { |
123 | session.getUser().checkRight(table, Right.INSERT); |
124 | setCurrentRowNumber(0); |
125 | table.fire(session, Trigger.INSERT, true); |
126 | rowNumber = 0; |
127 | int listSize = list.size(); |
128 | if (listSize > 0) { |
129 | int columnLen = columns.length; |
130 | for (int x = 0; x < listSize; x++) { |
131 | session.startStatementWithinTransaction(); |
132 | Row newRow = table.getTemplateRow(); |
133 | Expression[] expr = list.get(x); |
134 | setCurrentRowNumber(x + 1); |
135 | for (int i = 0; i < columnLen; i++) { |
136 | Column c = columns[i]; |
137 | int index = c.getColumnId(); |
138 | Expression e = expr[i]; |
139 | if (e != null) { |
140 | // e can be null (DEFAULT) |
141 | e = e.optimize(session); |
142 | try { |
143 | Value v = c.convert(e.getValue(session)); |
144 | newRow.setValue(index, v); |
145 | } catch (DbException ex) { |
146 | throw setRow(ex, x, getSQL(expr)); |
147 | } |
148 | } |
149 | } |
150 | rowNumber++; |
151 | table.validateConvertUpdateSequence(session, newRow); |
152 | boolean done = table.fireBeforeRow(session, null, newRow); |
153 | if (!done) { |
154 | table.lock(session, true, false); |
155 | try { |
156 | table.addRow(session, newRow); |
157 | } catch (DbException de) { |
158 | handleOnDuplicate(de); |
159 | } |
160 | session.log(table, UndoLogRecord.INSERT, newRow); |
161 | table.fireAfterRow(session, null, newRow, false); |
162 | } |
163 | } |
164 | } else { |
165 | table.lock(session, true, false); |
166 | if (insertFromSelect) { |
167 | query.query(0, this); |
168 | } else { |
169 | ResultInterface rows = query.query(0); |
170 | while (rows.next()) { |
171 | Value[] r = rows.currentRow(); |
172 | addRow(r); |
173 | } |
174 | rows.close(); |
175 | } |
176 | } |
177 | table.fire(session, Trigger.INSERT, false); |
178 | return rowNumber; |
179 | } |
180 | |
181 | @Override |
182 | public void addRow(Value[] values) { |
183 | Row newRow = table.getTemplateRow(); |
184 | setCurrentRowNumber(++rowNumber); |
185 | for (int j = 0, len = columns.length; j < len; j++) { |
186 | Column c = columns[j]; |
187 | int index = c.getColumnId(); |
188 | try { |
189 | Value v = c.convert(values[j]); |
190 | newRow.setValue(index, v); |
191 | } catch (DbException ex) { |
192 | throw setRow(ex, rowNumber, getSQL(values)); |
193 | } |
194 | } |
195 | table.validateConvertUpdateSequence(session, newRow); |
196 | boolean done = table.fireBeforeRow(session, null, newRow); |
197 | if (!done) { |
198 | table.addRow(session, newRow); |
199 | session.log(table, UndoLogRecord.INSERT, newRow); |
200 | table.fireAfterRow(session, null, newRow, false); |
201 | } |
202 | } |
203 | |
204 | @Override |
205 | public int getRowCount() { |
206 | return rowNumber; |
207 | } |
208 | |
209 | @Override |
210 | public String getPlanSQL() { |
211 | StatementBuilder buff = new StatementBuilder("INSERT INTO "); |
212 | buff.append(table.getSQL()).append('('); |
213 | for (Column c : columns) { |
214 | buff.appendExceptFirst(", "); |
215 | buff.append(c.getSQL()); |
216 | } |
217 | buff.append(")\n"); |
218 | if (insertFromSelect) { |
219 | buff.append("DIRECT "); |
220 | } |
221 | if (sortedInsertMode) { |
222 | buff.append("SORTED "); |
223 | } |
224 | if (list.size() > 0) { |
225 | buff.append("VALUES "); |
226 | int row = 0; |
227 | if (list.size() > 1) { |
228 | buff.append('\n'); |
229 | } |
230 | for (Expression[] expr : list) { |
231 | if (row++ > 0) { |
232 | buff.append(",\n"); |
233 | } |
234 | buff.append('('); |
235 | buff.resetCount(); |
236 | for (Expression e : expr) { |
237 | buff.appendExceptFirst(", "); |
238 | if (e == null) { |
239 | buff.append("DEFAULT"); |
240 | } else { |
241 | buff.append(e.getSQL()); |
242 | } |
243 | } |
244 | buff.append(')'); |
245 | } |
246 | } else { |
247 | buff.append(query.getPlanSQL()); |
248 | } |
249 | return buff.toString(); |
250 | } |
251 | |
252 | @Override |
253 | public void prepare() { |
254 | if (columns == null) { |
255 | if (list.size() > 0 && list.get(0).length == 0) { |
256 | // special case where table is used as a sequence |
257 | columns = new Column[0]; |
258 | } else { |
259 | columns = table.getColumns(); |
260 | } |
261 | } |
262 | if (list.size() > 0) { |
263 | for (Expression[] expr : list) { |
264 | if (expr.length != columns.length) { |
265 | throw DbException.get(ErrorCode.COLUMN_COUNT_DOES_NOT_MATCH); |
266 | } |
267 | for (int i = 0, len = expr.length; i < len; i++) { |
268 | Expression e = expr[i]; |
269 | if (e != null) { |
270 | e = e.optimize(session); |
271 | if (e instanceof Parameter) { |
272 | Parameter p = (Parameter) e; |
273 | p.setColumn(columns[i]); |
274 | } |
275 | expr[i] = e; |
276 | } |
277 | } |
278 | } |
279 | } else { |
280 | query.prepare(); |
281 | if (query.getColumnCount() != columns.length) { |
282 | throw DbException.get(ErrorCode.COLUMN_COUNT_DOES_NOT_MATCH); |
283 | } |
284 | } |
285 | } |
286 | |
287 | @Override |
288 | public boolean isTransactional() { |
289 | return true; |
290 | } |
291 | |
292 | @Override |
293 | public ResultInterface queryMeta() { |
294 | return null; |
295 | } |
296 | |
297 | public void setSortedInsertMode(boolean sortedInsertMode) { |
298 | this.sortedInsertMode = sortedInsertMode; |
299 | } |
300 | |
301 | @Override |
302 | public int getType() { |
303 | return CommandInterface.INSERT; |
304 | } |
305 | |
306 | public void setInsertFromSelect(boolean value) { |
307 | this.insertFromSelect = value; |
308 | } |
309 | |
310 | @Override |
311 | public boolean isCacheable() { |
312 | return duplicateKeyAssignmentMap == null || |
313 | duplicateKeyAssignmentMap.isEmpty(); |
314 | } |
315 | |
316 | private void handleOnDuplicate(DbException de) { |
317 | if (de.getErrorCode() != ErrorCode.DUPLICATE_KEY_1) { |
318 | throw de; |
319 | } |
320 | if (duplicateKeyAssignmentMap == null || |
321 | duplicateKeyAssignmentMap.isEmpty()) { |
322 | throw de; |
323 | } |
324 | |
325 | ArrayList<String> variableNames = new ArrayList<String>( |
326 | duplicateKeyAssignmentMap.size()); |
327 | for (int i = 0; i < columns.length; i++) { |
328 | String key = session.getCurrentSchemaName() + "." + |
329 | table.getName() + "." + columns[i].getName(); |
330 | variableNames.add(key); |
331 | session.setVariable(key, |
332 | list.get(getCurrentRowNumber() - 1)[i].getValue(session)); |
333 | } |
334 | |
335 | StatementBuilder buff = new StatementBuilder("UPDATE "); |
336 | buff.append(table.getSQL()).append(" SET "); |
337 | for (Column column : duplicateKeyAssignmentMap.keySet()) { |
338 | buff.appendExceptFirst(", "); |
339 | Expression ex = duplicateKeyAssignmentMap.get(column); |
340 | buff.append(column.getSQL()).append("=").append(ex.getSQL()); |
341 | } |
342 | buff.append(" WHERE "); |
343 | Index foundIndex = searchForUpdateIndex(); |
344 | if (foundIndex == null) { |
345 | throw DbException.getUnsupportedException( |
346 | "Unable to apply ON DUPLICATE KEY UPDATE, no index found!"); |
347 | } |
348 | buff.append(prepareUpdateCondition(foundIndex).getSQL()); |
349 | String sql = buff.toString(); |
350 | Prepared command = session.prepare(sql); |
351 | for (Parameter param : command.getParameters()) { |
352 | Parameter insertParam = parameters.get(param.getIndex()); |
353 | param.setValue(insertParam.getValue(session)); |
354 | } |
355 | command.update(); |
356 | for (String variableName : variableNames) { |
357 | session.setVariable(variableName, ValueNull.INSTANCE); |
358 | } |
359 | } |
360 | |
361 | private Index searchForUpdateIndex() { |
362 | Index foundIndex = null; |
363 | for (Index index : table.getIndexes()) { |
364 | if (index.getIndexType().isPrimaryKey() || index.getIndexType().isUnique()) { |
365 | for (Column indexColumn : index.getColumns()) { |
366 | for (Column insertColumn : columns) { |
367 | if (indexColumn.getName().equals(insertColumn.getName())) { |
368 | foundIndex = index; |
369 | break; |
370 | } |
371 | foundIndex = null; |
372 | } |
373 | if (foundIndex == null) { |
374 | break; |
375 | } |
376 | } |
377 | if (foundIndex != null) { |
378 | break; |
379 | } |
380 | } |
381 | } |
382 | return foundIndex; |
383 | } |
384 | |
385 | private Expression prepareUpdateCondition(Index foundIndex) { |
386 | Expression condition = null; |
387 | for (Column column : foundIndex.getColumns()) { |
388 | ExpressionColumn expr = new ExpressionColumn(session.getDatabase(), |
389 | session.getCurrentSchemaName(), table.getName(), column.getName()); |
390 | for (int i = 0; i < columns.length; i++) { |
391 | if (expr.getColumnName().equals(columns[i].getName())) { |
392 | if (condition == null) { |
393 | condition = new Comparison(session, Comparison.EQUAL, |
394 | expr, list.get(getCurrentRowNumber() - 1)[i++]); |
395 | } else { |
396 | condition = new ConditionAndOr(ConditionAndOr.AND, condition, |
397 | new Comparison(session, Comparison.EQUAL, |
398 | expr, list.get(0)[i++])); |
399 | } |
400 | } |
401 | } |
402 | } |
403 | return condition; |
404 | } |
405 | |
406 | } |