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