csv.java

非常棒的java数据库

/*
 * Copyright 2004-2008 H2 Group. Licensed under the H2 License, Version 1.0
 * (http://h2database.com/html/license.html).
 * Initial Developer: H2 Group
 */
package org.h2.tools;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.Reader;
import java.io.Writer;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Statement;
import java.sql.Types;
import java.util.ArrayList;

import org.h2.constant.SysProperties;
import org.h2.message.Message;
import org.h2.util.FileUtils;
import org.h2.util.IOUtils;
import org.h2.util.JdbcUtils;
import org.h2.util.StringCache;
import org.h2.util.StringUtils;

/**
 * A facility to read from and write to CSV (comma separated values) files.
 *
 * @author Thomas Mueller, Sylvain Cuaz
 */
public class Csv implements SimpleRowSource {

    private String charset = StringUtils.getDefaultCharset();
    private int bufferSize = 8 * 1024;
    private String[] columnNames;
    private char fieldSeparatorRead = ',';
    private char commentLineStart = '#';
    private String fieldSeparatorWrite = ",";
    private String rowSeparatorWrite;
    private char fieldDelimiter = '\"';
    private char escapeCharacter = '\"';
    private String lineSeparator = SysProperties.LINE_SEPARATOR;
    private String nullString = "";
    private String fileName;
    private Reader reader;
    private Writer writer;
    private int back;
    private boolean endOfLine, endOfFile;

    /**
     * Get a new object of this class.
     *
     * @return the new instance
     */
    public static Csv getInstance() {
        return new Csv();
    }

    private int writeResultSet(ResultSet rs) throws SQLException {
        try {
            ResultSetMetaData meta = rs.getMetaData();
            int rows = 0;
            int columnCount = meta.getColumnCount();
            String[] row = new String[columnCount];
            for (int i = 0; i < columnCount; i++) {
                row[i] = meta.getColumnLabel(i + 1);
            }
            writeRow(row);
            while (rs.next()) {
                for (int i = 0; i < columnCount; i++) {
                    row[i] = rs.getString(i + 1);
                }
                writeRow(row);
                rows++;
            }
            writer.close();
            return rows;
        } catch (IOException e) {
            throw Message.convertIOException(e, null);
        } finally {
            close();
            JdbcUtils.closeSilently(rs);
        }
    }

    /**
     * Writes the result set to a file in the CSV format.
     *
     * @param writer
     *            the writer
     * @param rs
     *            the result set
     * @return the number of rows written
     * @throws SQLException,
     *             IOException
     */
    public int write(Writer writer, ResultSet rs) throws SQLException, IOException {
        this.writer = writer;
        return writeResultSet(rs);
    }

    /**
     * Writes the result set to a file in the CSV format.
     *
     * @param fileName
     *            the name of the csv file
     * @param rs
     *            the result set
     * @param charset
     *            the charset or null to use UTF-8
     * @return the number of rows written
     * @throws SQLException
     */
    public int write(String fileName, ResultSet rs, String charset) throws SQLException {
        init(fileName, charset);
        try {
            initWrite();
            return writeResultSet(rs);
        } catch (IOException e) {
            throw convertException("IOException writing " + fileName, e);
        }
    }

    /**
     * Writes the result set of a query to a file in the CSV format.
     *
     * @param conn
     *            the connection
     * @param fileName
     *            the file name
     * @param sql
     *            the query
     * @param charset
     *            the charset or null to use UTF-8
     * @return the number of rows written
     * @throws SQLException
     */
    public int write(Connection conn, String fileName, String sql, String charset) throws SQLException {
        Statement stat = conn.createStatement();
        ResultSet rs = stat.executeQuery(sql);
        int rows = write(fileName, rs, charset);
        stat.close();
        return rows;
    }

    /**
     * Reads from the CSV file and returns a result set. The rows in the result
     * set are created on demand, that means the file is kept open until all
     * rows are read or the result set is closed.
     *
     * @param fileName the file name
     * @param colNames or null if the column names should be read from the CSV file
     * @param charset the charset or null to use UTF-8
     * @return the result set
     * @throws SQLException
     */
    public ResultSet read(String fileName, String[] colNames, String charset) throws SQLException {
        init(fileName, charset);
        try {
            return readResultSet(colNames);
        } catch (IOException e) {
            throw convertException("IOException reading " + fileName, e);
        }
    }

