eu.drus.jpa.unit.sql.dbunit.dataset.DataSetProducer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of jpa-unit-rdbms Show documentation
Makes JPA Unit support SQL databases
There is a newer version: 0.5.0
package eu.drus.jpa.unit.sql.dbunit.dataset;

import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.toList;
import static java.util.stream.Collectors.toSet;

import java.io.InputStream;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;

import org.dbunit.dataset.Column;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.DefaultTableMetaData;
import org.dbunit.dataset.ITableMetaData;
import org.dbunit.dataset.datatype.DataType;
import org.dbunit.dataset.stream.DefaultConsumer;
import org.dbunit.dataset.stream.IDataSetConsumer;
import org.dbunit.dataset.stream.IDataSetProducer;

public abstract class DataSetProducer implements IDataSetProducer {

    private IDataSetConsumer consumer = new DefaultConsumer();

    protected final InputStream input;

    public DataSetProducer(final InputStream input) {
        requireNonNull(input);
        this.input = input;
    }

    protected abstract Map>> loadDataSet() throws DataSetException;

    @Override
    public void setConsumer(final IDataSetConsumer consumer) throws DataSetException {
        this.consumer = consumer;
    }

    @Override
    public void produce() throws DataSetException {

        consumer.startDataSet();

        final Map>> dataset = loadDataSet();

        for (final Map.Entry>> entry : dataset.entrySet()) {
            // an entry represents a table
            final List> rows = entry.getValue();
            // each row represents a record in a table
            final Collection columnNames = extractColumnNames(rows);
            final ITableMetaData tableMetaData = new DefaultTableMetaData(entry.getKey(), createColumns(columnNames));

            consumer.startTable(tableMetaData);

            for (final Map row : rows) {
                final List values = new ArrayList<>();
                for (final Column column : tableMetaData.getColumns()) {
                    final Object rawValue = row.get(column.getColumnName());
                    final String value = rawValue == null ? null : String.valueOf(rawValue);
                    values.add(value);
                }
                consumer.row(values.toArray());
            }

            consumer.endTable();

        }

        consumer.endDataSet();
    }

    private Column[] createColumns(final Collection columnNames) {
        final List columns = columnNames.stream().map(e -> new Column(e, DataType.UNKNOWN)).collect(toList());
        return columns.toArray(new Column[columns.size()]);
    }

    private Collection extractColumnNames(final List> rows) {
        return rows.stream().flatMap(e -> e.keySet().stream()).collect(toSet());
    }

}