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