spring-ai/text-to-sql/src/main/java/com/ankurm/texttosql/QueryRunner.java

package com.ankurm.texttosql;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;

/** Runs one query and returns at most {@code maxRows} rows, saying whether it stopped early. */
public final class QueryRunner {

    public record Rows(List<String> columns, List<List<String>> rows, boolean truncated) {
    }

    private QueryRunner() {
    }

    public static Rows run(Connection c, String sql, int maxRows) throws SQLException {
        c.setAutoCommit(false);   // lets the JDBC driver ask the server for only maxRows + 1 rows
        try (Statement st = c.createStatement()) {
            st.setMaxRows(maxRows + 1);
            boolean hasResult = st.execute(sql);
            while (!hasResult && st.getUpdateCount() != -1) {
                hasResult = st.getMoreResults();
            }
            if (!hasResult) {
                return new Rows(List.of(), List.of(), false);
            }
            try (ResultSet rs = st.getResultSet()) {
                ResultSetMetaData md = rs.getMetaData();
                List<String> cols = new ArrayList<>();
                for (int i = 1; i <= md.getColumnCount(); i++) {
                    cols.add(md.getColumnLabel(i));
                }
                List<List<String>> out = new ArrayList<>();
                boolean truncated = false;
                while (rs.next()) {
                    if (out.size() == maxRows) {
                        truncated = true;
                        break;
                    }
                    List<String> row = new ArrayList<>();
                    for (int i = 1; i <= cols.size(); i++) {
                        row.add(rs.getString(i));
                    }
                    out.add(row);
                }
                return new Rows(cols, out, truncated);
            }
        } finally {
            c.rollback();
        }
    }
}