wamblee.org Git - utils/blob - test/enterprise/src/main/java/org/wamblee/support/persistence/DatabaseUtils.java

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