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:
Claude
@@ -0,0 +1,24 @@
|
||||
package com.ankurm.texttosql;
|
||||
|
||||
import java.sql.Connection;
|
||||
import java.sql.DriverManager;
|
||||
import java.sql.SQLException;
|
||||
|
||||
/** Connection settings for the throwaway PostgreSQL started by scripts/pg-up.sh. */
|
||||
public final class Db {
|
||||
|
||||
public static final String URL = "jdbc:postgresql://127.0.0.1:" + System.getProperty("t2s.port", "5451") + "/t2s";
|
||||
|
||||
private Db() {
|
||||
}
|
||||
|
||||
/** The role the assistant uses: SELECT on four tables, read-only by default, 2 second statement timeout. */
|
||||
public static Connection reader() throws SQLException {
|
||||
return DriverManager.getConnection(URL, "t2s_reader", "reader");
|
||||
}
|
||||
|
||||
/** A superuser. Used only to set things up and to show what happens WITHOUT the read-only role. */
|
||||
public static Connection admin() throws SQLException {
|
||||
return DriverManager.getConnection(URL, "t2s_admin", "admin");
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package com.ankurm.texttosql;
|
||||
|
||||
import java.sql.Connection;
|
||||
import java.sql.PreparedStatement;
|
||||
import java.sql.ResultSet;
|
||||
import java.sql.SQLException;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* Builds the schema section of the prompt from the database itself (information_schema plus the
|
||||
* COMMENTs on tables and columns), for an explicit list of tables. A table that is not in the list
|
||||
* is not described, so the model is never shown that it exists.
|
||||
*/
|
||||
public final class SchemaPrompt {
|
||||
|
||||
private SchemaPrompt() {
|
||||
}
|
||||
|
||||
public static String describe(Connection c, List<String> tables) throws SQLException {
|
||||
StringBuilder sb = new StringBuilder();
|
||||
for (String table : tables) {
|
||||
String tableComment = scalar(c, "SELECT obj_description(?::regclass, 'pg_class')", table);
|
||||
sb.append("TABLE ").append(table);
|
||||
if (tableComment != null) {
|
||||
sb.append(" -- ").append(tableComment);
|
||||
}
|
||||
sb.append('\n');
|
||||
try (PreparedStatement ps = c.prepareStatement(
|
||||
"SELECT column_name, data_type, col_description(?::regclass, ordinal_position) "
|
||||
+ "FROM information_schema.columns WHERE table_schema = 'public' AND table_name = ? "
|
||||
+ "ORDER BY ordinal_position")) {
|
||||
ps.setString(1, table);
|
||||
ps.setString(2, table);
|
||||
try (ResultSet rs = ps.executeQuery()) {
|
||||
while (rs.next()) {
|
||||
sb.append(" ").append(rs.getString(1)).append(' ').append(rs.getString(2));
|
||||
if (rs.getString(3) != null) {
|
||||
sb.append(" -- ").append(rs.getString(3));
|
||||
}
|
||||
sb.append('\n');
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return sb.toString();
|
||||
}
|
||||
|
||||
private static String scalar(Connection c, String sql, String arg) throws SQLException {
|
||||
try (PreparedStatement ps = c.prepareStatement(sql)) {
|
||||
ps.setString(1, arg);
|
||||
try (ResultSet rs = ps.executeQuery()) {
|
||||
return rs.next() ? rs.getString(1) : null;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public static List<String> columnNames(Connection c, String table) throws SQLException {
|
||||
List<String> out = new ArrayList<>();
|
||||
try (PreparedStatement ps = c.prepareStatement(
|
||||
"SELECT column_name FROM information_schema.columns WHERE table_schema='public' AND table_name=? ORDER BY ordinal_position")) {
|
||||
ps.setString(1, table);
|
||||
try (ResultSet rs = ps.executeQuery()) {
|
||||
while (rs.next()) {
|
||||
out.add(rs.getString(1));
|
||||
}
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package com.ankurm.texttosql;
|
||||
|
||||
import java.util.HashSet;
|
||||
import java.util.Locale;
|
||||
import java.util.Set;
|
||||
|
||||
import net.sf.jsqlparser.JSQLParserException;
|
||||
import net.sf.jsqlparser.expression.Function;
|
||||
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
|
||||
import net.sf.jsqlparser.statement.Statement;
|
||||
import net.sf.jsqlparser.statement.Statements;
|
||||
import net.sf.jsqlparser.statement.select.ParenthesedSelect;
|
||||
import net.sf.jsqlparser.statement.select.PlainSelect;
|
||||
import net.sf.jsqlparser.statement.select.Select;
|
||||
import net.sf.jsqlparser.statement.select.TableFunction;
|
||||
import net.sf.jsqlparser.statement.select.WithItem;
|
||||
import net.sf.jsqlparser.util.TablesNamesFinder;
|
||||
|
||||
/**
|
||||
* Decides, BEFORE anything reaches the database, whether a model-written string is acceptable:
|
||||
* exactly one statement, a plain SELECT, no SELECT INTO or row locks, no data-modifying WITH,
|
||||
* only listed tables, only listed functions. Anything it cannot parse is rejected (fail closed).
|
||||
*
|
||||
* <p>This is a seat belt, not the wall. The wall is the database role. The two are tested together.
|
||||
*/
|
||||
public final class SqlGuard {
|
||||
|
||||
public record Verdict(boolean ok, String reason) {
|
||||
|
||||
static Verdict yes() {
|
||||
return new Verdict(true, "ok");
|
||||
}
|
||||
|
||||
static Verdict no(String why) {
|
||||
return new Verdict(false, why);
|
||||
}
|
||||
}
|
||||
|
||||
public static final Set<String> DEFAULT_FUNCTIONS = Set.of("count", "sum", "avg", "min", "max", "round", "coalesce",
|
||||
"nullif", "lower", "upper", "length", "date_trunc", "to_char", "abs", "row_number", "rank", "dense_rank");
|
||||
|
||||
private final Set<String> tables;
|
||||
|
||||
private final Set<String> functions;
|
||||
|
||||
public SqlGuard(Set<String> allowedTables, Set<String> allowedFunctions) {
|
||||
this.tables = lower(allowedTables);
|
||||
this.functions = lower(allowedFunctions);
|
||||
}
|
||||
|
||||
public Verdict check(String sql) {
|
||||
Statements parsed;
|
||||
try {
|
||||
parsed = CCJSqlParserUtil.parseStatements(sql);
|
||||
} catch (JSQLParserException e) {
|
||||
return Verdict.no("not parseable as a single SQL statement");
|
||||
}
|
||||
if (parsed.size() != 1) {
|
||||
return Verdict.no("more than one statement (" + parsed.size() + ")");
|
||||
}
|
||||
Statement st = parsed.get(0);
|
||||
if (!(st instanceof Select select)) {
|
||||
return Verdict.no("only SELECT is allowed, found " + st.getClass().getSimpleName());
|
||||
}
|
||||
Set<String> cteNames = new HashSet<>();
|
||||
if (select.getWithItemsList() != null) {
|
||||
for (WithItem<?> w : select.getWithItemsList()) {
|
||||
if (!(w.getParenthesedStatement() instanceof ParenthesedSelect)) {
|
||||
return Verdict.no("WITH item is not a SELECT (data-modifying CTE)");
|
||||
}
|
||||
cteNames.add(w.getAlias().getName().toLowerCase(Locale.ROOT));
|
||||
}
|
||||
}
|
||||
if (select instanceof PlainSelect ps) {
|
||||
if (ps.getIntoTables() != null && !ps.getIntoTables().isEmpty()) {
|
||||
return Verdict.no("SELECT INTO creates a table");
|
||||
}
|
||||
if (ps.getForMode() != null) {
|
||||
return Verdict.no("row locking clause (" + ps.getForMode() + ")");
|
||||
}
|
||||
}
|
||||
Set<String> usedFunctions = new HashSet<>();
|
||||
TablesNamesFinder<Void> finder = new TablesNamesFinder<>() {
|
||||
@Override
|
||||
public <S> Void visit(Function function, S context) {
|
||||
usedFunctions.add(function.getName().toLowerCase(Locale.ROOT));
|
||||
return super.visit(function, context);
|
||||
}
|
||||
|
||||
@Override
|
||||
public <S> Void visit(TableFunction tableFunction, S context) {
|
||||
usedFunctions.add(tableFunction.getFunction().getName().toLowerCase(Locale.ROOT));
|
||||
return super.visit(tableFunction, context);
|
||||
}
|
||||
};
|
||||
Set<String> used;
|
||||
try {
|
||||
used = finder.getTables(st);
|
||||
} catch (RuntimeException e) {
|
||||
return Verdict.no("could not analyse the statement: " + e.getClass().getSimpleName());
|
||||
}
|
||||
for (String t : used) {
|
||||
String name = t.toLowerCase(Locale.ROOT).replace("\"", "");
|
||||
if (name.startsWith("public.")) {
|
||||
name = name.substring("public.".length());
|
||||
}
|
||||
if (!tables.contains(name) && !cteNames.contains(name)) {
|
||||
return Verdict.no("table " + t + " is not allowed");
|
||||
}
|
||||
}
|
||||
for (String f : usedFunctions) {
|
||||
if (!functions.contains(f)) {
|
||||
return Verdict.no("function " + f + " is not allowed");
|
||||
}
|
||||
}
|
||||
return Verdict.yes();
|
||||
}
|
||||
|
||||
private static Set<String> lower(Set<String> in) {
|
||||
Set<String> out = new HashSet<>();
|
||||
in.forEach(s -> out.add(s.toLowerCase(Locale.ROOT)));
|
||||
return out;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package com.ankurm.texttosql;
|
||||
|
||||
import java.sql.Connection;
|
||||
import java.sql.SQLException;
|
||||
import java.util.List;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
|
||||
/** question -> SQL (model) -> guard -> database (read-only role) -> at most N rows. */
|
||||
public class TextToSql {
|
||||
|
||||
public static final List<String> TABLES = List.of("customers", "products", "orders", "order_items");
|
||||
|
||||
public record Answer(String sql, String rejectedBy, String databaseError, QueryRunner.Rows rows) {
|
||||
|
||||
public boolean answered() {
|
||||
return rows != null;
|
||||
}
|
||||
}
|
||||
|
||||
private final ChatClient client;
|
||||
|
||||
private final SqlGuard guard;
|
||||
|
||||
private final String system;
|
||||
|
||||
private final int maxRows;
|
||||
|
||||
public TextToSql(ChatClient client, SqlGuard guard, String schema, int maxRows) {
|
||||
this.client = client;
|
||||
this.guard = guard;
|
||||
this.maxRows = maxRows;
|
||||
this.system = """
|
||||
You translate a question into ONE PostgreSQL SELECT statement.
|
||||
Use only the tables and columns below. Return only the SQL, no explanation, no markdown.
|
||||
Never write INSERT, UPDATE, DELETE, DDL or more than one statement.
|
||||
|
||||
""" + schema;
|
||||
}
|
||||
|
||||
public String systemPrompt() {
|
||||
return system;
|
||||
}
|
||||
|
||||
public Answer answer(String question, Connection db) {
|
||||
String raw = client.prompt().system(system).user(question).call().content();
|
||||
String sql = extractSql(raw);
|
||||
SqlGuard.Verdict verdict = guard.check(sql);
|
||||
if (!verdict.ok()) {
|
||||
return new Answer(sql, verdict.reason(), null, null);
|
||||
}
|
||||
try {
|
||||
return new Answer(sql, null, null, QueryRunner.run(db, sql, maxRows));
|
||||
} catch (SQLException e) {
|
||||
return new Answer(sql, null, firstLine(e.getMessage()), null);
|
||||
}
|
||||
}
|
||||
|
||||
/** Models like to wrap SQL in a markdown fence even when told not to. */
|
||||
public static String extractSql(String raw) {
|
||||
String s = raw == null ? "" : raw.strip();
|
||||
if (s.startsWith("```")) {
|
||||
s = s.replaceFirst("^```[a-zA-Z]*\\s*", "").replaceFirst("\\s*```\\s*$", "");
|
||||
}
|
||||
return s.strip();
|
||||
}
|
||||
|
||||
static String firstLine(String message) {
|
||||
return message == null ? "" : message.lines().findFirst().orElse("").strip();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package com.ankurm.texttosql;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.sql.Connection;
|
||||
import java.sql.SQLException;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
import com.ankurm.texttosql.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
/**
|
||||
* Hostile or merely careless SQL, as a model might write it, against four setups on a REAL PostgreSQL:
|
||||
* neither protection (a superuser connection), the parser guard only, the read-only role only, both.
|
||||
* Everything runs inside a transaction that is rolled back.
|
||||
*/
|
||||
class AttackMatrixTest {
|
||||
|
||||
record Attack(String name, String sql, boolean harmful) {
|
||||
}
|
||||
|
||||
static final List<Attack> ATTACKS = List.of(
|
||||
new Attack("DELETE FROM order_items", "DELETE FROM order_items", true),
|
||||
new Attack("second statement drops a table", "SELECT * FROM orders; DROP TABLE order_items", true),
|
||||
new Attack("table the assistant must not see", "SELECT * FROM api_keys", true),
|
||||
new Attack("email column of an allowed table", "SELECT email FROM customers", true),
|
||||
new Attack("system catalog: list database roles", "SELECT usename FROM pg_user", true),
|
||||
new Attack("pg_sleep(3)", "SELECT pg_sleep(3)", true),
|
||||
new Attack("read a server file", "SELECT pg_read_file('/etc/passwd')", true),
|
||||
new Attack("data-modifying CTE", "WITH d AS (DELETE FROM order_items RETURNING *) SELECT count(*) FROM d", true),
|
||||
new Attack("row locks on every order", "SELECT * FROM orders FOR UPDATE", true),
|
||||
new Attack("SELECT INTO makes a table", "SELECT * INTO newtab FROM orders", true),
|
||||
new Attack("2 billion generated rows", "SELECT count(*) FROM generate_series(1, 2000000000)", true),
|
||||
new Attack("turn read-only off for the session", "SELECT set_config('default_transaction_read_only', 'off', false)", true),
|
||||
new Attack("cartesian product, allowed names only", "SELECT count(*) FROM orders a, orders b, orders c, orders d", true),
|
||||
new Attack("(legit) join and group", "SELECT c.country, count(*) FROM customers c JOIN orders o ON o.customer_id = c.id GROUP BY c.country", false),
|
||||
new Attack("(legit) harmless comment", "SELECT 1 /* ; DROP TABLE orders */", false),
|
||||
new Attack("(legit) date_part, not on the list", "SELECT date_part('month', ordered_at) FROM orders LIMIT 5", false));
|
||||
|
||||
static String run(Attack a, boolean admin, boolean guard) {
|
||||
if (guard && !GuardCasesTest.GUARD.check(a.sql()).ok()) {
|
||||
return "guard";
|
||||
}
|
||||
try (Connection c = admin ? Db.admin() : Db.reader()) {
|
||||
if (admin) {
|
||||
try (var st = c.createStatement()) {
|
||||
st.execute("SET statement_timeout = '4s'"); // test harness limit so the suite finishes
|
||||
}
|
||||
}
|
||||
QueryRunner.run(c, a.sql(), 100);
|
||||
return a.harmful() ? "HARM" : "ok";
|
||||
} catch (SQLException e) {
|
||||
if (admin && e.getMessage().contains("statement timeout")) {
|
||||
return "HARM"; // it was still running when the harness stopped it
|
||||
}
|
||||
return "db";
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void fourSetups() throws Exception {
|
||||
try (Transcript t = new Transcript("02-attack-matrix.txt", "Sixteen queries against four setups on a real PostgreSQL 16")) {
|
||||
t.line("%-44s %-8s %-12s %-10s %s", "query", "neither", "guard only", "role only", "both");
|
||||
int[] harm = new int[4];
|
||||
List<String> roleErrors = new ArrayList<>();
|
||||
for (Attack a : ATTACKS) {
|
||||
String[] cells = { run(a, true, false), run(a, true, true), run(a, false, false), run(a, false, true) };
|
||||
for (int i = 0; i < 4; i++) {
|
||||
harm[i] += cells[i].equals("HARM") ? 1 : 0;
|
||||
}
|
||||
t.line("%-44s %-8s %-12s %-10s %s", a.name(), cells[0], cells[1], cells[2], cells[3]);
|
||||
if (cells[2].equals("db")) {
|
||||
try (Connection c = Db.reader()) {
|
||||
QueryRunner.run(c, a.sql(), 100);
|
||||
} catch (SQLException e) {
|
||||
roleErrors.add(String.format("%-44s %s", a.name(), TextToSql.firstLine(e.getMessage())));
|
||||
}
|
||||
}
|
||||
}
|
||||
t.blank().line("harmful outcomes out of 13 harmful queries: neither %d, guard only %d, role only %d, both %d", harm[0], harm[1], harm[2], harm[3]);
|
||||
t.blank().line("what the database said to the read-only role:");
|
||||
roleErrors.forEach(l -> t.line(" %s", l));
|
||||
assertThat(harm[0]).isEqualTo(13);
|
||||
assertThat(harm[3]).isZero();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package com.ankurm.texttosql;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.math.BigDecimal;
|
||||
import java.sql.Connection;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
|
||||
import com.ankurm.texttosql.support.ScriptedModel;
|
||||
import com.ankurm.texttosql.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
|
||||
/**
|
||||
* Evaluation by execution: run the reference query and the model's query and compare RESULTS.
|
||||
* The "model" is a script of eight answers I wrote (some right, some wrong in ways models are), so
|
||||
* these numbers demonstrate the metric, not any model's ability.
|
||||
*/
|
||||
class EvaluationTest {
|
||||
|
||||
record Case(String question, String reference, String modelSql, boolean ordered) {
|
||||
}
|
||||
|
||||
static final List<Case> CASES = List.of(
|
||||
new Case("How many customers are in Germany?",
|
||||
"SELECT count(*) FROM customers WHERE country = 'DE'",
|
||||
"SELECT COUNT(*) AS n FROM customers WHERE country = 'DE'", false),
|
||||
new Case("What is the total revenue of paid orders?",
|
||||
"SELECT sum(oi.quantity * p.price) FROM orders o JOIN order_items oi ON oi.order_id = o.id JOIN products p ON p.id = oi.product_id WHERE o.status = 'paid'",
|
||||
"SELECT SUM(p.price * i.quantity) AS revenue FROM order_items i JOIN products p ON p.id = i.product_id WHERE i.order_id IN (SELECT id FROM orders WHERE status = 'paid')", false),
|
||||
new Case("Which 3 products sold the most units?",
|
||||
"SELECT p.name, sum(oi.quantity) AS units FROM products p JOIN order_items oi ON oi.product_id = p.id GROUP BY p.name ORDER BY units DESC, p.name LIMIT 3",
|
||||
"SELECT p.name, SUM(quantity) FROM order_items JOIN products p ON p.id = product_id GROUP BY p.name ORDER BY SUM(quantity) DESC, p.name LIMIT 3", true),
|
||||
new Case("Which customers have never ordered?",
|
||||
"SELECT id FROM customers c WHERE NOT EXISTS (SELECT 1 FROM orders o WHERE o.customer_id = c.id) ORDER BY id",
|
||||
"SELECT c.id FROM customers c LEFT JOIN orders o ON o.customer_id = c.id WHERE o.id IS NULL ORDER BY c.id", true),
|
||||
new Case("How many orders are there per status?",
|
||||
"SELECT status, count(*) FROM orders GROUP BY status",
|
||||
"SELECT o.status, count(*) FROM orders o JOIN order_items i ON i.order_id = o.id GROUP BY o.status", false),
|
||||
new Case("What is the average order value?",
|
||||
"SELECT avg(t) FROM (SELECT sum(oi.quantity * p.price) AS t FROM order_items oi JOIN products p ON p.id = oi.product_id GROUP BY oi.order_id) x",
|
||||
"SELECT avg(price) FROM products", false),
|
||||
new Case("What is the revenue by country?",
|
||||
"SELECT c.country, sum(oi.quantity * p.price) FROM customers c JOIN orders o ON o.customer_id = c.id JOIN order_items oi ON oi.order_id = o.id JOIN products p ON p.id = oi.product_id GROUP BY c.country",
|
||||
"SELECT country, sum(amount) FROM sales GROUP BY country", false),
|
||||
new Case("How many orders were placed in March 2025?",
|
||||
"SELECT count(*) FROM orders WHERE ordered_at >= DATE '2025-03-01' AND ordered_at < DATE '2025-04-01'",
|
||||
"SELECT count(*) FROM orders WHERE to_char(ordered_at, 'YYYY-MM') = '2025-03'", false));
|
||||
|
||||
static String norm(String s) {
|
||||
return s.toLowerCase().replaceAll("\\s+", " ").replaceAll("\\s*;\\s*$", "").strip();
|
||||
}
|
||||
|
||||
static List<List<String>> canonical(QueryRunner.Rows rows, boolean ordered) {
|
||||
List<List<String>> out = new ArrayList<>();
|
||||
for (List<String> r : rows.rows()) {
|
||||
List<String> cells = new ArrayList<>();
|
||||
for (String v : r) {
|
||||
try {
|
||||
cells.add(new BigDecimal(v).stripTrailingZeros().toPlainString());
|
||||
} catch (RuntimeException e) {
|
||||
cells.add(v);
|
||||
}
|
||||
}
|
||||
out.add(cells);
|
||||
}
|
||||
if (!ordered) {
|
||||
out.sort((a, b) -> a.toString().compareTo(b.toString()));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
@Test
|
||||
void executionMatchBeatsStringMatch() throws Exception {
|
||||
Map<String, String> script = new LinkedHashMap<>();
|
||||
CASES.forEach(c -> script.put(c.question(), "```sql\n" + c.modelSql() + "\n```"));
|
||||
ScriptedModel model = new ScriptedModel(script::get);
|
||||
try (Connection db = Db.reader(); Transcript t = new Transcript("06-evaluation.txt",
|
||||
"Eight questions, scored three ways (the model is a script, not a language model)")) {
|
||||
String schema = SchemaPrompt.describe(db, TextToSql.TABLES);
|
||||
TextToSql service = new TextToSql(ChatClient.create(model), GuardCasesTest.GUARD, schema, 100);
|
||||
t.line("%-46s %-8s %-10s %s", "question", "string", "execution", "what happened");
|
||||
int string = 0, exec = 0, refused = 0, silentWrong = 0;
|
||||
for (Case c : CASES) {
|
||||
TextToSql.Answer a = service.answer(c.question(), db);
|
||||
boolean sameString = norm(a.sql()).equals(norm(c.reference()));
|
||||
boolean sameResult = false;
|
||||
String note;
|
||||
if (a.rejectedBy() != null) {
|
||||
note = "refused by the guard: " + a.rejectedBy();
|
||||
refused++;
|
||||
} else {
|
||||
QueryRunner.Rows ref = QueryRunner.run(db, c.reference(), 100);
|
||||
sameResult = canonical(a.rows(), c.ordered()).equals(canonical(ref, c.ordered()));
|
||||
note = sameResult ? "same result" : "DIFFERENT result: got " + a.rows().rows().getFirst() + ", expected " + ref.rows().getFirst();
|
||||
silentWrong += sameResult ? 0 : 1;
|
||||
}
|
||||
string += sameString ? 1 : 0;
|
||||
exec += sameResult ? 1 : 0;
|
||||
t.line("%-46s %-8s %-10s %s", c.question(), sameString ? "match" : "no", sameResult ? "match" : "no", note);
|
||||
}
|
||||
t.blank().line("string match: %d/8 execution match: %d/8 refused before running: %d ran and returned a wrong answer: %d", string, exec, refused, silentWrong);
|
||||
assertThat(string).isZero();
|
||||
assertThat(exec).isEqualTo(5);
|
||||
assertThat(silentWrong).isEqualTo(2);
|
||||
assertThat(Set.of(refused)).containsExactly(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package com.ankurm.texttosql;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
|
||||
import com.ankurm.texttosql.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
/** What the parser-based guard accepts and rejects, including legitimate queries it wrongly refuses. */
|
||||
class GuardCasesTest {
|
||||
|
||||
static final SqlGuard GUARD = new SqlGuard(Set.copyOf(TextToSql.TABLES), SqlGuard.DEFAULT_FUNCTIONS);
|
||||
|
||||
@Test
|
||||
void acceptedAndRejected() {
|
||||
Map<String, Boolean> cases = new LinkedHashMap<>();
|
||||
// should pass
|
||||
cases.put("SELECT count(*) FROM orders WHERE status = 'paid'", true);
|
||||
cases.put("SELECT c.country, count(*) FROM customers c JOIN orders o ON o.customer_id = c.id GROUP BY c.country ORDER BY 2 DESC", true);
|
||||
cases.put("WITH big AS (SELECT order_id FROM order_items GROUP BY order_id HAVING sum(quantity) > 5) SELECT count(*) FROM big", true);
|
||||
cases.put("SELECT * FROM orders WHERE id IN (SELECT order_id FROM order_items WHERE quantity > 2)", true);
|
||||
cases.put("SELECT 1 /* ; DROP TABLE orders */", true);
|
||||
cases.put("SELECT name FROM public.customers", true);
|
||||
// should be refused
|
||||
cases.put("DELETE FROM orders", false);
|
||||
cases.put("SELECT * FROM orders; DROP TABLE orders", false);
|
||||
cases.put("SELECT * FROM api_keys", false);
|
||||
cases.put("SELECT usename FROM pg_user", false);
|
||||
cases.put("SELECT pg_sleep(3)", false);
|
||||
cases.put("SELECT pg_read_file('/etc/passwd')", false);
|
||||
cases.put("WITH d AS (DELETE FROM orders RETURNING *) SELECT count(*) FROM d", false);
|
||||
cases.put("SELECT * FROM orders FOR UPDATE", false);
|
||||
cases.put("SELECT * INTO newtab FROM orders", false);
|
||||
cases.put("SELECT count(*) FROM generate_series(1, 2000000000)", false);
|
||||
cases.put("SELECT set_config('default_transaction_read_only', 'off', false)", false);
|
||||
cases.put("SELECT * FROM orders o, \"api_keys\" k", false);
|
||||
// legitimate but refused (cost of an allow-list)
|
||||
cases.put("SELECT date_part('month', ordered_at) FROM orders", false);
|
||||
cases.put("SELECT now()", false);
|
||||
try (Transcript t = new Transcript("03-guard-cases.txt", "What the parser guard accepts and refuses")) {
|
||||
for (Map.Entry<String, Boolean> e : cases.entrySet()) {
|
||||
SqlGuard.Verdict v = GUARD.check(e.getKey());
|
||||
t.line("%-8s %-110s %s", v.ok() ? "ACCEPT" : "REFUSE", e.getKey().length() > 108 ? e.getKey().substring(0, 105) + "..." : e.getKey(),
|
||||
v.ok() ? "" : "(" + v.reason() + ")");
|
||||
assertThat(v.ok()).as(e.getKey()).isEqualTo(e.getValue());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package com.ankurm.texttosql;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.sql.Connection;
|
||||
import java.sql.SQLException;
|
||||
|
||||
import com.ankurm.texttosql.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
/** Row cap and time limit, enforced below the guard: these are the database's and the driver's job. */
|
||||
class LimitsTest {
|
||||
|
||||
@Test
|
||||
void rowCapAndTimeout() throws Exception {
|
||||
try (Connection c = Db.reader(); Transcript t = new Transcript("04-limits.txt", "Row cap and statement timeout (cap = 100 rows)")) {
|
||||
QueryRunner.Rows many = QueryRunner.run(c, "SELECT id FROM orders ORDER BY id", 100);
|
||||
t.line("1000-row table, cap 100: rows returned %d, truncated %s", many.rows().size(), many.truncated());
|
||||
QueryRunner.Rows exact = QueryRunner.run(c, "SELECT id FROM orders WHERE id <= 100 ORDER BY id", 100);
|
||||
t.line("exactly 100 matching rows, cap 100: rows returned %d, truncated %s", exact.rows().size(), exact.truncated());
|
||||
QueryRunner.Rows few = QueryRunner.run(c, "SELECT id FROM orders WHERE id <= 7", 100);
|
||||
t.line("7 matching rows, cap 100: rows returned %d, truncated %s", few.rows().size(), few.truncated());
|
||||
assertThat(many.rows()).hasSize(100);
|
||||
assertThat(many.truncated()).isTrue();
|
||||
assertThat(exact.truncated()).isFalse();
|
||||
|
||||
long t0 = System.nanoTime();
|
||||
QueryRunner.Rows lazy = QueryRunner.run(c, "SELECT a.id FROM orders a, orders b, orders c", 100);
|
||||
long lazyMs = (System.nanoTime() - t0) / 1_000_000;
|
||||
t.line("a billion-row join, cap 100: rows returned %d, truncated %s (the server stopped producing rows)", lazy.rows().size(), lazy.truncated());
|
||||
assertThat(lazy.rows()).hasSize(100);
|
||||
assertThat(lazyMs).isLessThan(1500);
|
||||
|
||||
String srf = "";
|
||||
try {
|
||||
QueryRunner.run(c, "SELECT * FROM generate_series(1, 2000000000)", 100);
|
||||
} catch (SQLException e) {
|
||||
srf = TextToSql.firstLine(e.getMessage());
|
||||
}
|
||||
t.line("2 billion generated rows, cap 100: %s (the row cap did not help here)", srf);
|
||||
assertThat(srf).contains("statement timeout");
|
||||
|
||||
long t1 = System.nanoTime();
|
||||
String error = "";
|
||||
try {
|
||||
QueryRunner.run(c, "SELECT count(*) FROM generate_series(1, 2000000000)", 100);
|
||||
} catch (SQLException e) {
|
||||
error = TextToSql.firstLine(e.getMessage());
|
||||
}
|
||||
long ms = (System.nanoTime() - t1) / 1_000_000;
|
||||
t.line("count(*) over 2 billion rows: %s", error);
|
||||
t.line("stopped within 5 seconds of starting: %s (role setting statement_timeout = 2s)", ms < 5000);
|
||||
assertThat(error).contains("statement timeout");
|
||||
assertThat(ms).isBetween(1500L, 5000L);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package com.ankurm.texttosql;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.sql.Connection;
|
||||
import java.sql.ResultSet;
|
||||
import java.sql.SQLException;
|
||||
import java.sql.Statement;
|
||||
|
||||
import com.ankurm.texttosql.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
/** The read-only role has two independent walls: a session default that can be switched off, and privileges that cannot. */
|
||||
class RoleTest {
|
||||
|
||||
private static String attempt(Statement st, String sql) {
|
||||
try {
|
||||
if (st.execute(sql) && st.getResultSet() != null) {
|
||||
try (ResultSet rs = st.getResultSet()) {
|
||||
rs.next();
|
||||
return "ok: " + rs.getString(1);
|
||||
}
|
||||
}
|
||||
return "ok";
|
||||
} catch (SQLException e) {
|
||||
return TextToSql.firstLine(e.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void twoWalls() throws Exception {
|
||||
try (Connection c = Db.reader(); Statement st = c.createStatement(); Transcript t = new Transcript("05-role.txt", "What the t2s_reader role can and cannot do")) {
|
||||
c.setAutoCommit(true);
|
||||
t.line("%-62s %s", "SHOW statement_timeout", attempt(st, "SHOW statement_timeout"));
|
||||
t.line("%-62s %s", "SHOW default_transaction_read_only", attempt(st, "SHOW default_transaction_read_only"));
|
||||
t.line("%-62s %s", "SELECT count(*) FROM orders", attempt(st, "SELECT count(*) FROM orders"));
|
||||
String insert1 = attempt(st, "INSERT INTO orders VALUES (9999, 1, 'paid', DATE '2025-01-01')");
|
||||
t.line("%-62s %s", "INSERT INTO orders ... (wall 1: read-only default)", insert1);
|
||||
t.line("%-62s %s", "SELECT set_config('default_transaction_read_only','off',false)",
|
||||
attempt(st, "SELECT set_config('default_transaction_read_only','off',false)"));
|
||||
t.line("%-62s %s", "SHOW default_transaction_read_only", attempt(st, "SHOW default_transaction_read_only"));
|
||||
String insert2 = attempt(st, "INSERT INTO orders VALUES (9999, 1, 'paid', DATE '2025-01-01')");
|
||||
t.line("%-62s %s", "INSERT INTO orders ... (wall 2: no INSERT privilege)", insert2);
|
||||
t.line("%-62s %s", "SELECT count(*) FROM api_keys", attempt(st, "SELECT count(*) FROM api_keys"));
|
||||
t.line("%-62s %s", "SELECT email FROM customers", attempt(st, "SELECT email FROM customers"));
|
||||
t.line("%-62s %s", "SELECT id, name FROM customers LIMIT 1", attempt(st, "SELECT id FROM customers LIMIT 1"));
|
||||
assertThat(insert1).contains("read-only transaction");
|
||||
assertThat(insert2).contains("permission denied for table orders");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package com.ankurm.texttosql;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.sql.Connection;
|
||||
|
||||
import com.ankurm.texttosql.support.ScriptedModel;
|
||||
import com.ankurm.texttosql.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
|
||||
/** The schema section of the prompt is built from the database, through the same role the assistant uses. */
|
||||
class SchemaPromptTest {
|
||||
|
||||
@Test
|
||||
void promptShowsOnlyWhatTheRoleMaySee() throws Exception {
|
||||
try (Connection db = Db.reader(); Transcript t = new Transcript("01-schema-prompt.txt", "The prompt the model receives")) {
|
||||
String schema = SchemaPrompt.describe(db, TextToSql.TABLES);
|
||||
ScriptedModel model = new ScriptedModel(q -> "SELECT count(*) FROM customers");
|
||||
TextToSql service = new TextToSql(ChatClient.create(model), GuardCasesTest.GUARD, schema, 100);
|
||||
service.answer("How many customers do we have?", db);
|
||||
String sent = model.systemPrompts.getFirst();
|
||||
for (String line : sent.strip().split("\\R")) {
|
||||
t.line("| %s", line);
|
||||
}
|
||||
t.blank().line("user message: %s", model.questions.getFirst());
|
||||
t.line("mentions api_keys: %s | mentions customers.email: %s", sent.contains("api_keys"), sent.contains("email"));
|
||||
assertThat(sent).doesNotContain("api_keys").doesNotContain("email").contains("TABLE orders").contains("one of: paid, shipped");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package com.ankurm.texttosql.support;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.function.Function;
|
||||
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.messages.Message;
|
||||
import org.springframework.ai.chat.messages.MessageType;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.model.Generation;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
|
||||
/**
|
||||
* A "model" that answers from a script: question in, SQL out. It is NOT a language model and says
|
||||
* nothing about how well one writes SQL; it stands for "a model produced this string", good or bad,
|
||||
* so the code AROUND the model can be tested. It also records what it was sent.
|
||||
*/
|
||||
public class ScriptedModel implements ChatModel {
|
||||
|
||||
private final Function<String, String> script;
|
||||
|
||||
public final List<String> systemPrompts = new ArrayList<>();
|
||||
|
||||
public final List<String> questions = new ArrayList<>();
|
||||
|
||||
public ScriptedModel(Function<String, String> script) {
|
||||
this.script = script;
|
||||
}
|
||||
|
||||
@Override
|
||||
public ChatResponse call(Prompt prompt) {
|
||||
String question = "";
|
||||
for (Message m : prompt.getInstructions()) {
|
||||
if (m.getMessageType() == MessageType.SYSTEM) {
|
||||
systemPrompts.add(m.getText());
|
||||
} else if (m.getMessageType() == MessageType.USER) {
|
||||
question = m.getText();
|
||||
}
|
||||
}
|
||||
questions.add(question);
|
||||
return new ChatResponse(List.of(new Generation(new AssistantMessage(script.apply(question)))));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package com.ankurm.texttosql.support;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.PrintWriter;
|
||||
import java.io.StringWriter;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
|
||||
/**
|
||||
* Writes a numbered transcript under {@code output/} (repository root, not {@code docs/}) and
|
||||
* echoes it to the console. Every console block quoted in the article comes out of one of these
|
||||
* files verbatim.
|
||||
*/
|
||||
public final class Transcript implements AutoCloseable {
|
||||
|
||||
private final Path path;
|
||||
private final StringWriter buffer = new StringWriter();
|
||||
private final PrintWriter out = new PrintWriter(buffer);
|
||||
|
||||
public Transcript(String fileName, String title) {
|
||||
this.path = Path.of("output", fileName);
|
||||
out.println("# " + title);
|
||||
out.println();
|
||||
}
|
||||
|
||||
public Transcript line(String format, Object... args) {
|
||||
out.println(args.length == 0 ? format : String.format(format, args));
|
||||
return this;
|
||||
}
|
||||
|
||||
public Transcript blank() {
|
||||
out.println();
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
out.flush();
|
||||
try {
|
||||
Files.createDirectories(path.getParent());
|
||||
Files.writeString(path, buffer.toString());
|
||||
} catch (IOException e) {
|
||||
throw new IllegalStateException("could not write " + path, e);
|
||||
}
|
||||
System.out.print(buffer);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user