    /**
     * Reads CSV data from a reader and returns a result set. The rows in the
     * result set are created on demand, that means the reader is kept open
     * until all rows are read or the result set is closed.
     *
     * @param reader the reader
     * @param colNames or null if the column names should be read from the CSV file
     * @return the result set
     * @throws SQLException, IOException
     */
    public ResultSet read(Reader reader, String[] colNames) throws SQLException, IOException {
        init(null, null);
        this.reader = reader;
        return readResultSet(colNames);
    }

    private ResultSet readResultSet(String[] colNames) throws SQLException, IOException {
        this.columnNames = colNames;
        initRead();
        SimpleResultSet result = new SimpleResultSet(this);
        makeColumnNamesUnique();
        for (int i = 0; i < columnNames.length; i++) {
            result.addColumn(columnNames[i], Types.VARCHAR, Integer.MAX_VALUE, 0);
        }
        return result;
    }

    private void makeColumnNamesUnique() {
        for (int i = 0; i < columnNames.length; i++) {
            String x = columnNames[i];
            if (x == null || x.length() == 0) {
                x = "C" + (i + 1);
            }
            for (int j = 0; j < i; j++) {
                String y = columnNames[j];
                if (x.equals(y)) {
                    x = x + "1";
                    j = -1;
                }
            }
            columnNames[i] = x;
        }
    }

    private Csv() {
    }

    private void init(String fileName, String charset) {
        this.fileName = fileName;
        if (charset != null) {
            this.charset = charset;
        }
    }

    private void initWrite() throws IOException {
        if (writer == null) {
            try {
                OutputStream out = new FileOutputStream(fileName);
                out = new BufferedOutputStream(out, bufferSize);
                writer = new BufferedWriter(new OutputStreamWriter(out, charset));
            } catch (IOException e) {
                close();
                throw e;
            }
        }
    }

    private void writeRow(String[] values) throws IOException {
        for (int i = 0; i < values.length; i++) {
            if (i > 0) {
                if (fieldSeparatorWrite != null) {
                    writer.write(fieldSeparatorWrite);
                }
            }
            String s = values[i];
            if (s != null) {
                if (escapeCharacter != 0) {
                    if (fieldDelimiter != 0) {
                        writer.write(fieldDelimiter);

                    }
                    writer.write(escape(s));
                    if (fieldDelimiter != 0) {
                        writer.write(fieldDelimiter);
                    }
                } else {
                    writer.write(s);
                }
            } else if (nullString.length() > 0) {
                writer.write(nullString);
            }
        }
        if (rowSeparatorWrite != null) {
            writer.write(rowSeparatorWrite);
        }
        writer.write(lineSeparator);
    }

    private String escape(String data) {
        if (data.indexOf(fieldDelimiter) < 0) {
            if (escapeCharacter == fieldDelimiter || data.indexOf(escapeCharacter) < 0) {
                return data;
            }
        }
        StringBuffer buff = new StringBuffer(data.length());
        for (int i = 0; i < data.length(); i++) {
            char ch = data.charAt(i);
            if (ch == fieldDelimiter || ch == escapeCharacter) {
                buff.append(escapeCharacter);
            }
            buff.append(ch);
        }
        return buff.toString();
    }

    private void initRead() throws IOException {
        if (reader == null) {
            try {
                InputStream in = FileUtils.openFileInputStream(fileName);
                in = new BufferedInputStream(in, bufferSize);
                reader = new InputStreamReader(in, charset);
                reader = new BufferedReader(reader);
            } catch (IOException e) {
                close();
                throw e;
            }
        }
        if (columnNames == null) {
            readHeader();
        }
    }

    private void readHeader() throws IOException {
        ArrayList list = new ArrayList();
        while (true) {
            String v = readValue();
            if (v == null) {
                if (endOfLine) {
                    if (endOfFile || list.size() > 0) {
                        break;
                    }
                } else {
                    list.add("COLUMN" + list.size());
                }
            } else {
                list.add(v);
            }
        }
        columnNames = new String[list.size()];
        list.toArray(columnNames);
    }