DBUnit 2.4.9+ 支持数据库唯一键扩展

深藏阁楼爱情的钟 2022-05-21 00:00 179阅读 0赞

由于DBUnit在做refresh或者delete操作时是根据数据库的connection获取到meta信息，然后通过meta信息来确定主键列，最后根据这些列构造where条件，进行更新或者删除处理，但是如果表没有定义主键，只定义了唯一键，DBUnit无法处理，而目前我所遇到的就有这种情况，所以我对DBUnit进行了扩展，支持oracle和mysql的唯一键处理，但是要注意唯一键的键值如果有null值，则无法处理，这是一个局限，后续我会做修复。

1、继承DefaultMetadataHandler覆盖getPrimaryKeys方法，如果有主键，则直接用主键，没有就用唯一键，

直接上代码：

/**
 * 由于DBUNIT对于没有主键的表不支持，更新，单条删除等，这里做一下扩展
 * 支持oracle/mysql的唯一键更新
 * 不过注意一点局限，如果唯一键中有某个字段为null,那么可能找不到这条记录
 *
 * @author qingxu
 * @version $$Id: UniquekeyMetadataHandler.java, v 0.1 12-10-17 下午12:46 qingxu Exp $$
 */
 public class UniquekeyMetadataHandler extends DefaultMetadataHandler {
 /**
 * Logger for this class
 */
 private static final Logger logger = LoggerFactory.getLogger(UniquekeyMetadataHandler.class);
 
 private IDatabaseConnection databaseConnection;
 
 public UniquekeyMetadataHandler(IDatabaseConnection databaseConnection) {
 this.databaseConnection = databaseConnection;
 }
 
 @Override
 public ResultSet getPrimaryKeys(DatabaseMetaData metaData, String schemaName, String tableName)
 throws SQLException {
  ResultSet rs = super.getPrimaryKeys(metaData, schemaName, tableName);
         if (rs.next()) {
             //            rs.previous();由于不是可移动游标，无法前移，这里只有再查一次了
             rs.close();//must colsed
             rs = super.getPrimaryKeys(metaData, schemaName, tableName);
             return rs;
 
         }
         rs.close();//must colsed
 
 if (logger.isTraceEnabled())
 logger
 .trace(
 "UniquekeyMetadataHandler#getPrimaryKeys(metaData={}, schemaName={}, tableName={}) - start",
 new Object[] { metaData, schemaName, tableName });
 
 String productName = metaData.getDatabaseProductName();
 String sql;
 if (StringUtils.equalsIgnoreCase("oracle", productName)) {
 sql = "select col.owner,col.constraint_name,col.table_name,col.column_name,col.position\n"
 + " from user_constraints con, user_cons_columns col\n"
 + " where con.constraint_name = col.constraint_name\n"
 + " and con.constraint_type = 'U'\n" + " and col.table_name = ?";
 } else if (StringUtils.equalsIgnoreCase("mysql", productName)) {
 sql = "select distinct a.table_schema,a.constraint_name,a.table_name,a.column_name,a.ordinal_position\n"
 + " from information_schema.key_column_usage a,\n"
 + " information_schema.table_constraints b\n"
 + " where a.constraint_name = b.constraint_name\n"
 + " and b.constraint_type = 'UNIQUE'\n" + " and a.table_name = ?";
 
 } else {
 throw new RuntimeException("暂时不支持 数据库[" + productName + "]的唯一键获取~");
 }
 
 Connection connection = this.databaseConnection.getConnection();
 PreparedStatement preparedStatement = connection.prepareStatement(sql);
 preparedStatement.setString(1, tableName);
 rs = preparedStatement.executeQuery();
 
 return rs;
 
 }
 }

2、在构建Tester时更换配置处理器

databaseTester = new DataSourceDatabaseTester(dataSource);
    databaseTester.setOperationListener(getOperationListener());
     //替换meta处理器，如果数据库没有主键，则用唯一键代替主键
     protected IOperationListener getOperationListener() {
            logger.debug("getOperationListener() - start");
            if (this.operationListener == null) {
                this.operationListener = new DefaultOperationListener() {
                    public void connectionRetrieved(IDatabaseConnection connection) {
                        super.connectionRetrieved(connection);
                        // When a new connection has been created then invoke the setUp method
                        // so that user defined DatabaseConfig parameters can be set.
                        setUpDatabaseConfig(connection, connection.getConfig());
                    }
                };
            }
            return this.operationListener;
        }
    
        /**
         * Designed to be overridden by subclasses in order to set additional configuration
         * parameters for the {@link IDatabaseConnection}.
         *
         * @param connection
         * @param config     The settings of the current {@link IDatabaseConnection} to be configured
         */
        protected void setUpDatabaseConfig(IDatabaseConnection connection, DatabaseConfig config) {
            // Designed to be overridden.
            config.setProperty(DatabaseConfig.PROPERTY_METADATA_HANDLER, new UniquekeyMetadataHandler(
                connection));
            config.setProperty(DatabaseConfig.PROPERTY_DATATYPE_FACTORY, new OracleDataTypeFactory());
            config.setProperty(DatabaseConfig.FEATURE_BATCHED_STATEMENTS, Boolean.TRUE);
        }

OK，其他处理方式照旧，DBUNIT开始支持唯一键处理。