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