Add text-to-sql module: schema prompt through a restricted role, JSqlParser guard, read-only role with column grants and timeout, row cap, execution-based evaluation on a real PostgreSQL
Co-Authored-By: Claude Sonnet 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01JXVi2GMQ7bR5EmbUFdDj7N
This commit is contained in:
@@ -0,0 +1,56 @@
|
||||
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();
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user