kern/src/de/superx/spring/batch/tasklet/PgCopyLoadTasklet.java

package de.superx.spring.batch.tasklet;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.Reader;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Statement;

import javax.sql.DataSource;

import org.postgresql.PGConnection;
import org.postgresql.copy.CopyManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.batch.core.ExitStatus;
import org.springframework.batch.core.StepContribution;
import org.springframework.batch.core.StepListener;
import org.springframework.batch.core.scope.context.ChunkContext;
import org.springframework.batch.core.step.tasklet.Tasklet;
import org.springframework.batch.repeat.RepeatStatus;
import org.springframework.core.io.Resource;
import org.springframework.jdbc.core.JdbcTemplate;

import de.superx.util.TrailingDelimiterRemovingReader;

public class PgCopyLoadTasklet implements Tasklet, StepListener {
    
    Logger logger = LoggerFactory.getLogger(this.getClass());
    
    private String table;
    private Resource unlFile;
    private String delimiter;
    private boolean isCsv;
    private String encoding = "UTF8";
    private boolean header = false;
    private boolean truncateTargetTable;
    private DataSource dataSource;
    private boolean continueOnError = true;
    
    
    public PgCopyLoadTasklet(Resource unlFile, String delimiter, String table, DataSource dataSource) {
    	this(unlFile, delimiter, table, dataSource, true, true, false);
    }
    
    public PgCopyLoadTasklet(Resource unlFile, String delimiter, String table, DataSource dataSource, boolean truncateTable, boolean isCsv) {
    	this(unlFile, delimiter, table, dataSource, truncateTable, true, isCsv);
    }
    
    public PgCopyLoadTasklet(Resource unlFile, String delimiter, String table, DataSource dataSource, boolean truncateTable, boolean continueOnError, boolean isCsv) {
        this.table = table;
        this.unlFile = unlFile;
        this.delimiter = delimiter;
        this.dataSource = dataSource;
        this.truncateTargetTable = truncateTable;
        this.continueOnError = continueOnError;
        this.isCsv = isCsv;
    }

    @Override
    public RepeatStatus execute(StepContribution contribution, ChunkContext chunkContext) throws Exception {
        logger.info("Load from " + unlFile.getFile().getAbsolutePath());
        if (!unlFile.exists()) {
            logger.warn("Not loading because unl file does not exist: " + unlFile.getFilename());
            return RepeatStatus.FINISHED;
        }
        logger.info("Into table " + table);
        String copySql   = "COPY " + table + " FROM STDIN WITH ("+ (isCsv ? "FORMAT csv, " : "") + "DELIMITER '" + delimiter 
                        + "', NULL '', ENCODING '" + encoding + "'" + (header ? " HEADER" : "") + ")";
        try (Connection con = dataSource.getConnection();
             FileReader unl = new FileReader(unlFile.getFile());
             BufferedReader reader = new BufferedReader(unl);
                        ) {
            int columnCount = getColumnCount(con);
            long count = 0;
            if(!isCsv) {
                try( Reader in = new TrailingDelimiterRemovingReader(reader, delimiter, columnCount, true, unlFile.getFilename()) ) {
                    CopyManager copyManager = getPgConnection(con).getCopyAPI();
                    count = copyManager.copyIn(copySql, in);
                }
            } else {
                CopyManager copyManager = getPgConnection(con).getCopyAPI();
                count = copyManager.copyIn(copySql, reader);
            }
            JdbcTemplate jt = new JdbcTemplate(dataSource);
            jt.execute("ANALYZE " + table);
            logger.info(count + " records loaded with analyze.");
        } catch (final Exception e) {
            String errorMsg = e.getMessage();
            logger.error("Error loading table: " + table);
            logger.error(errorMsg);
		    if (!continueOnError) {
		        throw e;
		    }
            ExitStatus exitStatus = ExitStatus.FAILED;
            exitStatus = exitStatus.addExitDescription(errorMsg);
            contribution.setExitStatus(exitStatus);
        }
        return RepeatStatus.FINISHED;
    }

    private static PGConnection getPgConnection(Connection connection) throws SQLException {
        PGConnection pgConnection = null;
        if (connection instanceof PGConnection) {
            pgConnection = (PGConnection) connection;
        } else {
            pgConnection = connection.unwrap(PGConnection.class);
        }
        return pgConnection;
    }
    
    private int getColumnCount(Connection con) throws SQLException {
        int count = -1;
        try(Statement st = con.createStatement();
                        ResultSet rs = st.executeQuery("select * from " + table + " limit 1")) {
            ResultSetMetaData rsm = rs.getMetaData();
            count = rsm.getColumnCount();
            if (truncateTargetTable) {
                st.execute("delete from " + table);
            }
        }
        return count;
    }
    
}