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