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 {
48 public static interface TableSet {
49 boolean contains(String aTableName);
52 public static interface JdbcUnitOfWork<T> {
53 T execute(Connection aConnection) throws Exception;
56 public static interface TableSetOperation {
57 void execute(String aTable) throws Exception;
60 private static final Logger LOG = Logger.getLogger(DatabaseUtils.class
66 private static final String SCHEMA_PATTERN = "%";
67 private DataSource dataSource;
68 private ITableFilterSimple tables;
70 private IDatabaseTester dbtester;
72 * List of connections that were created for dbtesters.
73 * This list will be closed in the {@link #stop()} method.
75 private List<IDatabaseConnection> connections;
78 * Constructs the database utils.
79 * Before use, {@link #start()} must be called.
80 * @param aDataSource Datasource.
81 * @param aTables Tables to work with.
83 public DatabaseUtils(DataSource aDataSource, ITableFilterSimple aTables) {
84 dataSource = aDataSource;
86 dbtester = new DataSourceDatabaseTester(dataSource);
87 connections = new ArrayList<IDatabaseConnection>();
91 * Starts the database utils.
94 // Empty. No operation currently.
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.
104 for (IDatabaseConnection connection: connections) {
107 } catch (SQLException e) {
108 LOG.log(Level.WARNING, "Could not close connection", e);
113 public IDatabaseTester createDbTester() throws Exception {
114 return createDbTester(getTableNames(tables));
117 public IDatabaseTester createDbTester(String[] aTables) throws Exception {
118 IDatabaseConnection connection = dbtester.getConnection();
119 connections.add(connection);
120 dbtester.setDataSet(connection.createDataSet(aTables));
124 public void cleanDatabase() throws Exception {
125 cleanDatabase(tables);
128 public void executeOnTables(ITableFilterSimple aTables,
129 final TableSetOperation aOperation) throws Exception {
130 final String[] tables = getTableNames(aTables);
131 executeInTransaction(new JdbcUnitOfWork<Void>() {
132 public Void execute(Connection aConnection) throws Exception {
133 for (int i = tables.length - 1; i >= 0; i--) {
134 aOperation.execute(tables[i]);
141 public void cleanDatabase(ITableFilterSimple aSelection) throws Exception {
143 final String[] tables = getTableNames(aSelection);
144 executeInTransaction(new JdbcUnitOfWork<Void>() {
146 public Void execute(Connection aConnection) throws Exception {
147 IDatabaseConnection connection = new DatabaseConnection(
149 ITableFilter filter = new DatabaseSequenceFilter(connection,
151 IDataSet dataset = new FilteredDataSet(filter, connection
152 .createDataSet(tables));
153 DatabaseOperation.DELETE_ALL.execute(connection, dataset);
160 public <T> T executeInTransaction(JdbcUnitOfWork<T> aCallback)
162 Connection connection = dataSource.getConnection();
163 connection.setAutoCommit(false);
165 T value = aCallback.execute(connection);
173 public String[] getTableNames() throws Exception {
174 return getTableNames(tables);
178 * @throws SQLException
180 public String[] getTableNames(ITableFilterSimple aSelection)
183 List<String> result = new ArrayList<String>();
184 LOG.fine("Getting database table names to clean (schema: '" +
185 SCHEMA_PATTERN + "'");
187 Connection connection = dataSource.getConnection();
189 ResultSet tables = connection.getMetaData().getTables(null,
190 SCHEMA_PATTERN, "%", new String[] { "TABLE" });
191 while (tables.next()) {
192 String table = tables.getString("TABLE_NAME");
193 if (aSelection.accept(table)) {
197 return (String[]) result.toArray(new String[0]);
203 public void emptyTables() throws Exception {
204 executeOnTables(tables, new TableSetOperation() {
205 public void execute(String aTable) throws Exception {
213 * @throws SQLException
215 public void emptyTables(final ITableFilterSimple aSelection)
217 executeOnTables(aSelection, new TableSetOperation() {
218 public void execute(String aTable) throws Exception {
226 * @throws SQLException
228 public void emptyTable(String aTable) throws Exception {
229 executeSql("delete from " + aTable);
232 public void dropTables() throws Exception {
236 public void dropTables(ITableFilterSimple aTables) throws Exception {
237 final String[] tables = getTableNames(aTables);
238 String[] sortedTables = executeInTransaction(new JdbcUnitOfWork<String[]>() {
240 public String[] execute(Connection aConnection) throws Exception {
241 IDatabaseConnection connection = new DatabaseConnection(
243 ITableFilter filter = new DatabaseSequenceFilter(connection,
245 IDataSet dataset = new FilteredDataSet(filter, connection
246 .createDataSet(tables));
247 return dataset.getTableNames();
250 for (int i = sortedTables.length - 1; i >= 0; i--) {
251 dropTable(sortedTables[i]);
257 * @throws SQLException
259 public void dropTable(final String aTable) throws Exception {
260 executeInTransaction(new JdbcUnitOfWork<Void>() {
261 public Void execute(Connection aConnection) throws Exception {
262 executeUpdate(aConnection, "drop table " + aTable);
270 * Executes an SQL statement within a transaction.
274 * @return Return code of the corresponding JDBC call.
276 public int executeSql(final String aSql) throws Exception {
277 return executeSql(aSql, new Object[0]);
281 * Executes an SQL statement within a transaction. See
282 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
283 * supported argument types.
288 * Argument of the sql statement.
289 * @return Return code of the corresponding JDBC call.
291 public int executeSql(final String aSql, final Object aArg)
293 return executeSql(aSql, new Object[] { aArg });
297 * Executes an sql statement. See
298 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
299 * supported argument types.
302 * SQL query to execute.
305 * @return Number of rows updated.
307 public int executeSql(final String aSql, final Object[] aArgs)
309 return executeInTransaction(new JdbcUnitOfWork<Integer>() {
310 public Integer execute(Connection aConnection) throws Exception {
311 PreparedStatement stmt = aConnection.prepareStatement(aSql);
312 setPreparedParams(aArgs, stmt);
313 return stmt.executeUpdate();
319 * Executes an SQL query.
323 * @return Result set.
325 public ResultSet executeQuery(Connection aConnection, String aSql) {
326 return executeQuery(aConnection, aSql, new Object[0]);
330 * Executes a query with a single argument. See
331 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
332 * supported argument types.
338 * @return Result set.
340 public ResultSet executeQuery(Connection aConnection, String aSql,
342 return executeQuery(aConnection, aSql, new Object[] { aArg });
346 * Executes a query. See
347 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
348 * supported argument types.
353 * Arguments to the query.
354 * @return Result set.
356 public ResultSet executeQuery(Connection aConnection, final String aSql,
357 final Object[] aArgs) {
359 PreparedStatement statement = aConnection.prepareStatement(aSql);
360 setPreparedParams(aArgs, statement);
362 return statement.executeQuery();
363 } catch (SQLException e) {
364 throw new RuntimeException(e);
368 public int executeUpdate(Connection aConnection, final String aSql,
369 final Object... aArgs) {
371 PreparedStatement statement = aConnection.prepareStatement(aSql);
372 setPreparedParams(aArgs, statement);
374 return statement.executeUpdate();
375 } catch (SQLException e) {
376 throw new RuntimeException(e);
381 * Sets the values of a prepared statement. See
382 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
383 * supported argument types.
386 * Arguments to the prepared statement.
389 * @throws SQLException
391 private void setPreparedParams(final Object[] aArgs,
392 PreparedStatement aStatement) throws SQLException {
393 for (int i = 1; i <= aArgs.length; i++) {
394 setPreparedParam(i, aStatement, aArgs[i - 1]);
399 * Sets a prepared statement parameter.
402 * Index of the parameter.
404 * Prepared statement.
406 * Value Must be of type Integer, Long, or String.
407 * @throws SQLException
409 private void setPreparedParam(int aIndex, PreparedStatement aStatement,
410 Object aObject) throws SQLException {
411 if (aObject instanceof Integer) {
412 aStatement.setInt(aIndex, ((Integer) aObject).intValue());
413 } else if (aObject instanceof Long) {
414 aStatement.setLong(aIndex, ((Long) aObject).longValue());
415 } else if (aObject instanceof String) {
416 aStatement.setString(aIndex, (String) aObject);
418 TestCase.fail("Unsupported object type for prepared statement: " +
419 aObject.getClass() + " value: " + aObject + " statement: " +
426 * @throws SQLException
428 public int getTableSize(final String aTable) throws Exception {
429 return executeInTransaction(new JdbcUnitOfWork<Integer>() {
430 public Integer execute(Connection aConnection) throws Exception {
431 ResultSet resultSet = executeQuery(aConnection,
432 "select count(*) from " + aTable);
434 return resultSet.getInt(1);
440 public int countResultSet(ResultSet aResultSet) throws SQLException {
443 while (aResultSet.next()) {