wamblee.org Git - utils/blob - test/enterprise/src/main/java/org/wamblee/test/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.test.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     /**
  49      * Represents a set of tables.
  50      *
  51      * @author Erik Brakkee
  52      */
  53     public static interface TableSet {
  54         boolean contains(String aTableName);
  55     }
  56
  57     /**
  58      * Represents a unit of work (transaction).
  59      *
  60      * @author Erik Brakkee
  61      *
  62      * @param <T>
  63      *            Type of return value.
  64      */
  65     public static interface JdbcUnitOfWork<T> {
  66         /**
  67          * Executes statement within a transaction.
  68          *
  69          * @param aConnection
  70          *            Connection.
  71          * @return Result of the work.
  72          * @throws Exception
  73          */
  74         T execute(Connection aConnection) throws Exception;
  75     }
  76
  77     /**
  78      * Operation to be executed on a set of tables for each table individually.
  79      *
  80      * @author Erik Brakkee
  81      */
  82     public static interface TableSetOperation {
  83         /**
  84          * Executes on a table.
  85          *
  86          * @param aTable
  87          *            Table name.
  88          * @throws Exception
  89          */
  90         void execute(String aTable) throws Exception;
  91     }
  92
  93     private static final Logger LOG = Logger.getLogger(DatabaseUtils.class
  94         .getName());
  95
  96     /**
  97      * Schema pattern.
  98      */
  99     private static final String SCHEMA_PATTERN = "%";
 100     private DataSource dataSource;
 101
 102     private IDatabaseTester dbtester;
 103     /**
 104      * List of connections that were created for dbtesters. This list will be
 105      * closed in the {@link #stop()} method.
 106      */
 107     private List<IDatabaseConnection> connections;
 108
 109     /**
 110      * Constructs the database utils. Before use, {@link #start()} must be
 111      * called.
 112      *
 113      * @param aDataSource
 114      *            Datasource.
 115      */
 116     public DatabaseUtils(DataSource aDataSource) {
 117         dataSource = aDataSource;
 118         dbtester = new DataSourceDatabaseTester(dataSource);
 119         connections = new ArrayList<IDatabaseConnection>();
 120     }
 121
 122     /**
 123      * Starts the database utils.
 124      */
 125     public void start() {
 126         // Empty. No operation currently.
 127     }
 128
 129     /**
 130      * Stops the database utils, closing any JDBC connections that were created
 131      * by this utility. Note that connections obtained from the datasource
 132      * directly must still be closed by the user. The involved connections are
 133      * only those that are created by this utility.
 134      */
 135     public void stop() {
 136         for (IDatabaseConnection connection : connections) {
 137             try {
 138                 connection.close();
 139             } catch (SQLException e) {
 140                 LOG.log(Level.WARNING, "Could not close connection", e);
 141             }
 142         }
 143         connections.clear();
 144     }
 145
 146     /**
 147      * Creates database tester.
 148      *
 149      * @param aTables
 150      *            Tables to create the tester for.
 151      * @return Database tester.
 152      * @throws Exception
 153      */
 154     public IDatabaseTester createDbTester(ITableFilterSimple aTables)
 155         throws Exception {
 156         return createDbTester(getTableNames(aTables));
 157     }
 158
 159     /**
 160      * Creates database tester.
 161      *
 162      * @param aTables
 163      *            Tables to create the tester for.
 164      * @return Database tester.
 165      * @throws Exception
 166      */
 167     public IDatabaseTester createDbTester(String[] aTables) throws Exception {
 168         IDatabaseConnection connection = dbtester.getConnection();
 169         connections.add(connection);
 170         dbtester.setDataSet(connection.createDataSet(aTables));
 171         return dbtester;
 172     }
 173
 174     /**
 175      * Executes an operation on a set of tables.
 176      *
 177      * @param aTables
 178      *            Tables.
 179      * @param aOperation
 180      *            Operation.
 181      * @throws Exception
 182      */
 183     public void executeOnTables(ITableFilterSimple aTables,
 184         final TableSetOperation aOperation) throws Exception {
 185         final String[] tableNames = getTableNames(aTables);
 186         executeInTransaction(new JdbcUnitOfWork<Void>() {
 187             public Void execute(Connection aConnection) throws Exception {
 188                 for (int i = tableNames.length - 1; i >= 0; i--) {
 189                     aOperation.execute(tableNames[i]);
 190                 }
 191                 return null;
 192             }
 193         });
 194     }
 195
 196     /**
 197      * Cleans a number of database tables. This means deleting the content not
 198      * dropping the tables. This may fail in case of cyclic dependencies between
 199      * the tables (current limitation).
 200      *
 201      * @param aSelection
 202      *            Tables.
 203      * @throws Exception
 204      */
 205     public void cleanDatabase(ITableFilterSimple aSelection) throws Exception {
 206
 207         final String[] tableNames = getTableNames(aSelection);
 208         executeInTransaction(new JdbcUnitOfWork<Void>() {
 209
 210             public Void execute(Connection aConnection) throws Exception {
 211                 IDatabaseConnection connection = new DatabaseConnection(
 212                     aConnection);
 213                 ITableFilter filter = new DatabaseSequenceFilter(connection,
 214                     tableNames);
 215                 IDataSet dataset = new FilteredDataSet(filter, connection
 216                     .createDataSet(tableNames));
 217                 DatabaseOperation.DELETE_ALL.execute(connection, dataset);
 218                 return null;
 219             }
 220         });
 221
 222     }
 223
 224     /**
 225      * Executes a unit of work within a transaction.
 226      *
 227      * @param <T>
 228      *            Result type of th ework.
 229      * @param aWork
 230      *            Unit of work.
 231      * @return
 232      * @throws Exception
 233      */
 234     public <T> T executeInTransaction(JdbcUnitOfWork<T> aWork) throws Exception {
 235         Connection connection = dataSource.getConnection();
 236         connection.setAutoCommit(false);
 237         try {
 238             T value = aWork.execute(connection);
 239             connection.commit();
 240             return value;
 241         } finally {
 242             connection.close();
 243         }
 244     }
 245
 246     /**
 247      * Returns table names based on a table filter.
 248      *
 249      * @param aSelection
 250      *            Table filter.
 251      * @return Table names.
 252      * @throws Exception
 253      */
 254     public String[] getTableNames(ITableFilterSimple aSelection)
 255         throws Exception {
 256
 257         List<String> result = new ArrayList<String>();
 258         LOG.fine("Getting database table names to clean (schema: '" +
 259             SCHEMA_PATTERN + "'");
 260
 261         Connection connection = dataSource.getConnection();
 262         try {
 263             ResultSet tableNames = connection.getMetaData().getTables(null,
 264                 SCHEMA_PATTERN, "%", new String[] { "TABLE" });
 265             while (tableNames.next()) {
 266                 String table = tableNames.getString("TABLE_NAME");
 267                 if (aSelection.accept(table)) {
 268                     result.add(table);
 269                 }
 270             }
 271             return (String[]) result.toArray(new String[0]);
 272         } finally {
 273             connection.close();
 274         }
 275     }
 276
 277     /**
 278      * Use {@link #cleanDatabase(ITableFilterSimple)} instead.
 279      */
 280     @Deprecated
 281     public void emptyTables(final ITableFilterSimple aSelection)
 282         throws Exception {
 283         executeOnTables(aSelection, new TableSetOperation() {
 284             public void execute(String aTable) throws Exception {
 285                 emptyTable(aTable);
 286             }
 287         });
 288     }
 289
 290     /**
 291      * User {@link #cleanDatabase(ITableFilterSimple)} instead.
 292      */
 293     @Deprecated
 294     public void emptyTable(String aTable) throws Exception {
 295         executeSql("delete from " + aTable);
 296     }
 297
 298     /**
 299      * Drops tables. This only works if there are no cyclic dependencies between
 300      * the tables.
 301      *
 302      * @param aTables
 303      *            Tables to drop.
 304      * @throws Exception
 305      */
 306     public void dropTables(ITableFilterSimple aTables) throws Exception {
 307         final String[] tableNames = getTableNames(aTables);
 308         String[] sortedTables = executeInTransaction(new JdbcUnitOfWork<String[]>() {
 309
 310             public String[] execute(Connection aConnection) throws Exception {
 311                 IDatabaseConnection connection = new DatabaseConnection(
 312                     aConnection);
 313                 ITableFilter filter = new DatabaseSequenceFilter(connection,
 314                     tableNames);
 315                 IDataSet dataset = new FilteredDataSet(filter, connection
 316                     .createDataSet(tableNames));
 317                 return dataset.getTableNames();
 318             }
 319         });
 320         for (int i = sortedTables.length - 1; i >= 0; i--) {
 321             dropTable(sortedTables[i]);
 322         }
 323     }
 324
 325     /**
 326      * Drops a table.
 327      *
 328      * @param aTable
 329      *            Table to drop.
 330      * @throws Exception
 331      */
 332     public void dropTable(final String aTable) throws Exception {
 333         executeInTransaction(new JdbcUnitOfWork<Void>() {
 334             public Void execute(Connection aConnection) throws Exception {
 335                 executeUpdate(aConnection, "drop table " + aTable);
 336                 return null;
 337             }
 338         });
 339
 340     }
 341
 342     /**
 343      * Executes an SQL statement within a transaction.
 344      *
 345      * @param aSql
 346      *            SQL statement.
 347      * @return Return code of the corresponding JDBC call.
 348      */
 349     public int executeSql(final String aSql) throws Exception {
 350         return executeSql(aSql, new Object[0]);
 351     }
 352
 353     /**
 354      * Executes an SQL statement within a transaction. See
 355      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 356      * supported argument types.
 357      *
 358      * @param aSql
 359      *            SQL statement.
 360      * @param aArg
 361      *            Argument of the sql statement.
 362      * @return Return code of the corresponding JDBC call.
 363      */
 364     public int executeSql(final String aSql, final Object aArg)
 365         throws Exception {
 366         return executeSql(aSql, new Object[] { aArg });
 367     }
 368
 369     /**
 370      * Executes an sql statement. See
 371      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 372      * supported argument types.
 373      *
 374      * @param aSql
 375      *            SQL query to execute.
 376      * @param aArgs
 377      *            Arguments.
 378      * @return Number of rows updated.
 379      */
 380     public int executeSql(final String aSql, final Object[] aArgs)
 381         throws Exception {
 382         return executeInTransaction(new JdbcUnitOfWork<Integer>() {
 383             public Integer execute(Connection aConnection) throws Exception {
 384                 PreparedStatement stmt = aConnection.prepareStatement(aSql);
 385                 setPreparedParams(aArgs, stmt);
 386                 return stmt.executeUpdate();
 387             }
 388         });
 389     }
 390
 391     /**
 392      * Executes an SQL query.
 393      *
 394      * @param aSql
 395      *            Query to execute.
 396      * @return Result set.
 397      */
 398     public ResultSet executeQuery(Connection aConnection, String aSql) {
 399         return executeQuery(aConnection, aSql, new Object[0]);
 400     }
 401
 402     /**
 403      * Executes a query with a single argument. See
 404      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 405      * supported argument types.
 406      *
 407      * @param aSql
 408      *            Query.
 409      * @param aArg
 410      *            Argument.
 411      * @return Result set.
 412      */
 413     public ResultSet executeQuery(Connection aConnection, String aSql,
 414         Object aArg) {
 415         return executeQuery(aConnection, aSql, new Object[] { aArg });
 416     }
 417
 418     /**
 419      * Executes a query. See
 420      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 421      * supported argument types.
 422      *
 423      * @param aSql
 424      *            Sql query.
 425      * @param aArgs
 426      *            Arguments to the query.
 427      * @return Result set.
 428      */
 429     public ResultSet executeQuery(Connection aConnection, final String aSql,
 430         final Object[] aArgs) {
 431         try {
 432             PreparedStatement statement = aConnection.prepareStatement(aSql);
 433             setPreparedParams(aArgs, statement);
 434
 435             return statement.executeQuery();
 436         } catch (SQLException e) {
 437             throw new RuntimeException(e);
 438         }
 439     }
 440
 441     /**
 442      * Executes an update.
 443      *
 444      * @param aConnection
 445      *            Connection to use.
 446      * @param aSql
 447      *            SQL update to use.
 448      * @param aArgs
 449      *            Arguments to the update.
 450      * @return Number of rows updated.
 451      */
 452     public int executeUpdate(Connection aConnection, final String aSql,
 453         final Object... aArgs) {
 454         try {
 455             PreparedStatement statement = aConnection.prepareStatement(aSql);
 456             setPreparedParams(aArgs, statement);
 457
 458             return statement.executeUpdate();
 459         } catch (SQLException e) {
 460             throw new RuntimeException(e);
 461         }
 462     }
 463
 464     /**
 465      * Sets the values of a prepared statement. See
 466      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 467      * supported argument types.
 468      *
 469      * @param aArgs
 470      *            Arguments to the prepared statement.
 471      * @param aStatement
 472      *            Prepared statement
 473      * @throws SQLException
 474      */
 475     private void setPreparedParams(final Object[] aArgs,
 476         PreparedStatement aStatement) throws SQLException {
 477         for (int i = 1; i <= aArgs.length; i++) {
 478             setPreparedParam(i, aStatement, aArgs[i - 1]);
 479         }
 480     }
 481
 482     /**
 483      * Sets a prepared statement parameter.
 484      *
 485      * @param aIndex
 486      *            Index of the parameter.
 487      * @param aStatement
 488      *            Prepared statement.
 489      * @param aObject
 490      *            Value Must be of type Integer, Long, or String.
 491      * @throws SQLException
 492      */
 493     private void setPreparedParam(int aIndex, PreparedStatement aStatement,
 494         Object aObject) throws SQLException {
 495         if (aObject instanceof Integer) {
 496             aStatement.setInt(aIndex, ((Integer) aObject).intValue());
 497         } else if (aObject instanceof Long) {
 498             aStatement.setLong(aIndex, ((Long) aObject).longValue());
 499         } else if (aObject instanceof String) {
 500             aStatement.setString(aIndex, (String) aObject);
 501         } else {
 502             TestCase.fail("Unsupported object type for prepared statement: " +
 503                 aObject.getClass() + " value: " + aObject + " statement: " +
 504                 aStatement);
 505         }
 506     }
 507
 508     /**
 509      * Gets the table size.
 510      *
 511      * @param aTable
 512      *            Table.
 513      * @return Table size.
 514      * @throws SQLException
 515      */
 516     public int getTableSize(final String aTable) throws Exception {
 517         return executeInTransaction(new JdbcUnitOfWork<Integer>() {
 518             public Integer execute(Connection aConnection) throws Exception {
 519                 ResultSet resultSet = executeQuery(aConnection,
 520                     "select count(*) from " + aTable);
 521                 resultSet.next();
 522                 return resultSet.getInt(1);
 523             }
 524         });
 525
 526     }
 527
 528     /**
 529      * Counts the results in a result set.
 530      *
 531      * @param aResultSet
 532      *            Resultset.
 533      * @return Number of rows in the set.
 534      * @throws SQLException
 535      */
 536     public int countResultSet(ResultSet aResultSet) throws SQLException {
 537         int count = 0;
 538
 539         while (aResultSet.next()) {
 540             count++;
 541         }
 542
 543         return count;
 544     }
 545
 546 }