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.Logger;
26 import javax.sql.DataSource;
28 import junit.framework.TestCase;
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;
42 * Database utilities is a simple support class for common tasks in working with
45 public class DatabaseUtils {
47 public static interface TableSet {
48 boolean contains(String aTableName);
51 public static interface JdbcUnitOfWork<T> {
52 T execute(Connection aConnection) throws Exception;
55 public static interface TableSetOperation {
56 void execute(String aTable) throws Exception;
59 private static final Logger LOG = Logger.getLogger(DatabaseUtils.class
65 private static final String SCHEMA_PATTERN = "%";
66 private DataSource dataSource;
67 private ITableFilterSimple tables;
69 public DatabaseUtils(DataSource aDataSource, ITableFilterSimple aTables) {
70 dataSource = aDataSource;
74 public IDatabaseTester createDbTester() throws Exception {
75 return createDbTester(getTableNames(tables));
78 public IDatabaseTester createDbTester(String[] aTables) throws Exception {
79 IDatabaseTester dbtester = new DataSourceDatabaseTester(dataSource);
80 dbtester.setDataSet(dbtester.getConnection().createDataSet(aTables));
84 public void cleanDatabase() throws Exception {
85 cleanDatabase(tables);
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]);
101 public void cleanDatabase(ITableFilterSimple aSelection) throws Exception {
103 final String[] tables = getTableNames(aSelection);
104 executeInTransaction(new JdbcUnitOfWork<Void>() {
106 public Void execute(Connection aConnection) throws Exception {
107 IDatabaseConnection connection = new DatabaseConnection(
109 ITableFilter filter = new DatabaseSequenceFilter(connection,
111 IDataSet dataset = new FilteredDataSet(filter, connection
112 .createDataSet(tables));
113 DatabaseOperation.DELETE_ALL.execute(connection, dataset);
120 public <T> T executeInTransaction(JdbcUnitOfWork<T> aCallback)
122 Connection connection = dataSource.getConnection();
123 connection.setAutoCommit(false);
125 T value = aCallback.execute(connection);
133 public String[] getTableNames() throws Exception {
134 return getTableNames(tables);
138 * @throws SQLException
140 public String[] getTableNames(ITableFilterSimple aSelection)
143 List<String> result = new ArrayList<String>();
144 LOG.fine("Getting database table names to clean (schema: '" +
145 SCHEMA_PATTERN + "'");
147 Connection connection = dataSource.getConnection();
149 ResultSet tables = connection.getMetaData().getTables(null,
150 SCHEMA_PATTERN, "%", new String[] { "TABLE" });
151 while (tables.next()) {
152 String table = tables.getString("TABLE_NAME");
153 if (aSelection.accept(table)) {
157 return (String[]) result.toArray(new String[0]);
163 public void emptyTables() throws Exception {
164 executeOnTables(tables, new TableSetOperation() {
165 public void execute(String aTable) throws Exception {
173 * @throws SQLException
175 public void emptyTables(final ITableFilterSimple aSelection)
177 executeOnTables(aSelection, new TableSetOperation() {
178 public void execute(String aTable) throws Exception {
186 * @throws SQLException
188 public void emptyTable(String aTable) throws Exception {
189 executeSql("delete from " + aTable);
192 public void dropTables() throws Exception {
196 public void dropTables(ITableFilterSimple aTables) throws Exception {
197 final String[] tables = getTableNames(aTables);
198 String[] sortedTables = executeInTransaction(new JdbcUnitOfWork<String[]>() {
200 public String[] execute(Connection aConnection) throws Exception {
201 IDatabaseConnection connection = new DatabaseConnection(
203 ITableFilter filter = new DatabaseSequenceFilter(connection,
205 IDataSet dataset = new FilteredDataSet(filter, connection
206 .createDataSet(tables));
207 return dataset.getTableNames();
210 for (int i = sortedTables.length-1; i >= 0; i--) {
211 dropTable(sortedTables[i]);
217 * @throws SQLException
219 public void dropTable(final String aTable) throws Exception {
220 executeInTransaction(new JdbcUnitOfWork<Void>() {
221 public Void execute(Connection aConnection) throws Exception {
222 executeUpdate(aConnection, "drop table " + aTable);
230 * Executes an SQL statement within a transaction.
234 * @return Return code of the corresponding JDBC call.
236 public int executeSql(final String aSql) throws Exception {
237 return executeSql(aSql, new Object[0]);
241 * Executes an SQL statement within a transaction. See
242 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
243 * supported argument types.
248 * Argument of the sql statement.
249 * @return Return code of the corresponding JDBC call.
251 public int executeSql(final String aSql, final Object aArg)
253 return executeSql(aSql, new Object[] { aArg });
257 * Executes an sql statement. See
258 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
259 * supported argument types.
262 * SQL query to execute.
265 * @return Number of rows updated.
267 public int executeSql(final String aSql, final Object[] aArgs)
269 return executeInTransaction(new JdbcUnitOfWork<Integer>() {
270 public Integer execute(Connection aConnection) throws Exception {
271 PreparedStatement stmt = aConnection.prepareStatement(aSql);
272 setPreparedParams(aArgs, stmt);
273 return stmt.executeUpdate();
279 * Executes an SQL query.
283 * @return Result set.
285 public ResultSet executeQuery(Connection aConnection, String aSql) {
286 return executeQuery(aConnection, aSql, new Object[0]);
290 * Executes a query with a single argument. See
291 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
292 * supported argument types.
298 * @return Result set.
300 public ResultSet executeQuery(Connection aConnection, String aSql,
302 return executeQuery(aConnection, aSql, new Object[] { aArg });
306 * Executes a query. See
307 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
308 * supported argument types.
313 * Arguments to the query.
314 * @return Result set.
316 public ResultSet executeQuery(Connection aConnection, final String aSql,
317 final Object[] aArgs) {
319 PreparedStatement statement = aConnection.prepareStatement(aSql);
320 setPreparedParams(aArgs, statement);
322 return statement.executeQuery();
323 } catch (SQLException e) {
324 throw new RuntimeException(e);
328 public int executeUpdate(Connection aConnection, final String aSql,
329 final Object... aArgs) {
331 PreparedStatement statement = aConnection.prepareStatement(aSql);
332 setPreparedParams(aArgs, statement);
334 return statement.executeUpdate();
335 } catch (SQLException e) {
336 throw new RuntimeException(e);
341 * Sets the values of a prepared statement. See
342 * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
343 * supported argument types.
346 * Arguments to the prepared statement.
349 * @throws SQLException
351 private void setPreparedParams(final Object[] aArgs,
352 PreparedStatement aStatement) throws SQLException {
353 for (int i = 1; i <= aArgs.length; i++) {
354 setPreparedParam(i, aStatement, aArgs[i - 1]);
359 * Sets a prepared statement parameter.
362 * Index of the parameter.
364 * Prepared statement.
366 * Value Must be of type Integer, Long, or String.
367 * @throws SQLException
369 private void setPreparedParam(int aIndex, PreparedStatement aStatement,
370 Object aObject) throws SQLException {
371 if (aObject instanceof Integer) {
372 aStatement.setInt(aIndex, ((Integer) aObject).intValue());
373 } else if (aObject instanceof Long) {
374 aStatement.setLong(aIndex, ((Long) aObject).longValue());
375 } else if (aObject instanceof String) {
376 aStatement.setString(aIndex, (String) aObject);
378 TestCase.fail("Unsupported object type for prepared statement: " +
379 aObject.getClass() + " value: " + aObject + " statement: " +
386 * @throws SQLException
388 public int getTableSize(final String aTable) throws Exception {
389 return executeInTransaction(new JdbcUnitOfWork<Integer>() {
390 public Integer execute(Connection aConnection) throws Exception {
391 ResultSet resultSet = executeQuery(aConnection,
392 "select count(*) from " + aTable);
394 return resultSet.getInt(1);
400 public int countResultSet(ResultSet aResultSet) throws SQLException {
403 while (aResultSet.next()) {