com.metricstream.jdbc.MockSQLBuilderProvider Maven / Gradle / Ivy

Go to download
/*
 * Copyright © 2020, MetricStream, Inc. All rights reserved.
 */
package com.metricstream.jdbc;

import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiFunction;
import java.util.function.Supplier;

import java.math.BigDecimal;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;


public final class MockSQLBuilderProvider implements SQLBuilderProvider {

    private static final Queue mockResultSets = new ConcurrentLinkedQueue<>();

    private static BiFunction intByColumnIndex;
    private static BiFunction intByColumnLabel;
    private static BiFunction longByColumnIndex;
    private static BiFunction longByColumnLabel;
    private static BiFunction stringByColumnIndex;
    private static BiFunction stringByColumnLabel;
    private static BiFunction bigDecimalByColumnIndex;
    private static BiFunction bigDecimalByColumnLabel;
    private static BiFunction objectByColumnIndex;
    private static BiFunction objectByColumnLabel;
    private static Supplier executeSupplier;
    private final boolean generateSingleRowResultSet;

    public MockSQLBuilderProvider() {
        this(true);
    }

    public MockSQLBuilderProvider(boolean generateSingleRowResultSet) {
        this.generateSingleRowResultSet = generateSingleRowResultSet;
        reset();
    }

    public static void addResultSet(ResultSet rs) {
        mockResultSets.add(rs);
    }

    public static void addResultSet(final String tag, final Object[][] data) throws SQLException {
        mockResultSets.add(MockResultSet.create(tag, data));
    }

    public static void addResultSet(final String tag, final String labels, final String... csvs) throws SQLException {
        mockResultSets.add(MockResultSet.create(tag, labels, csvs));
    }

    public static void addResultSet(final String tag, final String csv, boolean withLabels) throws SQLException {
        mockResultSets.add(MockResultSet.create(tag, csv, withLabels));
    }

    public static void addResultSet(final String tag, final String csv) throws SQLException {
        mockResultSets.add(MockResultSet.create(tag, csv, false));
    }

    public static void addResultSet(final String tag, final InputStream csv, boolean withLabels) throws SQLException {
        mockResultSets.add(MockResultSet.create(tag, csv, withLabels));
    }

    public static void addResultSet(final String tag, final InputStream csv) throws SQLException {
        mockResultSets.add(MockResultSet.create(tag, csv, true));
    }

    public static void setIntByColumnIndex(BiFunction intByColumnIndex) {
        MockSQLBuilderProvider.intByColumnIndex = intByColumnIndex;
    }

    public static void setIntByColumnLabel(BiFunction intByColumnLabel) {
        MockSQLBuilderProvider.intByColumnLabel = intByColumnLabel;
    }

    public static void setLongByColumnIndex(BiFunction longByColumnIndex) {
        MockSQLBuilderProvider.longByColumnIndex = longByColumnIndex;
    }

    public static void setLongByColumnLabel(BiFunction longByColumnLabel) {
        MockSQLBuilderProvider.longByColumnLabel = longByColumnLabel;
    }

    public static void setStringByColumnIndex(BiFunction stringByColumnIndex) {
        MockSQLBuilderProvider.stringByColumnIndex = stringByColumnIndex;
    }

    public static void setStringByColumnLabel(BiFunction stringByColumnLabel) {
        MockSQLBuilderProvider.stringByColumnLabel = stringByColumnLabel;
    }

    public static void setBigDecimalByColumnIndex(BiFunction bigDecimalByColumnIndex) {
        MockSQLBuilderProvider.bigDecimalByColumnIndex = bigDecimalByColumnIndex;
    }

    public static void setBigDecimalByColumnLabel(BiFunction bigDecimalByColumnLabel) {
        MockSQLBuilderProvider.bigDecimalByColumnLabel = bigDecimalByColumnLabel;
    }

    public static void setObjectByColumnIndex(BiFunction objectByColumnIndex) {
        MockSQLBuilderProvider.objectByColumnIndex = objectByColumnIndex;
    }

    public static void setObjectByColumnLabel(BiFunction objectByColumnLabel) {
        MockSQLBuilderProvider.objectByColumnLabel = objectByColumnLabel;
    }

    public static void setExecute(Supplier supplier) {
        executeSupplier = supplier;
    }

    public static void setExecute(int value) {
        executeSupplier = () -> value;
    }

    public static void setExecute(int... values) {
        final AtomicInteger count = new AtomicInteger();
        executeSupplier = () -> count.get() < values.length ? values[count.getAndIncrement()] : 42;
    }

    public ResultSet getResultSet(SQLBuilder sqlBuilder, Connection connection, boolean wrapConnection) throws SQLException {
        return getRs();
    }

    @Override
    public int getInt(SQLBuilder sqlBuilder, Connection connection, int columnNumber, int defaultValue) throws SQLException {
        if (intByColumnIndex != null) {
            return intByColumnIndex.apply(columnNumber, defaultValue);
        }
        final ResultSet rs = getRs();
        return rs.next() ? rs.getInt(columnNumber) : defaultValue;
    }

