2 * Copyright 2005-2010 the original author or authors.
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
16 package org.wamblee.support.persistence;
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;
27 import javax.sql.DataSource;
29 import junit.framework.TestCase;
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;
43 * Database utilities is a simple support class for common tasks in working with
46 public class DatabaseUtils {
49 * Represents a set of tables.
51 * @author Erik Brakkee
53 public static interface TableSet {
54 boolean contains(String aTableName);
58 * Represents a unit of work (transaction).
60 * @author Erik Brakkee
62 * @param <T> Type of return value.
64 public static interface JdbcUnitOfWork<T> {
66 * Executes statement within a transaction.
67 * @param aConnection Connection.
68 * @return Result of the work.
71 T execute(Connection aConnection) throws Exception;
75 * Operation to be executed on a set of tables for each table
78 * @author Erik Brakkee
80 public static interface TableSetOperation {
82 * Executes on a table.
83 * @param aTable Table name.
86 void execute(String aTable) throws Exception;
89 private static final Logger LOG = Logger.getLogger(DatabaseUtils.class
95 private static final String SCHEMA_PATTERN = "%";
96 private DataSource dataSource;
98 private IDatabaseTester dbtester;
100 * List of connections that were created for dbtesters.
101 * This list will be closed in the {@link #stop()} method.
103 private List<IDatabaseConnection> connections;
106 * Constructs the database utils.
107 * Before use, {@link #start()} must be called.
108 * @param aDataSource Datasource.
110 public DatabaseUtils(DataSource aDataSource) {
111 dataSource = aDataSource;
112 dbtester = new DataSourceDatabaseTester(dataSource);
113 connections = new ArrayList<IDatabaseConnection>();
117 * Starts the database utils.
119 public void start() {
120 // Empty. No operation currently.
124 * Stops the database utils, closing any JDBC connections that were created
125 * by this utility. Note that connections obtained from the datasource directly
126 * must still be closed by the user. The involved connections are only those that
127 * are created by this utility.
130 for (IDatabaseConnection connection: connections) {
133 } catch (SQLException e) {
134 LOG.log(Level.WARNING, "Could not close connection", e);
141 * Creates database tester.
142 * @param aTables Tables to create the tester for.
143 * @return Database tester.
146 public IDatabaseTester createDbTester(ITableFilterSimple aTables) throws Exception {
147 return createDbTester(getTableNames(aTables));
151 * Creates database tester.
152 * @param aTables Tables to create the tester for.
153 * @return Database tester.
156 public IDatabaseTester createDbTester(String[] aTables) throws Exception {
157 IDatabaseConnection connection = dbtester.getConnection();
158 connections.add(connection);
159 dbtester.setDataSet(connection.createDataSet(aTables));
164 * Executes an operation on a set of tables.
165 * @param aTables Tables.
166 * @param aOperation Operation.
169 public void executeOnTables(ITableFilterSimple aTables,
170 final TableSetOperation aOperation) throws Exception {
171 final String[] tableNames = getTableNames(aTables);
172 executeInTransaction(new JdbcUnitOfWork<Void>() {
173 public Void execute(Connection aConnection) throws Exception {
174 for (int i = tableNames.length - 1; i >= 0; i--) {
175 aOperation.execute(tableNames[i]);
183 * Cleans a number of database tables. This means deleting the content not dropping the tables.
184 * This may fail in case of cyclic dependencies between the tables (current limitation).
185 * @param aSelection Tables.
188 public void cleanDatabase(ITableFilterSimple aSelection) throws Exception {
190 final String[] tableNames = getTableNames(aSelection);
191 executeInTransaction(new JdbcUnitOfWork<Void>() {
193 public Void execute(Connection aConnection) throws Exception {
194 IDatabaseConnection connection = new DatabaseConnection(
196 ITableFilter filter = new DatabaseSequenceFilter(connection,
198 IDataSet dataset = new FilteredDataSet(filter, connection
199 .createDataSet(tableNames));
200 DatabaseOperation.DELETE_ALL.execute(connection, dataset);
208 * Executes a unit of work within a transaction.
209 * @param <T> Result type of th ework.
210 * @param aWork Unit of work.
214 public <T> T executeInTransaction(JdbcUnitOfWork<T> aWork)
216 Connection connection = dataSource.getConnection();
217 connection.setAutoCommit(false);
219 T value = aWork.execute(connection);
228 * Returns table names based on a table filter.
229 * @param aSelection Table filter.
230 * @return Table names.
233 public String[] getTableNames(ITableFilterSimple aSelection)
236 List<String> result = new ArrayList<String>();
237 LOG.fine("Getting database table names to clean (schema: '" +
238 SCHEMA_PATTERN + "'");
240 Connection connection = dataSource.getConnection();
242 ResultSet tableNames = connection.getMetaData().getTables(null,
243 SCHEMA_PATTERN, "%", new String[] { "TABLE" });
244 while (tableNames.next()) {
245 String table = tableNames.getString("TABLE_NAME");
246 if (aSelection.accept(table)) {
250 return (String[]) result.toArray(new String[0]);
257 * Use {@link #cleanDatabase(ITableFilterSimple)} instead.
260 public void emptyTables(final ITableFilterSimple aSelection)
262 executeOnTables(aSelection, new TableSetOperation() {
263 public void execute(String aTable) throws Exception {
270 * User {@link #cleanDatabase(ITableFilterSimple)} instead.
273 public void emptyTable(String aTable) throws Exception {
274 executeSql("delete from " + aTable);
278 * Drops tables. This only works if there are no cyclic dependencies between the tables.
279 * @param aTables Tables to drop.
282 public void dropTables(ITableFilterSimple aTables) throws Exception {
283 final String[] tableNames = getTableNames(aTables);
284 String[] sortedTables = executeInTransaction(new JdbcUnitOfWork<String[]>() {
286 public String[] execute(Connection aConnection) throws Exception {
287 IDatabaseConnection connection = new DatabaseConnection(
289 ITableFilter filter = new DatabaseSequenceFilter(connection,
291 IDataSet dataset = new FilteredDataSet(filter, connection
292 .createDataSet(tableNames));
293 return dataset.getTableNames();
296 for (int i = sortedTables.length - 1; i >= 0; i--) {
297 dropTable(sortedTables[i]);
303 * @param aTable Table to drop.
306 public void dropTable(final String aTable) throws Exception {
307 executeInTransaction(new JdbcUnitOfWork<Void>() {
308 public Void execute(Connection aConnection) throws Exception {
309 executeUpdate(aConnection, "drop table " + aTable);
317 * Executes an SQL statement within a transaction.
321 * @return Return code of the corresponding JDBC call.
323 public int executeSql(final String aSql) throws Exception {
324 return executeSql(aSql, new Object[0]);
328 * Executes an SQL statement within a transaction. See
329 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
330 * supported argument types.
335 * Argument of the sql statement.
336 * @return Return code of the corresponding JDBC call.
338 public int executeSql(final String aSql, final Object aArg)
340 return executeSql(aSql, new Object[] { aArg });
344 * Executes an sql statement. See
345 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
346 * supported argument types.
349 * SQL query to execute.
352 * @return Number of rows updated.
354 public int executeSql(final String aSql, final Object[] aArgs)
356 return executeInTransaction(new JdbcUnitOfWork<Integer>() {
357 public Integer execute(Connection aConnection) throws Exception {
358 PreparedStatement stmt = aConnection.prepareStatement(aSql);
359 setPreparedParams(aArgs, stmt);
360 return stmt.executeUpdate();
366 * Executes an SQL query.
370 * @return Result set.
372 public ResultSet executeQuery(Connection aConnection, String aSql) {
373 return executeQuery(aConnection, aSql, new Object[0]);
377 * Executes a query with a single argument. See
378 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
379 * supported argument types.
385 * @return Result set.
387 public ResultSet executeQuery(Connection aConnection, String aSql,
389 return executeQuery(aConnection, aSql, new Object[] { aArg });
393 * Executes a query. See
394 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
395 * supported argument types.
400 * Arguments to the query.
401 * @return Result set.
403 public ResultSet executeQuery(Connection aConnection, final String aSql,
404 final Object[] aArgs) {
406 PreparedStatement statement = aConnection.prepareStatement(aSql);
407 setPreparedParams(aArgs, statement);
409 return statement.executeQuery();
410 } catch (SQLException e) {
411 throw new RuntimeException(e);
416 * Executes an update.
417 * @param aConnection Connection to use.
418 * @param aSql SQL update to use.
419 * @param aArgs Arguments to the update.
420 * @return Number of rows updated.
422 public int executeUpdate(Connection aConnection, final String aSql,
423 final Object... aArgs) {
425 PreparedStatement statement = aConnection.prepareStatement(aSql);
426 setPreparedParams(aArgs, statement);
428 return statement.executeUpdate();
429 } catch (SQLException e) {
430 throw new RuntimeException(e);
435 * Sets the values of a prepared statement. See
436 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
437 * supported argument types.
440 * Arguments to the prepared statement.
443 * @throws SQLException
445 private void setPreparedParams(final Object[] aArgs,
446 PreparedStatement aStatement) throws SQLException {
447 for (int i = 1; i <= aArgs.length; i++) {
448 setPreparedParam(i, aStatement, aArgs[i - 1]);
453 * Sets a prepared statement parameter.
456 * Index of the parameter.
458 * Prepared statement.
460 * Value Must be of type Integer, Long, or String.
461 * @throws SQLException
463 private void setPreparedParam(int aIndex, PreparedStatement aStatement,
464 Object aObject) throws SQLException {
465 if (aObject instanceof Integer) {
466 aStatement.setInt(aIndex, ((Integer) aObject).intValue());
467 } else if (aObject instanceof Long) {
468 aStatement.setLong(aIndex, ((Long) aObject).longValue());
469 } else if (aObject instanceof String) {
470 aStatement.setString(aIndex, (String) aObject);
472 TestCase.fail("Unsupported object type for prepared statement: " +
473 aObject.getClass() + " value: " + aObject + " statement: " +
479 * Gets the table size.
480 * @param aTable Table.
481 * @return Table size.
482 * @throws SQLException
484 public int getTableSize(final String aTable) throws Exception {
485 return executeInTransaction(new JdbcUnitOfWork<Integer>() {
486 public Integer execute(Connection aConnection) throws Exception {
487 ResultSet resultSet = executeQuery(aConnection,
488 "select count(*) from " + aTable);
490 return resultSet.getInt(1);
497 * Counts the results in a result set.
498 * @param aResultSet Resultset.
499 * @return Number of rows in the set.
500 * @throws SQLException
502 public int countResultSet(ResultSet aResultSet) throws SQLException {
505 while (aResultSet.next()) {