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