primarykeyfilter.java

采用 Java 编写的数据库系统单元测试程序。

    // first, obtain the current allowed ids for that table
    Set currentAllowedIds = (Set) this.allowedPKsPerTable.get( table );
    if ( currentAllowedIds == null ) {
      currentAllowedIds = new HashSet();
      this.allowedPKsPerTable.put( table, currentAllowedIds );
    }
    // then, add the new ids, but checking if it should be allowed to add them
    Set forcedAllowedPKs = (Set) this.allowedPKsInput.get( table );
    if ( forcedAllowedPKs == null || forcedAllowedPKs.isEmpty() ) {
      currentAllowedIds.addAll( newAllowedPKs );
    } else {
      Iterator iterator = newAllowedPKs.iterator();
      while ( iterator.hasNext() ) {
        Object id = iterator.next();
        if ( forcedAllowedPKs.contains(id) ) {
          currentAllowedIds.add(id);
        } else {
          if ( this.logger.isDebugEnabled() ) {
            this.logger.debug( "Discarding id " + id + " of table " + table + 
              " as it was not included in the input!" );
          }
        }
      }
    }
  }
  
  private void scanPKs( String table, String pkColumn, Set allowedIds ) throws SQLException {
        logger.debug("scanPKs(table=" + table + ", pkColumn=" + pkColumn + ", allowedIds=" + allowedIds + ") - start");

    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() );
    Iterator iterator = fkEdges.iterator();
    StringBuffer colsBuffer = new StringBuffer();
    while ( 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 );
    }
    PreparedStatement pstmt = null;
    ResultSet rs = null;
    try {
      pstmt = this.connection.getConnection().prepareStatement( sql );
      iterator = allowedIds.iterator();
      while( 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 ) {
              if ( this.logger.isDebugEnabled() ) {
                this.logger.debug("New ID: " + newTable + "->" + fk  );
              }
              addPKToScan( newTable, fk );
            } else {
              this.logger.warn( "Found null FK for relationship  " + 
                  table + "=>" + newTable );
            }
          }
        }
      }
    } catch (SQLException e) {
      logger.error("scanPKs()", e);
      SQLHelper.close( rs, pstmt );
    }        
  }
  
  private void scanReversePKs(String table, Set pksToScan) throws SQLException {
        logger.debug("scanReversePKs(table=" + table + ", pksToScan=" + pksToScan + ") - start");

    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=" + edge + ", idsToScan=" + idsToScan + ") - start");

    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;
    try {
      if ( this.logger.isDebugEnabled() ) {
        this.logger.debug( "Preparing SQL query '" + sql + "'" );
      }
      pstmt = this.connection.getConnection().prepareStatement( sql );
    } catch (SQLException e) {
      logger.error("addReverseEdge()", e);
      SQLHelper.close( pstmt );
    }        
    ResultSet rs = null;
    Iterator iterator = idsToScan.iterator();
    try {
      while ( 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 );
    }
  }

  // TODO: support PKs with multiple values
  private String getPKColumn( String table ) throws SQLException {
        logger.debug("getPKColumn(table=" + table + ") - start");

    String pkColumn = (String) this.pkColumnPerTable.get( table );
    if ( pkColumn == null ) {
      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=" + table + ", ids=" + ids + ") - start");

    Set pksToScan = (Set) 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=" + table + ", pk=" + pk + ") - start");

    // first, check if it wasn't added yet
    Set scannedIds = (Set) this.allowedPKsPerTable.get( table );
    if ( scannedIds != null && scannedIds.contains(pk)) {
      if ( this.logger.isDebugEnabled() ) {
        this.logger.debug( "Discarding already scanned id=" + pk + " for table " + table );
      }
      return;
    }
    
    Set pksToScan = (Set) this.pksToScanPerTable.get(table);
    if ( pksToScan == null ) {
      pksToScan = new HashSet();
      this.pksToScanPerTable.put( table, pksToScan );
    }
    pksToScan.add( pk );
  }

  private class FilterIterator implements ITableIterator {

        /**
         * Logger for this class
         */
        private final Logger logger = LoggerFactory.getLogger(FilterIterator.class);

    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 = (Set) allowedPKsPerTable.get( tableName );
      if ( allowedPKs != null ) {
        return new PrimaryKeyFilteredTableWrapper(table, allowedPKs);
      }
      return table;
    }
  }

}