(no commit message)
[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.Level;
25 import java.util.logging.Logger;
26
27 import javax.sql.DataSource;
28
29 import junit.framework.TestCase;
30
31 import org.dbunit.DataSourceDatabaseTester;
32 import org.dbunit.IDatabaseTester;
33 import org.dbunit.database.DatabaseConnection;
34 import org.dbunit.database.DatabaseSequenceFilter;
35 import org.dbunit.database.IDatabaseConnection;
36 import org.dbunit.dataset.FilteredDataSet;
37 import org.dbunit.dataset.IDataSet;
38 import org.dbunit.dataset.filter.ITableFilter;
39 import org.dbunit.dataset.filter.ITableFilterSimple;
40 import org.dbunit.operation.DatabaseOperation;
41
42 /**
43  * Database utilities is a simple support class for common tasks in working with
44  * databases.
45  */
46 public class DatabaseUtils {
47
48     public static interface TableSet {
49         boolean contains(String aTableName);
50     }
51
52     public static interface JdbcUnitOfWork<T> {
53         T execute(Connection aConnection) throws Exception;
54     }
55
56     public static interface TableSetOperation {
57         void execute(String aTable) throws Exception;
58     }
59
60     private static final Logger LOG = Logger.getLogger(DatabaseUtils.class
61         .getName());
62
63     /**
64      * Schema pattern.
65      */
66     private static final String SCHEMA_PATTERN = "%";
67     private DataSource dataSource;
68     private ITableFilterSimple tables;
69
70     private IDatabaseTester dbtester;
71     /**
72      * List of connections that were created for dbtesters. 
73      * This list will be closed in the {@link #stop()} method. 
74      */
75     private List<IDatabaseConnection> connections; 
76     
77     /**
78      * Constructs the database utils. 
79      * Before use, {@link #start()} must be called. 
80      * @param aDataSource Datasource. 
81      * @param aTables Tables to work with. 
82      */
83     public DatabaseUtils(DataSource aDataSource, ITableFilterSimple aTables) {
84         dataSource = aDataSource;
85         tables = aTables;
86         dbtester = new DataSourceDatabaseTester(dataSource);
87         connections = new ArrayList<IDatabaseConnection>(); 
88     }
89
90     /**
91      * Starts the database utils. 
92      */
93     public void start() { 
94         // Empty. No operation currently. 
95     }
96
97     /**
98      * Stops the database utils, closing any JDBC connections that were created
99      * by this utility. Note that connections obtained from the datasource directly 
100      * must still be closed by the user. The involved connections are only those that
101      * are created by this utility. 
102      */
103     public void stop() {
104         for (IDatabaseConnection connection: connections) { 
105             try {
106                 connection.close();
107             } catch (SQLException e) {
108                 LOG.log(Level.WARNING, "Could not close connection", e);
109             }
110         }
111         connections.clear();
112     }
113
114     public IDatabaseTester createDbTester() throws Exception {
115         return createDbTester(getTableNames(tables));
116     }
117
118     public IDatabaseTester createDbTester(String[] aTables) throws Exception {
119         IDatabaseConnection connection = dbtester.getConnection();
120         connections.add(connection);
121         dbtester.setDataSet(connection.createDataSet(aTables));
122         return dbtester;
123     }
124
125     public void cleanDatabase() throws Exception {
126         cleanDatabase(tables);
127     }
128
129     public void executeOnTables(ITableFilterSimple aTables,
130         final TableSetOperation aOperation) throws Exception {
131         final String[] tableNames = getTableNames(aTables);
132         executeInTransaction(new JdbcUnitOfWork<Void>() {
133             public Void execute(Connection aConnection) throws Exception {
134                 for (int i = tableNames.length - 1; i >= 0; i--) {
135                     aOperation.execute(tableNames[i]);
136                 }
137                 return null;
138             }
139         });
140     }
141
142     public void cleanDatabase(ITableFilterSimple aSelection) throws Exception {
143
144         final String[] tableNames = getTableNames(aSelection);
145         executeInTransaction(new JdbcUnitOfWork<Void>() {
146
147             public Void execute(Connection aConnection) throws Exception {
148                 IDatabaseConnection connection = new DatabaseConnection(
149                     aConnection);
150                 ITableFilter filter = new DatabaseSequenceFilter(connection,
151                     tableNames);
152                 IDataSet dataset = new FilteredDataSet(filter, connection
153                     .createDataSet(tableNames));
154                 DatabaseOperation.DELETE_ALL.execute(connection, dataset);
155                 return null;
156             }
157         });
158
159     }
160
161     public <T> T executeInTransaction(JdbcUnitOfWork<T> aCallback)
162         throws Exception {
163         Connection connection = dataSource.getConnection();
164         connection.setAutoCommit(false);
165         try {
166             T value = aCallback.execute(connection);
167             connection.commit();
168             return value;
169         } finally {
170             connection.close();
171         }
172     }
173
174     public String[] getTableNames() throws Exception {
175         return getTableNames(tables);
176     }
177
178     /**
179      * @throws SQLException
180      */
181     public String[] getTableNames(ITableFilterSimple aSelection)
182         throws Exception {
183
184         List<String> result = new ArrayList<String>();
185         LOG.fine("Getting database table names to clean (schema: '" +
186             SCHEMA_PATTERN + "'");
187
188         Connection connection = dataSource.getConnection();
189         try {
190             ResultSet tableNames = connection.getMetaData().getTables(null,
191                 SCHEMA_PATTERN, "%", new String[] { "TABLE" });
192             while (tableNames.next()) {
193                 String table = tableNames.getString("TABLE_NAME");
194                 if (aSelection.accept(table)) {
195                     result.add(table);
196                 }
197             }
198             return (String[]) result.toArray(new String[0]);
199         } finally {
200             connection.close();
201         }
202     }
203
204     public void emptyTables() throws Exception {
205         executeOnTables(tables, new TableSetOperation() {
206             public void execute(String aTable) throws Exception {
207                 emptyTable(aTable);
208             }
209         });
210     }
211
212     /**
213      * @return
214      * @throws SQLException
215      */
216     public void emptyTables(final ITableFilterSimple aSelection)
217         throws Exception {
218         executeOnTables(aSelection, new TableSetOperation() {
219             public void execute(String aTable) throws Exception {
220                 emptyTable(aTable);
221             }
222         });
223     }
224
225     /**
226      * @return
227      * @throws SQLException
228      */
229     public void emptyTable(String aTable) throws Exception {
230         executeSql("delete from " + aTable);
231     }
232
233     public void dropTables() throws Exception {
234         dropTables(tables);
235     }
236
237     public void dropTables(ITableFilterSimple aTables) throws Exception {
238         final String[] tableNames = getTableNames(aTables);
239         String[] sortedTables = executeInTransaction(new JdbcUnitOfWork<String[]>() {
240
241             public String[] execute(Connection aConnection) throws Exception {
242                 IDatabaseConnection connection = new DatabaseConnection(
243                     aConnection);
244                 ITableFilter filter = new DatabaseSequenceFilter(connection,
245                     tableNames);
246                 IDataSet dataset = new FilteredDataSet(filter, connection
247                     .createDataSet(tableNames));
248                 return dataset.getTableNames();
249             }
250         });
251         for (int i = sortedTables.length - 1; i >= 0; i--) {
252             dropTable(sortedTables[i]);
253         }
254     }
255
256     /**
257      * @return
258      * @throws SQLException
259      */
260     public void dropTable(final String aTable) throws Exception {
261         executeInTransaction(new JdbcUnitOfWork<Void>() {
262             public Void execute(Connection aConnection) throws Exception {
263                 executeUpdate(aConnection, "drop table " + aTable);
264                 return null;
265             }
266         });
267
268     }
269
270     /**
271      * Executes an SQL statement within a transaction.
272      * 
273      * @param aSql
274      *            SQL statement.
275      * @return Return code of the corresponding JDBC call.
276      */
277     public int executeSql(final String aSql) throws Exception {
278         return executeSql(aSql, new Object[0]);
279     }
280
281     /**
282      * Executes an SQL statement within a transaction. See
283      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
284      * supported argument types.
285      * 
286      * @param aSql
287      *            SQL statement.
288      * @param aArg
289      *            Argument of the sql statement.
290      * @return Return code of the corresponding JDBC call.
291      */
292     public int executeSql(final String aSql, final Object aArg)
293         throws Exception {
294         return executeSql(aSql, new Object[] { aArg });
295     }
296
297     /**
298      * Executes an sql statement. See
299      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
300      * supported argument types.
301      * 
302      * @param aSql
303      *            SQL query to execute.
304      * @param aArgs
305      *            Arguments.
306      * @return Number of rows updated.
307      */
308     public int executeSql(final String aSql, final Object[] aArgs)
309         throws Exception {
310         return executeInTransaction(new JdbcUnitOfWork<Integer>() {
311             public Integer execute(Connection aConnection) throws Exception {
312                 PreparedStatement stmt = aConnection.prepareStatement(aSql);
313                 setPreparedParams(aArgs, stmt);
314                 return stmt.executeUpdate();
315             }
316         });
317     }
318
319     /**
320      * Executes an SQL query.
321      * 
322      * @param aSql
323      *            Query to execute.
324      * @return Result set.
325      */
326     public ResultSet executeQuery(Connection aConnection, String aSql) {
327         return executeQuery(aConnection, aSql, new Object[0]);
328     }
329
330     /**
331      * Executes a query with a single argument. See
332      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
333      * supported argument types.
334      * 
335      * @param aSql
336      *            Query.
337      * @param aArg
338      *            Argument.
339      * @return Result set.
340      */
341     public ResultSet executeQuery(Connection aConnection, String aSql,
342         Object aArg) {
343         return executeQuery(aConnection, aSql, new Object[] { aArg });
344     }
345
346     /**
347      * Executes a query. See
348      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
349      * supported argument types.
350      * 
351      * @param aSql
352      *            Sql query.
353      * @param aArgs
354      *            Arguments to the query.
355      * @return Result set.
356      */
357     public ResultSet executeQuery(Connection aConnection, final String aSql,
358         final Object[] aArgs) {
359         try {
360             PreparedStatement statement = aConnection.prepareStatement(aSql);
361             setPreparedParams(aArgs, statement);
362
363             return statement.executeQuery();
364         } catch (SQLException e) {
365             throw new RuntimeException(e);
366         }
367     }
368
369     public int executeUpdate(Connection aConnection, final String aSql,
370         final Object... aArgs) {
371         try {
372             PreparedStatement statement = aConnection.prepareStatement(aSql);
373             setPreparedParams(aArgs, statement);
374
375             return statement.executeUpdate();
376         } catch (SQLException e) {
377             throw new RuntimeException(e);
378         }
379     }
380
381     /**
382      * Sets the values of a prepared statement. See
383      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
384      * supported argument types.
385      * 
386      * @param aArgs
387      *            Arguments to the prepared statement.
388      * @param aStatement
389      *            Prepared statement
390      * @throws SQLException
391      */
392     private void setPreparedParams(final Object[] aArgs,
393         PreparedStatement aStatement) throws SQLException {
394         for (int i = 1; i <= aArgs.length; i++) {
395             setPreparedParam(i, aStatement, aArgs[i - 1]);
396         }
397     }
398
399     /**
400      * Sets a prepared statement parameter.
401      * 
402      * @param aIndex
403      *            Index of the parameter.
404      * @param aStatement
405      *            Prepared statement.
406      * @param aObject
407      *            Value Must be of type Integer, Long, or String.
408      * @throws SQLException
409      */
410     private void setPreparedParam(int aIndex, PreparedStatement aStatement,
411         Object aObject) throws SQLException {
412         if (aObject instanceof Integer) {
413             aStatement.setInt(aIndex, ((Integer) aObject).intValue());
414         } else if (aObject instanceof Long) {
415             aStatement.setLong(aIndex, ((Long) aObject).longValue());
416         } else if (aObject instanceof String) {
417             aStatement.setString(aIndex, (String) aObject);
418         } else {
419             TestCase.fail("Unsupported object type for prepared statement: " +
420                 aObject.getClass() + " value: " + aObject + " statement: " +
421                 aStatement);
422         }
423     }
424
425     /**
426      * @return
427      * @throws SQLException
428      */
429     public int getTableSize(final String aTable) throws Exception {
430         return executeInTransaction(new JdbcUnitOfWork<Integer>() {
431             public Integer execute(Connection aConnection) throws Exception {
432                 ResultSet resultSet = executeQuery(aConnection,
433                     "select count(*) from " + aTable);
434                 resultSet.next();
435                 return resultSet.getInt(1);
436             }
437         });
438
439     }
440
441     public int countResultSet(ResultSet aResultSet) throws SQLException {
442         int count = 0;
443
444         while (aResultSet.next()) {
445             count++;
446         }
447
448         return count;
449     }
450
451 }