CsvUtil.java

package com.base.framework.utils;

import com.google.common.base.Joiner;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang.StringEscapeUtils;
import org.apache.commons.lang.StringUtils;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.zip.ZipOutputStream;

public class CsvUtil {
    private static final int MaxHeaderLength = 2048;

    public static void exportCsv(String fileAbsPath, CsvHeaderSource headerSource, CsvDataSource dataSource, boolean dynamicHeader) throws Exception {
        try (FileOutputStream fileOutputStream = new FileOutputStream(fileAbsPath);
             OutputStreamWriter outputStreamWriter = new OutputStreamWriter(fileOutputStream, StandardCharsets.UTF_8);
             BufferedWriter bufferedWriter = new BufferedWriter(outputStreamWriter, 1024 * 256)
        ) {
            // 写入临时表头
            String emptyHeader = StringUtils.repeat(" ", MaxHeaderLength);
            bufferedWriter.write(emptyHeader);
            bufferedWriter.write("\n");

            int page = 1;
            List<Map<String, String>> rows = dataSource.getData(page);
            LinkedHashMap<String, String> header = headerSource.getHeader();
            while (rows.size() > 0) {
                for (Map<String, String> row : rows) {
                    // 处理动态表头
                    if (dynamicHeader) {
                        for (Map.Entry<String, String> entry : row.entrySet()) {
                            if (!header.containsKey(entry.getKey())) {
                                headerSource.addHeader(entry.getKey(), entry.getKey());
                            }
                        }
                    }
                    List<String> values = new ArrayList<>(header.size());
                    for (String key : header.keySet()) {
                        String value = row.get(key);
                        values.add(null == value ? "" : StringEscapeUtils.escapeCsv(value));
                    }
                    bufferedWriter.write(Joiner.on(",").join(values));
                    bufferedWriter.write("\n");
                }
                rows = dataSource.getData(++page);
            }
        }

        // 写入表头
        try (RandomAccessFile file = new RandomAccessFile(new File(fileAbsPath), "rw")) {
            file.seek(0);
            // 写入UTF-8 BOM
            file.writeByte(239);
            file.writeByte(187);
            file.writeByte(191);
            LinkedHashMap<String, String> header = headerSource.getHeader();
            StringBuilder sb = new StringBuilder();
            for (String key : header.keySet()) {
                sb.append(header.get(key));
                sb.append(",");
            }
            file.write(sb.toString().substring(0, Math.min(sb.length() - 1, MaxHeaderLength - 2)).getBytes(StandardCharsets.UTF_8));
            file.write(",".getBytes(StandardCharsets.UTF_8));
        }
    }

    public static void zipFiles(List<String> srcFiles, String zipFilePath, boolean deleteSrcFiles) throws Exception {
        File zipFile = new File(zipFilePath);
        try (ZipOutputStream zipOutputStream = new ZipOutputStream(new FileOutputStream(zipFile))){
            for (String filePath : srcFiles) {
                try (FileInputStream fileInputStream = new FileInputStream(filePath)) {
                    IOUtils.copy(fileInputStream, zipOutputStream);
                }
            }
        }
        if (deleteSrcFiles) {
            for (String file : srcFiles) {
                new File(file).delete();
            }
        }
    }

    public interface CsvHeaderSource {
        // 获取当前表头
        LinkedHashMap<String, String> getHeader();
        // 添加一个动态表头
        void addHeader(String key, String text);
    }

    public interface CsvDataSource {
        List<Map<String, String>> getData(int page);
    }
}