jpauserset now working.
[utils] / test / enterprise / src / main / java / org / wamblee / support / persistence / DatabaseUtils.java
1 /*
2  * Copyright 2005-2010 the original author or authors.
3  * 
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  * 
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  * 
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package org.wamblee.support.persistence;
17
18 import java.sql.Connection;
19 import java.sql.PreparedStatement;
20 import java.sql.ResultSet;
21 import java.sql.SQLException;
22 import java.util.ArrayList;
23 import java.util.List;
24 import java.util.logging.Logger;
25
26 import javax.sql.DataSource;
27
28 import junit.framework.TestCase;
29
30 import org.dbunit.DataSourceDatabaseTester;
31 import org.dbunit.IDatabaseTester;
32 import org.dbunit.database.DatabaseConnection;
33 import org.dbunit.database.DatabaseSequenceFilter;
34 import org.dbunit.database.IDatabaseConnection;
35 import org.dbunit.dataset.FilteredDataSet;
36 import org.dbunit.dataset.IDataSet;
37 import org.dbunit.dataset.filter.ITableFilter;
38 import org.dbunit.dataset.filter.ITableFilterSimple;
39 import org.dbunit.operation.DatabaseOperation;
40
41 /**
42  * Database utilities is a simple support class for common tasks in working with
43  * databases.
44  */
45 public class DatabaseUtils {
46
47     public static interface TableSet {
48         boolean contains(String aTableName);
49     }
50
51     public static interface JdbcUnitOfWork<T> {
52         T execute(Connection aConnection) throws Exception;
53     }
54
55     public static interface TableSetOperation {
56         void execute(String aTable) throws Exception;
57     }
58
59     private static final Logger LOG = Logger.getLogger(DatabaseUtils.class
60         .getName());
61
62     /**
63      * Schema pattern.
64      */
65     private static final String SCHEMA_PATTERN = "%";
66     private DataSource dataSource;
67     private ITableFilterSimple tables;
68
69     public DatabaseUtils(DataSource aDataSource, ITableFilterSimple aTables) {
70         dataSource = aDataSource;
71         tables = aTables;
72     }
73
74     public IDatabaseTester createDbTester() throws Exception {
75         return createDbTester(getTableNames(tables));
76     }
77
78     public IDatabaseTester createDbTester(String[] aTables) throws Exception {
79         IDatabaseTester dbtester = new DataSourceDatabaseTester(dataSource);
80         dbtester.setDataSet(dbtester.getConnection().createDataSet(aTables));
81         return dbtester;
82     }
83
84     public void cleanDatabase() throws Exception {
85         cleanDatabase(tables);
86     }
87
88     public void executeOnTables(ITableFilterSimple aTables,
89         final TableSetOperation aOperation) throws Exception {
90         final String[] tables = getTableNames(aTables);
91         executeInTransaction(new JdbcUnitOfWork<Void>() {
92             public Void execute(Connection aConnection) throws Exception {
93                 for (int i = tables.length-1; i >= 0; i--) {
94                     aOperation.execute(tables[i]);
95                 }
96                 return null;
97             }
98         });
99         for (String table : tables) {
100
101         }
102     }
103
104     public void cleanDatabase(ITableFilterSimple aSelection) throws Exception {
105
106         final String[] tables = getTableNames(aSelection);
107         executeInTransaction(new JdbcUnitOfWork<Void>() {
108
109             public Void execute(Connection aConnection) throws Exception {
110                 IDatabaseConnection connection = new DatabaseConnection(
111                     aConnection);
112                 ITableFilter filter = new DatabaseSequenceFilter(connection,
113                     tables);
114                 IDataSet dataset = new FilteredDataSet(filter, connection
115                     .createDataSet(tables));
116                 DatabaseOperation.DELETE_ALL.execute(connection, dataset);
117                 return null;
118             }
119         });
120
121     }
122
123     public <T> T executeInTransaction(JdbcUnitOfWork<T> aCallback)
124         throws Exception {
125         Connection connection = dataSource.getConnection();
126         try {
127             T value = aCallback.execute(connection);
128             connection.commit();
129             return value;
130         } finally {
131             connection.close();
132         }
133     }
134
135     public String[] getTableNames() throws Exception {
136         return getTableNames(tables);
137     }
138
139     /**
140      * @throws SQLException
141      */
142     public String[] getTableNames(ITableFilterSimple aSelection)
143         throws Exception {
144
145         List<String> result = new ArrayList<String>();
146         LOG.fine("Getting database table names to clean (schema: '" +
147             SCHEMA_PATTERN + "'");
148
149         Connection connection = dataSource.getConnection();
150         try {
151             ResultSet tables = connection.getMetaData().getTables(null,
152                 SCHEMA_PATTERN, "%", new String[] { "TABLE" });
153             while (tables.next()) {
154                 String table = tables.getString("TABLE_NAME");
155                 if (aSelection.accept(table)) {
156                     result.add(table);
157                 }
158             }
159             return (String[]) result.toArray(new String[0]);
160         } finally {
161             connection.close();
162         }
163     }
164
165     public void emptyTables() throws Exception {
166         executeOnTables(tables, new TableSetOperation() {
167             public void execute(String aTable) throws Exception {
168                 emptyTable(aTable);
169             }
170         });
171     }
172
173     /**
174      * @return
175      * @throws SQLException
176      */
177     public void emptyTables(final ITableFilterSimple aSelection)
178         throws Exception {
179         executeOnTables(aSelection, new TableSetOperation() {
180             public void execute(String aTable) throws Exception {
181                 emptyTable(aTable);
182             }
183         });
184     }
185
186     /**
187      * @return
188      * @throws SQLException
189      */
190     public void emptyTable(String aTable) throws Exception {
191         executeSql("delete from " + aTable);
192     }
193
194     public void dropTables() throws Exception {
195         executeOnTables(tables, new TableSetOperation() {
196
197             public void execute(String aTable) throws Exception {
198                 dropTable(aTable);
199             }
200         });
201     }
202
203     public void dropTables(ITableFilterSimple aTables) throws Exception {
204         executeOnTables(aTables, new TableSetOperation() {
205
206             public void execute(String aTable) throws Exception {
207                 dropTable(aTable);
208             }
209         });
210     }
211
212     /**
213      * @return
214      * @throws SQLException
215      */
216     public void dropTable(final String aTable) throws Exception {
217         executeInTransaction(new JdbcUnitOfWork<Void>() {
218             public Void execute(Connection aConnection) throws Exception {
219                 executeUpdate(aConnection, "drop table " + aTable);
220                 return null;
221             }
222         });
223
224     }
225
226     /**
227      * Executes an SQL statement within a transaction.
228      * 
229      * @param aSql
230      *            SQL statement.
231      * @return Return code of the corresponding JDBC call.
232      */
233     public int executeSql(final String aSql) throws Exception {
234         return executeSql(aSql, new Object[0]);
235     }
236
237     /**
238      * Executes an SQL statement within a transaction. See
239      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
240      * supported argument types.
241      * 
242      * @param aSql
243      *            SQL statement.
244      * @param aArg
245      *            Argument of the sql statement.
246      * @return Return code of the corresponding JDBC call.
247      */
248     public int executeSql(final String aSql, final Object aArg)
249         throws Exception {
250         return executeSql(aSql, new Object[] { aArg });
251     }
252
253     /**
254      * Executes an sql statement. See
255      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
256      * supported argument types.
257      * 
258      * @param aSql
259      *            SQL query to execute.
260      * @param aArgs
261      *            Arguments.
262      * @return Number of rows updated.
263      */
264     public int executeSql(final String aSql, final Object[] aArgs)
265         throws Exception {
266         return executeInTransaction(new JdbcUnitOfWork<Integer>() {
267             public Integer execute(Connection aConnection) throws Exception {
268                 PreparedStatement stmt = aConnection.prepareStatement(aSql);
269                 setPreparedParams(aArgs, stmt);
270                 return stmt.executeUpdate();
271             }
272         });
273     }
274
275     /**
276      * Executes an SQL query.
277      * 
278      * @param aSql
279      *            Query to execute.
280      * @return Result set.
281      */
282     public ResultSet executeQuery(Connection aConnection, String aSql) {
283         return executeQuery(aConnection, aSql, new Object[0]);
284     }
285
286     /**
287      * Executes a query with a single argument. See
288      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
289      * supported argument types.
290      * 
291      * @param aSql
292      *            Query.
293      * @param aArg
294      *            Argument.
295      * @return Result set.
296      */
297     public ResultSet executeQuery(Connection aConnection, String aSql,
298         Object aArg) {
299         return executeQuery(aConnection, aSql, new Object[] { aArg });
300     }
301
302     /**
303      * Executes a query. See
304      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
305      * supported argument types.
306      * 
307      * @param aSql
308      *            Sql query.
309      * @param aArgs
310      *            Arguments to the query.
311      * @return Result set.
312      */
313     public ResultSet executeQuery(Connection aConnection, final String aSql,
314         final Object[] aArgs) {
315         try {
316             PreparedStatement statement = aConnection.prepareStatement(aSql);
317             setPreparedParams(aArgs, statement);
318
319             return statement.executeQuery();
320         } catch (SQLException e) {
321             throw new RuntimeException(e);
322         }
323     }
324
325     public int executeUpdate(Connection aConnection, final String aSql,
326         final Object... aArgs) {
327         try {
328             PreparedStatement statement = aConnection.prepareStatement(aSql);
329             setPreparedParams(aArgs, statement);
330
331             return statement.executeUpdate();
332         } catch (SQLException e) {
333             throw new RuntimeException(e);
334         }
335     }
336
337     /**
338      * Sets the values of a prepared statement. See
339      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
340      * supported argument types.
341      * 
342      * @param aArgs
343      *            Arguments to the prepared statement.
344      * @param aStatement
345      *            Prepared statement
346      * @throws SQLException
347      */
348     private void setPreparedParams(final Object[] aArgs,
349         PreparedStatement aStatement) throws SQLException {
350         for (int i = 1; i <= aArgs.length; i++) {
351             setPreparedParam(i, aStatement, aArgs[i - 1]);
352         }
353     }
354
355     /**
356      * Sets a prepared statement parameter.
357      * 
358      * @param aIndex
359      *            Index of the parameter.
360      * @param aStatement
361      *            Prepared statement.
362      * @param aObject
363      *            Value Must be of type Integer, Long, or String.
364      * @throws SQLException
365      */
366     private void setPreparedParam(int aIndex, PreparedStatement aStatement,
367         Object aObject) throws SQLException {
368         if (aObject instanceof Integer) {
369             aStatement.setInt(aIndex, ((Integer) aObject).intValue());
370         } else if (aObject instanceof Long) {
371             aStatement.setLong(aIndex, ((Long) aObject).longValue());
372         } else if (aObject instanceof String) {
373             aStatement.setString(aIndex, (String) aObject);
374         } else {
375             TestCase.fail("Unsupported object type for prepared statement: " +
376                 aObject.getClass() + " value: " + aObject + " statement: " +
377                 aStatement);
378         }
379     }
380
381     /**
382      * @return
383      * @throws SQLException
384      */
385     public int getTableSize(final String aTable) throws Exception {
386         return executeInTransaction(new JdbcUnitOfWork<Integer>() {
387             public Integer execute(Connection aConnection) throws Exception {
388                 ResultSet resultSet = executeQuery(aConnection,
389                     "select count(*) from " + aTable);
390                 resultSet.next();
391                 return resultSet.getInt(1);
392             }
393         });
394
395     }
396
397     public int countResultSet(ResultSet aResultSet) throws SQLException {
398         int count = 0;
399
400         while (aResultSet.next()) {
401             count++;
402         }
403
404         return count;
405     }
406
407 }