<commit_before>package mini.jdbc;

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Statement that uses named parameters and safe operations to fetch results by hiding result set manipulation.
 */
public class DbStatement<T> implements AutoCloseable {
    @NotNull
    private final PreparedStatement statement;

    private final Map<String, List<Integer>> indexMap = new HashMap<>();
    private final DbMapper<T> mapper;

    public DbStatement(DbConnection c, String query, DbMapper<T> mapper) throws SQLException {
        this.mapper = mapper;
        String parsedQuery = parse(query, indexMap);
        try {
            statement = c.sqlConnection.prepareStatement(parsedQuery);
        } finally {
            c.statementsToClose.add(this);
        }
    }


    public DbStatement<T> setString(String name, String value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setString(i, value);
        }
        return this;
    }


    public DbStatement<T> setInt(String name, int value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setInt(i, value);
        }
        return this;
    }

    public DbStatement<T> setLong(String name, long value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setLong(i, value);
        }
        return this;
    }


    public DbStatement<T> setBoolean(String name, boolean value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setBoolean(i, value);
        }
        return this;
    }


    @Nullable
    public T query() throws SQLException {
        try (ResultSet r = statement.executeQuery()) {
            if (r.next()) {
                return mapper.map(r);
            }
            return null;
        }
    }

    @NotNull
    public List<T> queryList() throws SQLException {
        List<T> res = new ArrayList<>();
        try (ResultSet r = statement.executeQuery()) {
            while (r.next()) {
                res.add(mapper.map(r));
            }
        }
        return res;
    }


    public void insert() throws SQLException {
        statement.executeUpdate();
    }

    @NotNull
    public T insertAndReturnId() throws SQLException {
        try (ResultSet r = statement.executeQuery()) {
            if (r.next()) {
                return mapper.map(r);
            }
            throw new SQLException("Insert returned no ID!");
        }
    }

    public int update() throws SQLException {
        return statement.executeUpdate();
    }


    @NotNull
    private List<Integer> getIndexes(String name) {
        List<Integer> indexes = indexMap.get(name);
        if (indexes == null) {
            throw new IllegalArgumentException("Parameter not found: " + name);
        }
        return indexes;
    }


    // TODO: add mappings cache for parsed maps
    private static String parse(String query, Map<String, List<Integer>> paramMap) {<commit_msg>DbStatement cleanup<commit_after>package mini.jdbc;

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.math.BigDecimal;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLType;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Statement that uses named parameters and safe operations to fetch results by hiding result set manipulation.
 */
public class DbStatement<T> implements AutoCloseable {
    @NotNull
    public final PreparedStatement statement;

    @NotNull
    public final DbMapper<T> mapper;

    private final Map<String, List<Integer>> indexMap = new HashMap<>();

    public DbStatement(@NotNull DbConnection c, @NotNull String query, @NotNull DbMapper<T> mapper) throws SQLException {
        this.mapper = mapper;
        String parsedQuery = parse(query, indexMap);
        try {
            statement = c.sqlConnection.prepareStatement(parsedQuery);
        } finally {
            c.statementsToClose.add(this);
        }
    }

    public DbStatement<T> setNull(@NotNull String name, @NotNull SQLType type) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setNull(i, type.getVendorTypeNumber());
        }
        return this;
    }

    @NotNull
    public DbStatement<T> setBoolean(@NotNull String name, boolean value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setBoolean(i, value);
        }
        return this;
    }

    @NotNull
    public DbStatement<T> setInt(@NotNull String name, int value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setInt(i, value);
        }
        return this;
    }

    @NotNull
    public DbStatement<T> setLong(@NotNull String name, long value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setLong(i, value);
        }
        return this;
    }

    @NotNull
    public DbStatement<T> setFloat(@NotNull String name, float value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setFloat(i, value);
        }
        return this;
    }

    @NotNull
    public DbStatement<T> setDouble(@NotNull String name, double value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setDouble(i, value);
        }
        return this;
    }

    @NotNull
    public DbStatement<T> setBigDecimal(@NotNull String name, BigDecimal value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setBigDecimal(i, value);
        }
        return this;
    }

    @NotNull
    public DbStatement<T> setString(@NotNull String name, String value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setString(i, value);
        }
        return this;
    }

    @NotNull
    public DbStatement<T> setBytes(@NotNull String name, byte[] value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setBytes(i, value);
        }
        return this;
    }

    @NotNull
    public DbStatement<T> setTimestamp(@NotNull String name, Timestamp value) throws SQLException {
        for (int i : getIndexes(name)) {
            statement.setTimestamp(i, value);
        }
        return this;
    }

    @Nullable
    public T query() throws SQLException {
        try (ResultSet r = statement.executeQuery()) {
            if (r.next()) {
                return mapper.map(r);
            }
            return null;
        }
    }

    @NotNull
    public List<T> queryList() throws SQLException {
        List<T> res = new ArrayList<>();
        try (ResultSet r = statement.executeQuery()) {
            while (r.next()) {
                res.add(mapper.map(r));
            }
        }
        return res;
    }


    public void insert() throws SQLException {
        statement.executeUpdate();
    }

    @NotNull
    public T insertAndReturnId() throws SQLException {
        statement.executeUpdate();
        try (ResultSet r = statement.getGeneratedKeys()) {
            if (r.next()) {
                return mapper.map(r);
            }
            throw new SQLException("Insert returned no ID!");
        }
    }

    public int update() throws SQLException {
        return statement.executeUpdate();
    }

    @NotNull
    public List<Integer> getIndexes(String name) {
        List<Integer> indexes = indexMap.get(name);
        if (indexes == null) {
            throw new IllegalArgumentException("Parameter not found: " + name);
        }
        return indexes;
    }


    // TODO: add mappings cache for parsed maps -> reuse DbConnection -> Db::cache!
    @NotNull
    private static String parse(@NotNull String query, @NotNull Map<String, List<Integer>> paramMap) {
{'fraction_non_alphanumeric': 0.0994488671157379, 'fraction_numerical': 0.0, 'mean_word_length': 2.2940701895925777, 'pattern_counts': {'":': 0, '<': 43, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2126446', 'n_tokens_mistral': 2254, 'n_tokens_neox': 2167, 'n_words': 620}