wamblee.org Git - utils/blob - 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.Logger;
  25
  26 import javax.sql.DataSource;
  27
  28 import junit.framework.TestCase;
  29
  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;
  40
  41 /**
  42  * Database utilities is a simple support class for common tasks in working with
  43  * databases.
  44  */
  45 public class DatabaseUtils {
  46
  47     public static interface TableSet {
  48         boolean contains(String aTableName);
  49     }
  50
  51     public static interface JdbcUnitOfWork<T> {
  52         T execute(Connection aConnection) throws Exception;
  53     }
  54
  55     public static interface TableSetOperation {
  56         void execute(String aTable) throws Exception;
  57     }
  58
  59     private static final Logger LOG = Logger.getLogger(DatabaseUtils.class
  60         .getName());
  61
  62     /**
  63      * Schema pattern.
  64      */
  65     private static final String SCHEMA_PATTERN = "%";
  66     private DataSource dataSource;
  67     private ITableFilterSimple tables;
  68
  69     public DatabaseUtils(DataSource aDataSource, ITableFilterSimple aTables) {
  70         dataSource = aDataSource;
  71         tables = aTables;
  72     }
  73
  74     public IDatabaseTester createDbTester() throws Exception {
  75         return createDbTester(getTableNames(tables));
  76     }
  77
  78     public IDatabaseTester createDbTester(String[] aTables) throws Exception {
  79         IDatabaseTester dbtester = new DataSourceDatabaseTester(dataSource);
  80         dbtester.setDataSet(dbtester.getConnection().createDataSet(aTables));
  81         return dbtester;
  82     }
  83
  84     public void cleanDatabase() throws Exception {
  85         cleanDatabase(tables);
  86     }
  87
  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]);
  95                 }
  96                 return null;
  97             }
  98         });
  99     }
 100
 101     public void cleanDatabase(ITableFilterSimple aSelection) throws Exception {
 102
 103         final String[] tables = getTableNames(aSelection);
 104         executeInTransaction(new JdbcUnitOfWork<Void>() {
 105
 106             public Void execute(Connection aConnection) throws Exception {
 107                 IDatabaseConnection connection = new DatabaseConnection(
 108                     aConnection);
 109                 ITableFilter filter = new DatabaseSequenceFilter(connection,
 110                     tables);
 111                 IDataSet dataset = new FilteredDataSet(filter, connection
 112                     .createDataSet(tables));
 113                 DatabaseOperation.DELETE_ALL.execute(connection, dataset);
 114                 return null;
 115             }
 116         });
 117
 118     }
 119
 120     public <T> T executeInTransaction(JdbcUnitOfWork<T> aCallback)
 121         throws Exception {
 122         Connection connection = dataSource.getConnection();
 123         connection.setAutoCommit(false);
 124         try {
 125             T value = aCallback.execute(connection);
 126             connection.commit();
 127             return value;
 128         } finally {
 129             connection.close();
 130         }
 131     }
 132
 133     public String[] getTableNames() throws Exception {
 134         return getTableNames(tables);
 135     }
 136
 137     /**
 138      * @throws SQLException
 139      */
 140     public String[] getTableNames(ITableFilterSimple aSelection)
 141         throws Exception {
 142
 143         List<String> result = new ArrayList<String>();
 144         LOG.fine("Getting database table names to clean (schema: '" +
 145             SCHEMA_PATTERN + "'");
 146
 147         Connection connection = dataSource.getConnection();
 148         try {
 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)) {
 154                     result.add(table);
 155                 }
 156             }
 157             return (String[]) result.toArray(new String[0]);
 158         } finally {
 159             connection.close();
 160         }
 161     }
 162
 163     public void emptyTables() throws Exception {
 164         executeOnTables(tables, new TableSetOperation() {
 165             public void execute(String aTable) throws Exception {
 166                 emptyTable(aTable);
 167             }
 168         });
 169     }
 170
 171     /**
 172      * @return
 173      * @throws SQLException
 174      */
 175     public void emptyTables(final ITableFilterSimple aSelection)
 176         throws Exception {
 177         executeOnTables(aSelection, new TableSetOperation() {
 178             public void execute(String aTable) throws Exception {
 179                 emptyTable(aTable);
 180             }
 181         });
 182     }
 183
 184     /**
 185      * @return
 186      * @throws SQLException
 187      */
 188     public void emptyTable(String aTable) throws Exception {
 189         executeSql("delete from " + aTable);
 190     }
 191
 192     public void dropTables() throws Exception {
 193         dropTables(tables);
 194     }
 195
 196     public void dropTables(ITableFilterSimple aTables) throws Exception {
 197         final String[] tables = getTableNames(aTables);
 198         String[] sortedTables = executeInTransaction(new JdbcUnitOfWork<String[]>() {
 199
 200             public String[] execute(Connection aConnection) throws Exception {
 201                 IDatabaseConnection connection = new DatabaseConnection(
 202                     aConnection);
 203                 ITableFilter filter = new DatabaseSequenceFilter(connection,
 204                     tables);
 205                 IDataSet dataset = new FilteredDataSet(filter, connection
 206                     .createDataSet(tables));
 207                 return dataset.getTableNames();
 208             }
 209         });
 210         for (int i = sortedTables.length-1; i >= 0; i--) {
 211             dropTable(sortedTables[i]);
 212         }
 213     }
 214
 215     /**
 216      * @return
 217      * @throws SQLException
 218      */
 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);
 223                 return null;
 224             }
 225         });
 226
 227     }
 228
 229     /**
 230      * Executes an SQL statement within a transaction.
 231      *
 232      * @param aSql
 233      *            SQL statement.
 234      * @return Return code of the corresponding JDBC call.
 235      */
 236     public int executeSql(final String aSql) throws Exception {
 237         return executeSql(aSql, new Object[0]);
 238     }
 239
 240     /**
 241      * Executes an SQL statement within a transaction. See
 242      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 243      * supported argument types.
 244      *
 245      * @param aSql
 246      *            SQL statement.
 247      * @param aArg
 248      *            Argument of the sql statement.
 249      * @return Return code of the corresponding JDBC call.
 250      */
 251     public int executeSql(final String aSql, final Object aArg)
 252         throws Exception {
 253         return executeSql(aSql, new Object[] { aArg });
 254     }
 255
 256     /**
 257      * Executes an sql statement. See
 258      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 259      * supported argument types.
 260      *
 261      * @param aSql
 262      *            SQL query to execute.
 263      * @param aArgs
 264      *            Arguments.
 265      * @return Number of rows updated.
 266      */
 267     public int executeSql(final String aSql, final Object[] aArgs)
 268         throws Exception {
 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();
 274             }
 275         });
 276     }
 277
 278     /**
 279      * Executes an SQL query.
 280      *
 281      * @param aSql
 282      *            Query to execute.
 283      * @return Result set.
 284      */
 285     public ResultSet executeQuery(Connection aConnection, String aSql) {
 286         return executeQuery(aConnection, aSql, new Object[0]);
 287     }
 288
 289     /**
 290      * Executes a query with a single argument. See
 291      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 292      * supported argument types.
 293      *
 294      * @param aSql
 295      *            Query.
 296      * @param aArg
 297      *            Argument.
 298      * @return Result set.
 299      */
 300     public ResultSet executeQuery(Connection aConnection, String aSql,
 301         Object aArg) {
 302         return executeQuery(aConnection, aSql, new Object[] { aArg });
 303     }
 304
 305     /**
 306      * Executes a query. See
 307      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 308      * supported argument types.
 309      *
 310      * @param aSql
 311      *            Sql query.
 312      * @param aArgs
 313      *            Arguments to the query.
 314      * @return Result set.
 315      */
 316     public ResultSet executeQuery(Connection aConnection, final String aSql,
 317         final Object[] aArgs) {
 318         try {
 319             PreparedStatement statement = aConnection.prepareStatement(aSql);
 320             setPreparedParams(aArgs, statement);
 321
 322             return statement.executeQuery();
 323         } catch (SQLException e) {
 324             throw new RuntimeException(e);
 325         }
 326     }
 327
 328     public int executeUpdate(Connection aConnection, final String aSql,
 329         final Object... aArgs) {
 330         try {
 331             PreparedStatement statement = aConnection.prepareStatement(aSql);
 332             setPreparedParams(aArgs, statement);
 333
 334             return statement.executeUpdate();
 335         } catch (SQLException e) {
 336             throw new RuntimeException(e);
 337         }
 338     }
 339
 340     /**
 341      * Sets the values of a prepared statement. See
 342      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 343      * supported argument types.
 344      *
 345      * @param aArgs
 346      *            Arguments to the prepared statement.
 347      * @param aStatement
 348      *            Prepared statement
 349      * @throws SQLException
 350      */
 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]);
 355         }
 356     }
 357
 358     /**
 359      * Sets a prepared statement parameter.
 360      *
 361      * @param aIndex
 362      *            Index of the parameter.
 363      * @param aStatement
 364      *            Prepared statement.
 365      * @param aObject
 366      *            Value Must be of type Integer, Long, or String.
 367      * @throws SQLException
 368      */
 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);
 377         } else {
 378             TestCase.fail("Unsupported object type for prepared statement: " +
 379                 aObject.getClass() + " value: " + aObject + " statement: " +
 380                 aStatement);
 381         }
 382     }
 383
 384     /**
 385      * @return
 386      * @throws SQLException
 387      */
 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);
 393                 resultSet.next();
 394                 return resultSet.getInt(1);
 395             }
 396         });
 397
 398     }
 399
 400     public int countResultSet(ResultSet aResultSet) throws SQLException {
 401         int count = 0;
 402
 403         while (aResultSet.next()) {
 404             count++;
 405         }
 406
 407         return count;
 408     }
 409
 410 }