com.turbospaces.spark.AbstractRawSQLIterator Maven / Gradle / Ivy

Go to download
package com.turbospaces.spark;

import java.io.Closeable;
import java.io.IOException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Iterator;
import java.util.Map;
import java.util.Objects;

import javax.sql.DataSource;

import org.apache.commons.lang3.exception.ExceptionUtils;
import org.apache.commons.lang3.mutable.MutableInt;
import org.apache.commons.lang3.time.StopWatch;
import org.flywaydb.core.internal.jdbc.JdbcUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.ebean.RowMapper;

public abstract class AbstractRawSQLIterator implements Iterator, RowMapper, Closeable {
    private final Logger logger = LoggerFactory.getLogger(getClass());
    private final MutableInt totalRows = new MutableInt();
    private final StopWatch stopWatch;
    private Connection connection;
    private PreparedStatement stmt;
    private ResultSet resultSet;
    private T toReturn;

    protected final Map options;

    public AbstractRawSQLIterator(DataSource ds, Map options, String query) {
        super();
        this.options = Objects.requireNonNull(options);
        this.stopWatch = StopWatch.createStarted();

        try {
            connection = ds.getConnection();
            stmt = connection.prepareStatement(query);
            bindParameters(stmt);
            resultSet = stmt.executeQuery();
        } catch (SQLException err) {
            ExceptionUtils.rethrow(err);
        }
    }
    @Override
    public boolean hasNext() {
        try {
            if (resultSet.next()) {
                int rowNumber = totalRows.incrementAndGet();
                toReturn = map(resultSet, rowNumber);
                return true;
            }
        } catch (SQLException err) {
            ExceptionUtils.rethrow(err);
        }

        stopWatch.stop();
        logger.info("totally fetched {} rows in {}", totalRows.getValue(), stopWatch);
        totalRows.setValue(0);
        return false;
    }
    @Override
    public T next() {
        return toReturn;
    }
    @Override
    public void close() throws IOException {
        JdbcUtils.closeConnection(connection);
        JdbcUtils.closeStatement(stmt);
        JdbcUtils.closeResultSet(resultSet);
    }

    protected abstract void bindParameters(PreparedStatement statement) throws SQLException;
}