001package com.unitils.boot.xls;
002
003import org.apache.commons.lang.StringUtils;
004import org.dbunit.database.AmbiguousTableNameException;
005import org.dbunit.dataset.*;
006import org.dbunit.dataset.excel.XlsDataSet;
007import org.unitils.core.UnitilsException;
008import org.unitils.dbunit.util.MultiSchemaDataSet;
009
010import java.io.File;
011import java.io.FileInputStream;
012import java.util.*;
013
014/**
015 * @Author: yangjianzhou
016 * @Description:
017 * @Date:Created in 2018-07-08
018 */
019public class MultiSchemaXlsDataSetReader {
020    private String pattern = ".";
021    private String defaultSchemaName;
022
023    public MultiSchemaXlsDataSetReader(String defaultSchemaName) {
024        this.defaultSchemaName = defaultSchemaName;
025    }
026
027    public MultiSchemaDataSet readDataSetXls(File... dataSetFiles) {
028        try {
029            Map<String, List<ITable>> tbMap = getTables(dataSetFiles);
030            MultiSchemaDataSet dataSets = new MultiSchemaDataSet();
031
032            for (Map.Entry<String, List<ITable>> entry : tbMap.entrySet()) {
033                List<ITable> tables = entry.getValue();
034                try {
035                    DefaultDataSet ds = new DefaultDataSet(tables.toArray(new ITable[]{}));
036                    dataSets.setDataSetForSchema(entry.getKey(), ds);
037                } catch (AmbiguousTableNameException e) {
038                    throw new UnitilsException("构造DataSet失败！", e);
039                }
040            }
041            return dataSets;
042        } catch (Exception e) {
043            throw new UnitilsException("解析Excel文件出错：", e);
044        }
045    }
046
047    private Map<String, List<ITable>> getTables(File... dataSetFiles) {
048        Map<String, List<ITable>> tableMap = new HashMap<String, List<ITable>>();
049        // 需要根据schema把Table重新组合一下
050        try {
051            String schema, tableName;
052            for (File file : dataSetFiles) {
053                IDataSet dataSet = new XlsDataSet(new FileInputStream(file));
054                String[] tableNames = dataSet.getTableNames();
055                for (String tn : tableNames) {
056                    String[] temp = tn.split(pattern);
057                    if (temp.length == 2) {
058                        schema = temp[0];
059                        tableName = temp[1];
060                    } else {
061                        schema = this.defaultSchemaName;
062                        tableName = tn;
063                    }
064
065                    ITable table = dataSet.getTable(tn);
066                    if (!tableMap.containsKey(schema)) {
067                        tableMap.put(schema, new ArrayList<ITable>());
068                    }
069                    tableMap.get(schema).add(new XlsTable(tableName, table));
070                }
071            }
072        } catch (Exception e) {
073            throw new UnitilsException("Unable to create DbUnit dataset for data set files: " + Arrays.toString(dataSetFiles), e);
074        }
075        return tableMap;
076    }
077
078    //构造XslTable表
079    class XlsTable extends AbstractTable {
080        private ITable delegate;
081        private String tableName;
082
083        public XlsTable(String tableName, ITable table) {
084            this.delegate = table;
085            this.tableName = tableName;
086        }
087
088        public int getRowCount() {
089            return delegate.getRowCount();
090        }
091
092        public ITableMetaData getTableMetaData() {
093            ITableMetaData meta = delegate.getTableMetaData();
094            try {
095                return new DefaultTableMetaData(tableName, meta.getColumns(),
096                    meta.getPrimaryKeys());
097            } catch (DataSetException e) {
098                throw new UnitilsException("Don't get the meta info from  "
099                    + meta, e);
100            }
101        }
102
103        public Object getValue(int row, String column) throws DataSetException {
104            Object delta = delegate.getValue(row, column);
105            if (delta instanceof String) {
106                if (StringUtils.isEmpty((String) delta)) {
107                    return null;
108                }
109            }
110            return delta;
111        }
112
113    }
114}