    /**
     * Returns a value from the first row returned when executing the query.
     * @param connection The Connection from which the PreparedStatement is created
     * @param columnName The name of the column from which to return the value
     * @param defaultValue The default value that is returned if the query did not return any rows
     * @return the value from the query
     * @the exception thrown when generating or accessing the ResultSet
     */
    public int getInt(SQLBuilder sqlBuilder, Connection connection, String columnName, int defaultValue) throws SQLException {
        if (intByColumnLabel != null) {
            return intByColumnLabel.apply(columnName, defaultValue);
        }
        final ResultSet rs = getRs();
        return rs.next() ? rs.getInt(columnName) : defaultValue;
    }

    @Override
    public long getLong(SQLBuilder sqlBuilder, Connection connection, int columnNumber, long defaultValue) throws SQLException {
        if (longByColumnIndex != null) {
            return longByColumnIndex.apply(columnNumber, defaultValue);
        }
        final ResultSet rs = getRs();
        return rs.next() ? rs.getLong(columnNumber) : defaultValue;
    }

    @Override
    public long getLong(SQLBuilder sqlBuilder, Connection connection, String columnName, long defaultValue) throws SQLException {
        if (longByColumnLabel != null) {
            return longByColumnLabel.apply(columnName, defaultValue);
        }
        final ResultSet rs = getRs();
        return rs.next() ? rs.getLong(columnName) : defaultValue;
    }

    @Override
    public String getString(SQLBuilder sqlBuilder, Connection connection, int columnNumber, String defaultValue) throws SQLException {
        if (stringByColumnIndex != null) {
            return stringByColumnIndex.apply(columnNumber, defaultValue);
        }
        final ResultSet rs = getRs();
        return rs.next() ? rs.getString(columnNumber) : defaultValue;
    }

    @Override
    public String getString(SQLBuilder sqlBuilder, Connection connection, String columnName, String defaultValue) throws SQLException {
        if (stringByColumnLabel != null) {
            return stringByColumnLabel.apply(columnName, defaultValue);
        }
        final ResultSet rs = getRs();
        return rs.next() ? rs.getString(columnName) : defaultValue;
    }

    @Override
    public BigDecimal getBigDecimal(SQLBuilder sqlBuilder, Connection connection, int columnNumber, BigDecimal defaultValue) throws SQLException {
        if (bigDecimalByColumnIndex != null) {
            return bigDecimalByColumnIndex.apply(columnNumber, defaultValue);
        }
        final ResultSet rs = getRs();
        return rs.next() ? rs.getBigDecimal(columnNumber) : defaultValue;
    }

    @Override
    public BigDecimal getBigDecimal(SQLBuilder sqlBuilder, Connection connection, String columnName, BigDecimal defaultValue) throws SQLException {
        if (bigDecimalByColumnLabel != null) {
            return bigDecimalByColumnLabel.apply(columnName, defaultValue);
        }
        final ResultSet rs = getRs();
        return rs.next() ? rs.getBigDecimal(columnName) : defaultValue;
    }

    @Override
    public Object getObject(SQLBuilder sqlBuilder, Connection connection, int columnNumber, Object defaultValue) throws SQLException {
        if (objectByColumnIndex != null) {
            return objectByColumnIndex.apply(columnNumber, defaultValue);
        }
        final ResultSet rs = getRs();
        return rs.next() ? rs.getObject(columnNumber) : defaultValue;
    }

    @Override
    public Object getObject(SQLBuilder sqlBuilder, Connection connection, String columnName, Object defaultValue) throws SQLException {
        if (objectByColumnLabel != null) {
            return objectByColumnLabel.apply(columnName, defaultValue);
        }
        final ResultSet rs = getRs();
        return rs.next() ? rs.getObject(columnName) : defaultValue;
    }

    @Override
    public int execute(SQLBuilder sqlBuilder, Connection connection) {
        return executeSupplier.get();
    }

    @Override
    public  List getList(SQLBuilder sqlBuilder, Connection connection, SQLBuilder.RowMapper rowMapper, boolean withNull) throws SQLException {
        final ResultSet rs = getRs();
        final List list = new ArrayList<>();
        while (rs.next()) {
            final T item = rowMapper.map(rs);
            if (withNull || item != null) {
                list.add(item);
            }
        }
        return list;
    }

    @Override
    public  Optional getSingle(SQLBuilder sqlBuilder, Connection connection, SQLBuilder.RowMapper rowMapper) throws SQLException {
        final ResultSet rs = getRs();
        return Optional.ofNullable(rs.next() ? rowMapper.map(rs) : null);
    }

    @Override
    public  T getSingle(SQLBuilder sqlBuilder, Connection connection, SQLBuilder.RowMapper rowMapper, T defaultValue) throws SQLException {
        final ResultSet rs = getRs();
        return rs.next() ? rowMapper.map(rs) : defaultValue;
    }

    private ResultSet getRs() throws SQLException {
        ResultSet rs = mockResultSets.poll();
        if (rs != null) {
            return rs;
        }
        return generateSingleRowResultSet ? MockResultSet.create("", "42", false) : MockResultSet.empty("");
    }

    public static void reset() {
        mockResultSets.clear();
        setExecute(42);
    }

}