1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22 package org.dbunit.operation;
23
24 import org.slf4j.Logger;
25 import org.slf4j.LoggerFactory;
26
27 import org.dbunit.DatabaseUnitException;
28 import org.dbunit.database.IDatabaseConnection;
29 import org.dbunit.database.statement.IPreparedBatchStatement;
30 import org.dbunit.database.statement.SimplePreparedStatement;
31 import org.dbunit.dataset.Column;
32 import org.dbunit.dataset.DataSetException;
33 import org.dbunit.dataset.IDataSet;
34 import org.dbunit.dataset.ITable;
35 import org.dbunit.dataset.ITableIterator;
36 import org.dbunit.dataset.ITableMetaData;
37 import org.dbunit.dataset.NoPrimaryKeyException;
38 import org.dbunit.dataset.RowOutOfBoundsException;
39 import org.dbunit.dataset.datatype.DataType;
40
41 import java.sql.PreparedStatement;
42 import java.sql.ResultSet;
43 import java.sql.SQLException;
44 import java.util.BitSet;
45
46
47
48
49
50
51
52
53
54
55
56 public class RefreshOperation extends AbstractOperation
57 {
58
59
60
61
62 private static final Logger logger = LoggerFactory.getLogger(RefreshOperation.class);
63
64 private final InsertOperation _insertOperation;
65 private final UpdateOperation _updateOperation;
66
67 RefreshOperation()
68 {
69 _insertOperation = (InsertOperation)DatabaseOperation.INSERT;
70 _updateOperation = (UpdateOperation)DatabaseOperation.UPDATE;
71 }
72
73 private boolean isEmpty(ITable table) throws DataSetException
74 {
75 return AbstractBatchOperation.isEmpty(table);
76 }
77
78
79
80
81 public void execute(IDatabaseConnection connection, IDataSet dataSet)
82 throws DatabaseUnitException, SQLException
83 {
84 logger.debug("execute(connection={}, dataSet) - start", connection);
85
86
87 ITableIterator iterator = dataSet.iterator();
88 while (iterator.next())
89 {
90 ITable table = iterator.getTable();
91
92
93 if (isEmpty(table))
94 {
95 continue;
96 }
97
98 ITableMetaData metaData = getOperationMetaData(connection,
99 table.getTableMetaData());
100 RowOperation updateRowOperation = createUpdateOperation(connection,
101 metaData);
102 RowOperation insertRowOperation = new InsertRowOperation(connection,
103 metaData);
104
105 try
106 {
107
108 for (int i = 0; ; i++)
109 {
110 if (!updateRowOperation.execute(table, i))
111 {
112 insertRowOperation.execute(table, i);
113 }
114 }
115 }
116 catch (RowOutOfBoundsException e)
117 {
118
119
120
121
122 }
123 finally
124 {
125
126 updateRowOperation.close();
127 insertRowOperation.close();
128 }
129 }
130
131 }
132
133 private RowOperation createUpdateOperation(IDatabaseConnection connection,
134 ITableMetaData metaData)
135 throws DataSetException, SQLException
136 {
137 logger.debug("createUpdateOperation(connection={}, metaData={}) - start", connection, metaData);
138
139
140 if (metaData.getColumns().length > metaData.getPrimaryKeys().length)
141 {
142 return new UpdateRowOperation(connection, metaData);
143 }
144
145
146 return new RowExistOperation(connection, metaData);
147 }
148
149
150
151
152 class RowOperation
153 {
154
155
156
157
158 private final Logger logger = LoggerFactory.getLogger(RowOperation.class);
159
160 protected IPreparedBatchStatement _statement;
161 protected OperationData _operationData;
162 protected BitSet _ignoreMapping;
163
164
165
166
167
168 public boolean execute(ITable table, int row)
169 throws DataSetException, SQLException
170 {
171 logger.debug("execute(table={}, row={}) - start", table, String.valueOf(row));
172
173 Column[] columns = _operationData.getColumns();
174 for (int i = 0; i < columns.length; i++)
175 {
176
177 if (_ignoreMapping == null || !_ignoreMapping.get(i))
178 {
179 Object value = table.getValue(row, columns[i].getColumnName());
180 _statement.addValue(value, columns[i].getDataType());
181 }
182 }
183 _statement.addBatch();
184 int result = _statement.executeBatch();
185 _statement.clearBatch();
186
187 return result == 1;
188 }
189
190
191
192
193 public void close() throws SQLException
194 {
195 logger.debug("close() - start");
196
197 if (_statement != null)
198 {
199 _statement.close();
200 }
201 }
202 }
203
204
205
206
207 private class InsertRowOperation extends RowOperation
208 {
209
210
211
212
213 private final Logger logger = LoggerFactory.getLogger(InsertRowOperation.class);
214
215 private IDatabaseConnection _connection;
216 private ITableMetaData _metaData;
217
218 public InsertRowOperation(IDatabaseConnection connection,
219 ITableMetaData metaData)
220 throws DataSetException, SQLException
221 {
222 _connection = connection;
223 _metaData = metaData;
224 }
225
226 public boolean execute(ITable table, int row)
227 throws DataSetException, SQLException
228 {
229 logger.debug("execute(table={}, row={}) - start", table, String.valueOf(row));
230
231
232
233 if (_ignoreMapping == null ||
234 !_insertOperation.equalsIgnoreMapping(_ignoreMapping, table, row))
235 {
236
237 if (_statement != null)
238 {
239 _statement.close();
240 }
241
242 _ignoreMapping = _insertOperation.getIgnoreMapping(table, row);
243 _operationData = _insertOperation.getOperationData(_metaData,
244 _ignoreMapping, _connection);
245 _statement = new SimplePreparedStatement(_operationData.getSql(),
246 _connection.getConnection());
247 }
248
249 return super.execute(table, row);
250 }
251
252 }
253
254
255
256
257 private class UpdateRowOperation extends RowOperation
258 {
259 PreparedStatement _countStatement;
260
261 public UpdateRowOperation(IDatabaseConnection connection,
262 ITableMetaData metaData)
263 throws DataSetException, SQLException
264 {
265
266 _operationData = _updateOperation.getOperationData(
267 metaData, null, connection);
268 _statement = new SimplePreparedStatement(_operationData.getSql(),
269 connection.getConnection());
270 }
271 }
272
273
274
275
276 private class RowExistOperation extends RowOperation
277 {
278
279
280
281
282 private final Logger logger = LoggerFactory.getLogger(RowExistOperation.class);
283
284 PreparedStatement _countStatement;
285
286 public RowExistOperation(IDatabaseConnection connection,
287 ITableMetaData metaData)
288 throws DataSetException, SQLException
289 {
290
291 _operationData = getSelectCountData(metaData, connection);
292 _countStatement = connection.getConnection().prepareStatement(
293 _operationData.getSql());
294 }
295
296 private OperationData getSelectCountData(
297 ITableMetaData metaData, IDatabaseConnection connection) throws DataSetException
298 {
299 logger.debug("getSelectCountData(metaData={}, connection={}) - start", metaData, connection);
300
301 Column[] primaryKeys = metaData.getPrimaryKeys();
302
303
304 if (primaryKeys.length == 0)
305 {
306 throw new NoPrimaryKeyException(metaData.getTableName());
307 }
308
309
310 StringBuffer sqlBuffer = new StringBuffer(128);
311 sqlBuffer.append("select COUNT(*) from ");
312 sqlBuffer.append(getQualifiedName(connection.getSchema(), metaData.getTableName(), connection));
313
314
315 sqlBuffer.append(" where ");
316 for (int i = 0; i < primaryKeys.length; i++)
317 {
318 Column column = primaryKeys[i];
319
320 if (i > 0)
321 {
322 sqlBuffer.append(" and ");
323 }
324 sqlBuffer.append(getQualifiedName(null, column.getColumnName(), connection));
325 sqlBuffer.append(" = ?");
326 }
327
328 return new OperationData(sqlBuffer.toString(), primaryKeys);
329 }
330
331
332
333
334
335
336
337
338 public boolean execute(ITable table, int row)
339 throws DataSetException, SQLException
340 {
341 logger.debug("execute(table={}, row={}) - start", table, String.valueOf(row));
342
343 Column[] columns = _operationData.getColumns();
344 for (int i = 0; i < columns.length; i++)
345 {
346 Object value = table.getValue(row, columns[i].getColumnName());
347 DataType dataType = columns[i].getDataType();
348 dataType.setSqlValue(value, i + 1, _countStatement);
349 }
350
351 ResultSet resultSet = _countStatement.executeQuery();
352 try
353 {
354 resultSet.next();
355 return resultSet.getInt(1) > 0;
356 }
357 finally
358 {
359 resultSet.close();
360 }
361 }
362
363 public void close() throws SQLException
364 {
365 logger.debug("close() - start");
366
367 _countStatement.close();
368 }
369 }
370
371 }