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.test.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
63 * Type of return value.
65 public static interface JdbcUnitOfWork<T> {
67 * Executes statement within a transaction.
71 * @return Result of the work.
74 T execute(Connection aConnection) throws Exception;
78 * Operation to be executed on a set of tables for each table individually.
80 * @author Erik Brakkee
82 public static interface TableSetOperation {
84 * Executes on a table.
90 void execute(String aTable) throws Exception;
93 private static final Logger LOG = Logger.getLogger(DatabaseUtils.class
99 private static final String SCHEMA_PATTERN = "%";
100 private DataSource dataSource;
102 private IDatabaseTester dbtester;
104 * List of connections that were created for dbtesters. This list will be
105 * closed in the {@link #stop()} method.
107 private List<IDatabaseConnection> connections;
110 * Constructs the database utils. Before use, {@link #start()} must be
116 public DatabaseUtils(DataSource aDataSource) {
117 dataSource = aDataSource;
118 dbtester = new DataSourceDatabaseTester(dataSource);
119 connections = new ArrayList<IDatabaseConnection>();
123 * Starts the database utils.
125 public void start() {
126 // Empty. No operation currently.
130 * Stops the database utils, closing any JDBC connections that were created
131 * by this utility. Note that connections obtained from the datasource
132 * directly must still be closed by the user. The involved connections are
133 * only those that are created by this utility.
136 for (IDatabaseConnection connection : connections) {
139 } catch (SQLException e) {
140 LOG.log(Level.WARNING, "Could not close connection", e);
147 * Creates database tester.
150 * Tables to create the tester for.
151 * @return Database tester.
154 public IDatabaseTester createDbTester(ITableFilterSimple aTables)
156 return createDbTester(getTableNames(aTables));
160 * Creates database tester.
163 * Tables to create the tester for.
164 * @return Database tester.
167 public IDatabaseTester createDbTester(String[] aTables) throws Exception {
168 IDatabaseConnection connection = dbtester.getConnection();
169 connections.add(connection);
170 dbtester.setDataSet(connection.createDataSet(aTables));
175 * Executes an operation on a set of tables.
183 public void executeOnTables(ITableFilterSimple aTables,
184 final TableSetOperation aOperation) throws Exception {
185 final String[] tableNames = getTableNames(aTables);
186 executeInTransaction(new JdbcUnitOfWork<Void>() {
187 public Void execute(Connection aConnection) throws Exception {
188 for (int i = tableNames.length - 1; i >= 0; i--) {
189 aOperation.execute(tableNames[i]);
197 * Cleans a number of database tables. This means deleting the content not
198 * dropping the tables. This may fail in case of cyclic dependencies between
199 * the tables (current limitation).
205 public void cleanDatabase(ITableFilterSimple aSelection) throws Exception {
207 final String[] tableNames = getTableNames(aSelection);
208 executeInTransaction(new JdbcUnitOfWork<Void>() {
210 public Void execute(Connection aConnection) throws Exception {
211 IDatabaseConnection connection = new DatabaseConnection(
213 ITableFilter filter = new DatabaseSequenceFilter(connection,
215 IDataSet dataset = new FilteredDataSet(filter, connection
216 .createDataSet(tableNames));
217 DatabaseOperation.DELETE_ALL.execute(connection, dataset);
225 * Executes a unit of work within a transaction.
228 * Result type of th ework.
234 public <T> T executeInTransaction(JdbcUnitOfWork<T> aWork) throws Exception {
235 Connection connection = dataSource.getConnection();
236 connection.setAutoCommit(false);
238 T value = aWork.execute(connection);
247 * Returns table names based on a table filter.
251 * @return Table names.
254 public String[] getTableNames(ITableFilterSimple aSelection)
257 List<String> result = new ArrayList<String>();
258 LOG.fine("Getting database table names to clean (schema: '" +
259 SCHEMA_PATTERN + "'");
261 Connection connection = dataSource.getConnection();
263 ResultSet tableNames = connection.getMetaData().getTables(null,
264 SCHEMA_PATTERN, "%", new String[] { "TABLE" });
265 while (tableNames.next()) {
266 String table = tableNames.getString("TABLE_NAME");
267 if (aSelection.accept(table)) {
271 return (String[]) result.toArray(new String[0]);
278 * Use {@link #cleanDatabase(ITableFilterSimple)} instead.
281 public void emptyTables(final ITableFilterSimple aSelection)
283 executeOnTables(aSelection, new TableSetOperation() {
284 public void execute(String aTable) throws Exception {
291 * User {@link #cleanDatabase(ITableFilterSimple)} instead.
294 public void emptyTable(String aTable) throws Exception {
295 executeSql("delete from " + aTable);
299 * Drops tables. This only works if there are no cyclic dependencies between
306 public void dropTables(ITableFilterSimple aTables) throws Exception {
307 final String[] tableNames = getTableNames(aTables);
308 String[] sortedTables = executeInTransaction(new JdbcUnitOfWork<String[]>() {
310 public String[] execute(Connection aConnection) throws Exception {
311 IDatabaseConnection connection = new DatabaseConnection(
313 ITableFilter filter = new DatabaseSequenceFilter(connection,
315 IDataSet dataset = new FilteredDataSet(filter, connection
316 .createDataSet(tableNames));
317 return dataset.getTableNames();
320 for (int i = sortedTables.length - 1; i >= 0; i--) {
321 dropTable(sortedTables[i]);
332 public void dropTable(final String aTable) throws Exception {
333 executeInTransaction(new JdbcUnitOfWork<Void>() {
334 public Void execute(Connection aConnection) throws Exception {
335 executeUpdate(aConnection, "drop table " + aTable);
343 * Executes an SQL statement within a transaction.
347 * @return Return code of the corresponding JDBC call.
349 public int executeSql(final String aSql) throws Exception {
350 return executeSql(aSql, new Object[0]);
354 * Executes an SQL statement within a transaction. See
355 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
356 * supported argument types.
361 * Argument of the sql statement.
362 * @return Return code of the corresponding JDBC call.
364 public int executeSql(final String aSql, final Object aArg)
366 return executeSql(aSql, new Object[] { aArg });
370 * Executes an sql statement. See
371 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
372 * supported argument types.
375 * SQL query to execute.
378 * @return Number of rows updated.
380 public int executeSql(final String aSql, final Object[] aArgs)
382 return executeInTransaction(new JdbcUnitOfWork<Integer>() {
383 public Integer execute(Connection aConnection) throws Exception {
384 PreparedStatement stmt = aConnection.prepareStatement(aSql);
385 setPreparedParams(aArgs, stmt);
386 return stmt.executeUpdate();
392 * Executes an SQL query.
396 * @return Result set.
398 public ResultSet executeQuery(Connection aConnection, String aSql) {
399 return executeQuery(aConnection, aSql, new Object[0]);
403 * Executes a query with a single argument. See
404 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
405 * supported argument types.
411 * @return Result set.
413 public ResultSet executeQuery(Connection aConnection, String aSql,
415 return executeQuery(aConnection, aSql, new Object[] { aArg });
419 * Executes a query. See
420 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
421 * supported argument types.
426 * Arguments to the query.
427 * @return Result set.
429 public ResultSet executeQuery(Connection aConnection, final String aSql,
430 final Object[] aArgs) {
432 PreparedStatement statement = aConnection.prepareStatement(aSql);
433 setPreparedParams(aArgs, statement);
435 return statement.executeQuery();
436 } catch (SQLException e) {
437 throw new RuntimeException(e);
442 * Executes an update.
449 * Arguments to the update.
450 * @return Number of rows updated.
452 public int executeUpdate(Connection aConnection, final String aSql,
453 final Object... aArgs) {
455 PreparedStatement statement = aConnection.prepareStatement(aSql);
456 setPreparedParams(aArgs, statement);
458 return statement.executeUpdate();
459 } catch (SQLException e) {
460 throw new RuntimeException(e);
465 * Sets the values of a prepared statement. See
466 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
467 * supported argument types.
470 * Arguments to the prepared statement.
473 * @throws SQLException
475 private void setPreparedParams(final Object[] aArgs,
476 PreparedStatement aStatement) throws SQLException {
477 for (int i = 1; i <= aArgs.length; i++) {
478 setPreparedParam(i, aStatement, aArgs[i - 1]);
483 * Sets a prepared statement parameter.
486 * Index of the parameter.
488 * Prepared statement.
490 * Value Must be of type Integer, Long, or String.
491 * @throws SQLException
493 private void setPreparedParam(int aIndex, PreparedStatement aStatement,
494 Object aObject) throws SQLException {
495 if (aObject instanceof Integer) {
496 aStatement.setInt(aIndex, ((Integer) aObject).intValue());
497 } else if (aObject instanceof Long) {
498 aStatement.setLong(aIndex, ((Long) aObject).longValue());
499 } else if (aObject instanceof String) {
500 aStatement.setString(aIndex, (String) aObject);
502 TestCase.fail("Unsupported object type for prepared statement: " +
503 aObject.getClass() + " value: " + aObject + " statement: " +
509 * Gets the table size.
513 * @return Table size.
514 * @throws SQLException
516 public int getTableSize(final String aTable) throws Exception {
517 return executeInTransaction(new JdbcUnitOfWork<Integer>() {
518 public Integer execute(Connection aConnection) throws Exception {
519 ResultSet resultSet = executeQuery(aConnection,
520 "select count(*) from " + aTable);
522 return resultSet.getInt(1);
529 * Counts the results in a result set.
533 * @return Number of rows in the set.
534 * @throws SQLException
536 public int countResultSet(ResultSet aResultSet) throws SQLException {
539 while (aResultSet.next()) {