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