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         Connection connection = dataSource.getConnection();
 150         try {
 151             ResultSet tables = connection.getMetaData().getTables(null,
 152                 SCHEMA_PATTERN, "%", new String[] { "TABLE" });
 153             while (tables.next()) {
 154                 String table = tables.getString("TABLE_NAME");
 155                 if (aSelection.accept(table)) {
 156                     result.add(table);
 157                 }
 158             }
 159             return (String[]) result.toArray(new String[0]);
 160         } finally {
 161             connection.close();
 162         }
 163     }
 164
 165     public void emptyTables() throws Exception {
 166         executeOnTables(tables, new TableSetOperation() {
 167             public void execute(String aTable) throws Exception {
 168                 emptyTable(aTable);
 169             }
 170         });
 171     }
 172
 173     /**
 174      * @return
 175      * @throws SQLException
 176      */
 177     public void emptyTables(final ITableFilterSimple aSelection)
 178         throws Exception {
 179         executeOnTables(aSelection, new TableSetOperation() {
 180             public void execute(String aTable) throws Exception {
 181                 emptyTable(aTable);
 182             }
 183         });
 184     }
 185
 186     /**
 187      * @return
 188      * @throws SQLException
 189      */
 190     public void emptyTable(String aTable) throws Exception {
 191         executeSql("delete from " + aTable);
 192     }
 193
 194     public void dropTables() throws Exception {
 195         executeOnTables(tables, new TableSetOperation() {
 196
 197             public void execute(String aTable) throws Exception {
 198                 dropTable(aTable);
 199             }
 200         });
 201     }
 202
 203     public void dropTables(ITableFilterSimple aTables) throws Exception {
 204         executeOnTables(aTables, new TableSetOperation() {
 205
 206             public void execute(String aTable) throws Exception {
 207                 dropTable(aTable);
 208             }
 209         });
 210     }
 211
 212     /**
 213      * @return
 214      * @throws SQLException
 215      */
 216     public void dropTable(final String aTable) throws Exception {
 217         executeInTransaction(new JdbcUnitOfWork<Void>() {
 218             public Void execute(Connection aConnection) throws Exception {
 219                 executeUpdate(aConnection, "drop table " + aTable);
 220                 return null;
 221             }
 222         });
 223
 224     }
 225
 226     /**
 227      * Executes an SQL statement within a transaction.
 228      *
 229      * @param aSql
 230      *            SQL statement.
 231      * @return Return code of the corresponding JDBC call.
 232      */
 233     public int executeSql(final String aSql) throws Exception {
 234         return executeSql(aSql, new Object[0]);
 235     }
 236
 237     /**
 238      * Executes an SQL statement within a transaction. See
 239      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 240      * supported argument types.
 241      *
 242      * @param aSql
 243      *            SQL statement.
 244      * @param aArg
 245      *            Argument of the sql statement.
 246      * @return Return code of the corresponding JDBC call.
 247      */
 248     public int executeSql(final String aSql, final Object aArg)
 249         throws Exception {
 250         return executeSql(aSql, new Object[] { aArg });
 251     }
 252
 253     /**
 254      * Executes an sql statement. See
 255      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 256      * supported argument types.
 257      *
 258      * @param aSql
 259      *            SQL query to execute.
 260      * @param aArgs
 261      *            Arguments.
 262      * @return Number of rows updated.
 263      */
 264     public int executeSql(final String aSql, final Object[] aArgs)
 265         throws Exception {
 266         return executeInTransaction(new JdbcUnitOfWork<Integer>() {
 267             public Integer execute(Connection aConnection) throws Exception {
 268                 PreparedStatement stmt = aConnection.prepareStatement(aSql);
 269                 setPreparedParams(aArgs, stmt);
 270                 return stmt.executeUpdate();
 271             }
 272         });
 273     }
 274
 275     /**
 276      * Executes an SQL query.
 277      *
 278      * @param aSql
 279      *            Query to execute.
 280      * @return Result set.
 281      */
 282     public ResultSet executeQuery(Connection aConnection, String aSql) {
 283         return executeQuery(aConnection, aSql, new Object[0]);
 284     }
 285
 286     /**
 287      * Executes a query with a single argument. See
 288      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 289      * supported argument types.
 290      *
 291      * @param aSql
 292      *            Query.
 293      * @param aArg
 294      *            Argument.
 295      * @return Result set.
 296      */
 297     public ResultSet executeQuery(Connection aConnection, String aSql,
 298         Object aArg) {
 299         return executeQuery(aConnection, aSql, new Object[] { aArg });
 300     }
 301
 302     /**
 303      * Executes a query. See
 304      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 305      * supported argument types.
 306      *
 307      * @param aSql
 308      *            Sql query.
 309      * @param aArgs
 310      *            Arguments to the query.
 311      * @return Result set.
 312      */
 313     public ResultSet executeQuery(Connection aConnection, final String aSql,
 314         final Object[] aArgs) {
 315         try {
 316             PreparedStatement statement = aConnection.prepareStatement(aSql);
 317             setPreparedParams(aArgs, statement);
 318
 319             return statement.executeQuery();
 320         } catch (SQLException e) {
 321             throw new RuntimeException(e);
 322         }
 323     }
 324
 325     public int executeUpdate(Connection aConnection, final String aSql,
 326         final Object... aArgs) {
 327         try {
 328             PreparedStatement statement = aConnection.prepareStatement(aSql);
 329             setPreparedParams(aArgs, statement);
 330
 331             return statement.executeUpdate();
 332         } catch (SQLException e) {
 333             throw new RuntimeException(e);
 334         }
 335     }
 336
 337     /**
 338      * Sets the values of a prepared statement. See
 339      * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
 340      * supported argument types.
 341      *
 342      * @param aArgs
 343      *            Arguments to the prepared statement.
 344      * @param aStatement
 345      *            Prepared statement
 346      * @throws SQLException
 347      */
 348     private void setPreparedParams(final Object[] aArgs,
 349         PreparedStatement aStatement) throws SQLException {
 350         for (int i = 1; i <= aArgs.length; i++) {
 351             setPreparedParam(i, aStatement, aArgs[i - 1]);
 352         }
 353     }
 354
 355     /**
 356      * Sets a prepared statement parameter.
 357      *
 358      * @param aIndex
 359      *            Index of the parameter.
 360      * @param aStatement
 361      *            Prepared statement.
 362      * @param aObject
 363      *            Value Must be of type Integer, Long, or String.
 364      * @throws SQLException
 365      */
 366     private void setPreparedParam(int aIndex, PreparedStatement aStatement,
 367         Object aObject) throws SQLException {
 368         if (aObject instanceof Integer) {
 369             aStatement.setInt(aIndex, ((Integer) aObject).intValue());
 370         } else if (aObject instanceof Long) {
 371             aStatement.setLong(aIndex, ((Long) aObject).longValue());
 372         } else if (aObject instanceof String) {
 373             aStatement.setString(aIndex, (String) aObject);
 374         } else {
 375             TestCase.fail("Unsupported object type for prepared statement: " +
 376                 aObject.getClass() + " value: " + aObject + " statement: " +
 377                 aStatement);
 378         }
 379     }
 380
 381     /**
 382      * @return
 383      * @throws SQLException
 384      */
 385     public int getTableSize(final String aTable) throws Exception {
 386         return executeInTransaction(new JdbcUnitOfWork<Integer>() {
 387             public Integer execute(Connection aConnection) throws Exception {
 388                 ResultSet resultSet = executeQuery(aConnection,
 389                     "select count(*) from " + aTable);
 390                 resultSet.next();
 391                 return resultSet.getInt(1);
 392             }
 393         });
 394
 395     }
 396
 397     public int countResultSet(ResultSet aResultSet) throws SQLException {
 398         int count = 0;
 399
 400         while (aResultSet.next()) {
 401             count++;
 402         }
 403
 404         return count;
 405     }
 406
 407 }