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