PrimaryKeyFilter.java
/*
*
* The DbUnit Database Testing Framework
* Copyright (C)2002-2005, DbUnit.org
*
* This library is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* This library is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with this library; if not, write to the Free Software
* Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
*
*/
package org.dbunit.database;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;
import org.dbunit.database.search.ForeignKeyRelationshipEdge;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.ITable;
import org.dbunit.dataset.ITableIterator;
import org.dbunit.dataset.ITableMetaData;
import org.dbunit.dataset.filter.AbstractTableFilter;
import org.dbunit.util.SQLHelper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Filter a table given a map of the allowed rows based on primary key values.<br>
* It uses a depth-first algorithm (although not recursive - it might be refactored
* in the future) to define which rows are allowed, as well which rows are necessary
* (and hence allowed) because of dependencies with the allowed rows.<br>
* <strong>NOTE:</strong> multi-column primary keys are not supported at the moment.
* TODO: test cases
* @author Felipe Leme (dbunit@felipeal.net)
* @author Last changed by: $Author$
* @version $Revision$ $Date$
* @since Sep 9, 2005
*/
public class PrimaryKeyFilter extends AbstractTableFilter {
private final IDatabaseConnection connection;
private final PkTableMap allowedPKsPerTable;
private final PkTableMap allowedPKsInput;
private final PkTableMap pksToScanPerTable;
private final boolean reverseScan;
protected final Logger logger = LoggerFactory.getLogger(getClass());
// cache the primary keys
private final Map pkColumnPerTable = new HashMap();
private final Map fkEdgesPerTable = new HashMap();
private final Map fkReverseEdgesPerTable = new HashMap();
// name of the tables, in reverse order of dependency
private final List tableNames = new ArrayList();
/**
* Default constructor, it takes as input a map with desired rows in a final
* dataset; the filter will ensure that the rows necessary by these initial rows
* are also allowed (and so on...).
* @param connection database connection
* @param allowedPKs map of allowed rows, based on the primary keys (key is the name
* of a table; value is a Set with allowed primary keys for that table)
* @param reverseDependency flag indicating if the rows that depend on a row should
* also be allowed by the filter
*/
public PrimaryKeyFilter(IDatabaseConnection connection, PkTableMap allowedPKs, boolean reverseDependency) {
this.connection = connection;
this.allowedPKsPerTable = new PkTableMap();
this.allowedPKsInput = allowedPKs;
this.reverseScan = reverseDependency;
// we need a deep copy here
this.pksToScanPerTable = new PkTableMap(allowedPKs);
}
public void nodeAdded(Object node) {
this.tableNames.add( node );
if ( this.logger.isDebugEnabled() ) {
this.logger.debug("nodeAdded: " + node );
}
}
public void edgeAdded(ForeignKeyRelationshipEdge edge) {
if ( this.logger.isDebugEnabled() ) {
this.logger.debug("edgeAdded: " + edge );
}
// first add it to the "direct edges"
String from = (String) edge.getFrom();
Set edges = (Set) this.fkEdgesPerTable.get(from);
if ( edges == null ) {
edges = new HashSet();
this.fkEdgesPerTable.put( from, edges );
}
if ( ! edges.contains(edge) ) {
edges.add(edge);
}
// then add it to the "reverse edges"
String to = (String) edge.getTo();
edges = (Set) this.fkReverseEdgesPerTable.get(to);
if ( edges == null ) {
edges = new HashSet();
this.fkReverseEdgesPerTable.put(to, edges);
}
if ( ! edges.contains(edge) ) {
edges.add(edge);
}
// finally, update the PKs cache
updatePkCache(to, edge);
}
/**
* @see AbstractTableFilter
*/
public boolean isValidName(String tableName) throws DataSetException {
// boolean isValid = this.allowedIds.containsKey(tableName);
// return isValid;
return true;
}
public ITableIterator iterator(IDataSet dataSet, boolean reversed)
throws DataSetException {
if ( this.logger.isDebugEnabled() ) {
this.logger.debug("Filter.iterator()" );
}
try {
searchPKs(dataSet);
} catch (SQLException e) {
throw new DataSetException( e );
}
return new FilterIterator(reversed ? dataSet.reverseIterator() : dataSet
.iterator());
}
private void searchPKs(IDataSet dataSet) throws DataSetException, SQLException {
logger.debug("searchPKs(dataSet={}) - start", dataSet);
int counter = 0;
while ( !this.pksToScanPerTable.isEmpty() ) {
counter ++;
if ( this.logger.isDebugEnabled() ) {
this.logger.debug( "RUN # " + counter );
}
for( int i=this.tableNames.size()-1; i>=0; i-- ) {
String tableName = (String) this.tableNames.get(i);
// TODO: support multi-column PKs
String pkColumn = dataSet.getTable(tableName).getTableMetaData().getPrimaryKeys()[0].getColumnName();
Set tmpSet = this.pksToScanPerTable.get( tableName );
if ( tmpSet != null && ! tmpSet.isEmpty() ) {
Set pksToScan = new HashSet( tmpSet );
if ( this.logger.isDebugEnabled() ) {
this.logger.debug( "before search: "+ tableName + "=>" + pksToScan );
}
scanPKs( tableName, pkColumn, pksToScan );
scanReversePKs( tableName, pksToScan );
allowPKs( tableName, pksToScan );
removePKsToScan( tableName, pksToScan );
} // if
} // for
removeScannedTables();
} // while
if ( this.logger.isDebugEnabled() ) {
this.logger.debug( "Finished searchIds()" );
}
}
private void removeScannedTables() {
logger.debug("removeScannedTables() - start");
this.pksToScanPerTable.retainOnly(this.tableNames);
}
private void allowPKs(String table, Set newAllowedPKs) {
logger.debug("allowPKs(table={}, newAllowedPKs={}) - start", table, newAllowedPKs);
// then, add the new IDs, but checking if it should be allowed to add them
Set forcedAllowedPKs = this.allowedPKsInput.get( table );
if( forcedAllowedPKs == null || forcedAllowedPKs.isEmpty() ) {
allowedPKsPerTable.addAll(table, newAllowedPKs );
} else {
for(Iterator iterator = newAllowedPKs.iterator(); iterator.hasNext(); ) {
Object id = iterator.next();
if( forcedAllowedPKs.contains(id) ) {
allowedPKsPerTable.add(table, id);
}
else
{
logger.debug("Discarding id {} of table {} as it was not included in the input!", id, table);
}
}
}
}
private void scanPKs( String table, String pkColumn, Set allowedIds ) throws SQLException {
if (logger.isDebugEnabled())
{
logger.debug("scanPKs(table={}, pkColumn={}, allowedIds={}) - start",
new Object[]{ table, pkColumn, allowedIds });
}
Set fkEdges = (Set) this.fkEdgesPerTable.get( table );
if ( fkEdges == null || fkEdges.isEmpty() ) {
return;
}
// we need a temporary list as there is no warranty about the set order...
List fkTables = new ArrayList( fkEdges.size() );
final StringBuilder colsBuffer = new StringBuilder();
for(Iterator iterator = fkEdges.iterator(); iterator.hasNext(); ) {
ForeignKeyRelationshipEdge edge = (ForeignKeyRelationshipEdge) iterator.next();
fkTables.add( edge.getTo() );
colsBuffer.append( edge.getFKColumn() );
if ( iterator.hasNext() ) {
colsBuffer.append( ", " );
}
}
// NOTE: make sure the query below is compatible standard SQL
String sql = "SELECT " + colsBuffer + " FROM " + table +
" WHERE " + pkColumn + " = ? ";
if ( this.logger.isDebugEnabled() ) {
this.logger.debug( "SQL: " + sql );
}
scanPKs(table, sql, allowedIds, fkTables);
}
private void scanPKs(String table, String sql, Set allowedIds, List fkTables) throws SQLException
{
PreparedStatement pstmt = null;
ResultSet rs = null;
try {
pstmt = this.connection.getConnection().prepareStatement( sql );
for(Iterator iterator = allowedIds.iterator(); iterator.hasNext(); ) {
Object pk = iterator.next(); // id being scanned
if( this.logger.isDebugEnabled() ) {
this.logger.debug("Executing sql for ? = " + pk );
}
pstmt.setObject( 1, pk );
rs = pstmt.executeQuery();
while( rs.next() ) {
for( int i=0; i<fkTables.size(); i++ ) {
String newTable = (String) fkTables.get(i);
Object fk = rs.getObject(i+1);
if( fk != null ) {
logger.debug("New ID: {}->{}", newTable, fk);
addPKToScan( newTable, fk );
}
else {
logger.warn( "Found null FK for relationship {} =>{}", table, newTable );
}
}
}
}
} catch (SQLException e) {
logger.error("scanPKs()", e);
}
finally {
// new in the finally block. has been in the catch only before
SQLHelper.close( rs, pstmt );
}
}
private void scanReversePKs(String table, Set pksToScan) throws SQLException {
logger.debug("scanReversePKs(table={}, pksToScan={}) - start", table, pksToScan);
if ( ! this.reverseScan ) {
return;
}
Set fkReverseEdges = (Set) this.fkReverseEdgesPerTable.get( table );
if ( fkReverseEdges == null || fkReverseEdges.isEmpty() ) {
return;
}
Iterator iterator = fkReverseEdges.iterator();
while ( iterator.hasNext() ) {
ForeignKeyRelationshipEdge edge = (ForeignKeyRelationshipEdge) iterator.next();
addReverseEdge( edge, pksToScan );
}
}
private void addReverseEdge(ForeignKeyRelationshipEdge edge, Set idsToScan) throws SQLException {
logger.debug("addReverseEdge(edge={}, idsToScan=) - start", edge, idsToScan);
String fkTable = (String) edge.getFrom();
String fkColumn = edge.getFKColumn();
String pkColumn = getPKColumn( fkTable );
// NOTE: make sure the query below is compatible standard SQL
String sql = "SELECT " + pkColumn + " FROM " + fkTable + " WHERE " + fkColumn + " = ? ";
PreparedStatement pstmt = null;
ResultSet rs = null;
try {
logger.debug("Preparing SQL query '{}'", sql);
pstmt = this.connection.getConnection().prepareStatement(sql);
for(Iterator iterator = idsToScan.iterator(); iterator.hasNext(); ) {
Object pk = iterator.next();
if ( this.logger.isDebugEnabled() ) {
this.logger.debug( "executing query '" + sql + "' for ? = " + pk );
}
pstmt.setObject( 1, pk );
rs = pstmt.executeQuery();
while( rs.next() ) {
Object fk = rs.getObject(1);
addPKToScan( fkTable, fk );
}
}
} finally {
SQLHelper.close( rs, pstmt );
}
}
private void updatePkCache(String table, ForeignKeyRelationshipEdge edge) {
logger.debug("updatePkCache(to={}, edge={}) - start", table, edge);
Object pkTo = this.pkColumnPerTable.get(table);
if ( pkTo == null ) {
String pkColumn = edge.getPKColumn();
this.pkColumnPerTable.put( table, pkColumn );
}
}
// TODO: support PKs with multiple values
private String getPKColumn( String table ) throws SQLException {
logger.debug("getPKColumn(table={}) - start", table);
// Try to get the cached column
String pkColumn = (String) this.pkColumnPerTable.get( table );
if ( pkColumn == null ) {
// If the column has not been cached until now retrieve it from the database connection
pkColumn = SQLHelper.getPrimaryKeyColumn( this.connection.getConnection(), table );
this.pkColumnPerTable.put( table, pkColumn );
}
return pkColumn;
}
private void removePKsToScan(String table, Set ids) {
logger.debug("removePKsToScan(table={}, ids={}) - start", table, ids);
Set pksToScan = this.pksToScanPerTable.get(table);
if ( pksToScan != null ) {
if ( pksToScan == ids ) {
throw new RuntimeException( "INTERNAL ERROR on removeIdsToScan() for table " + table );
} else {
pksToScan.removeAll( ids );
}
}
}
private void addPKToScan(String table, Object pk) {
logger.debug("addPKToScan(table={}, pk={}) - start", table, pk);
// first, check if it wasn't added yet
if(this.allowedPKsPerTable.contains(table, pk)) {
if ( this.logger.isDebugEnabled() ) {
this.logger.debug( "Discarding already scanned id=" + pk + " for table " + table );
}
return;
}
this.pksToScanPerTable.add(table, pk);
}
public String toString() {
final StringBuilder sb = new StringBuilder();
sb.append("tableNames=").append(tableNames);
sb.append(", allowedPKsInput=").append(allowedPKsInput);
sb.append(", allowedPKsPerTable=").append(allowedPKsPerTable);
sb.append(", fkEdgesPerTable=").append(fkEdgesPerTable);
sb.append(", fkReverseEdgesPerTable=").append(fkReverseEdgesPerTable);
sb.append(", pkColumnPerTable=").append(pkColumnPerTable);
sb.append(", pksToScanPerTable=").append(pksToScanPerTable);
sb.append(", reverseScan=").append(reverseScan);
sb.append(", connection=").append(connection);
return sb.toString();
}
private class FilterIterator implements ITableIterator {
private final ITableIterator _iterator;
public FilterIterator(ITableIterator iterator) {
_iterator = iterator;
}
////////////////////////////////////////////////////////////////////////////
// ITableIterator interface
public boolean next() throws DataSetException {
if ( logger.isDebugEnabled() ) {
logger.debug("Iterator.next()" );
}
while (_iterator.next()) {
if (accept(_iterator.getTableMetaData().getTableName())) {
return true;
}
}
return false;
}
public ITableMetaData getTableMetaData() throws DataSetException {
if ( logger.isDebugEnabled() ) {
logger.debug("Iterator.getTableMetaData()" );
}
return _iterator.getTableMetaData();
}
public ITable getTable() throws DataSetException {
if ( logger.isDebugEnabled() ) {
logger.debug("Iterator.getTable()" );
}
ITable table = _iterator.getTable();
String tableName = table.getTableMetaData().getTableName();
Set allowedPKs = allowedPKsPerTable.get( tableName );
if ( allowedPKs != null ) {
return new PrimaryKeyFilteredTableWrapper(table, allowedPKs);
}
return table;
}
}
/**
* Map that associates a table with a set of primary key objects.
*
* @author gommma (gommma AT users.sourceforge.net)
* @author Last changed by: $Author$
* @version $Revision$ $Date$
* @since 2.3.0
*/
public static class PkTableMap
{
private final LinkedHashMap pksPerTable;
private final Logger logger = LoggerFactory.getLogger(PkTableMap.class);
public PkTableMap()
{
this.pksPerTable = new LinkedHashMap();
}
/**
* Copy constructor
* @param allowedPKs
*/
public PkTableMap(PkTableMap allowedPKs) {
this.pksPerTable = new LinkedHashMap();
Iterator iterator = allowedPKs.pksPerTable.entrySet().iterator();
while ( iterator.hasNext() ) {
Map.Entry entry = (Map.Entry) iterator.next();
String table = (String)entry.getKey();
SortedSet pkObjectSet = (SortedSet) entry.getValue();
SortedSet newSet = new TreeSet( pkObjectSet );
this.pksPerTable.put( table, newSet );
}
}
public int size() {
return pksPerTable.size();
}
public boolean isEmpty() {
return pksPerTable.isEmpty();
}
public boolean contains(String table, Object pkObject) {
Set pksPerTable = this.get(table);
return (pksPerTable != null && pksPerTable.contains(pkObject));
}
public void remove(String tableName) {
this.pksPerTable.remove(tableName);
}
public void put(String table, SortedSet pkObjects) {
this.pksPerTable.put(table, pkObjects);
}
public void add(String tableName, Object pkObject) {
Set pksPerTable = getCreateIfNeeded(tableName);
pksPerTable.add(pkObject);
}
public void addAll(String tableName, Set pkObjectsToAdd) {
Set pksPerTable = this.getCreateIfNeeded(tableName);
pksPerTable.addAll(pkObjectsToAdd);
}
public SortedSet get(String tableName) {
return (SortedSet) this.pksPerTable.get(tableName);
}
private SortedSet getCreateIfNeeded(String tableName){
SortedSet pksPerTable = this.get(tableName);
// Lazily create the set if it did not exist yet
if( pksPerTable == null ) {
pksPerTable = new TreeSet();
this.pksPerTable.put(tableName, pksPerTable);
}
return pksPerTable;
}
public String[] getTableNames() {
return (String[]) this.pksPerTable.keySet().toArray(new String[0]);
}
public void retainOnly(List tableNames) {
List tablesToRemove = new ArrayList();
for(Iterator iterator = this.pksPerTable.entrySet().iterator(); iterator.hasNext(); ) {
Map.Entry entry = (Map.Entry) iterator.next();
String table = (String) entry.getKey();
SortedSet pksToScan = (SortedSet) entry.getValue();
boolean removeIt = pksToScan.isEmpty();
if ( ! tableNames.contains(table) ) {
if ( this.logger.isWarnEnabled() ) {
this.logger.warn("Discarding ids " + pksToScan + " of table " + table +
"as this table has not been passed as input" );
}
removeIt = true;
}
if ( removeIt ) {
tablesToRemove.add( table );
}
}
for(Iterator iterator = tablesToRemove.iterator(); iterator.hasNext(); ) {
this.remove( (String)iterator.next() );
}
}
public String toString() {
final StringBuilder sb = new StringBuilder();
sb.append("pKsPerTable=").append(pksPerTable);
return sb.toString();
}
}
}