netget/utils/sql.rs
1//! What kind of answer an SQL statement expects, from its leading keyword.
2//!
3//! The SQL servers (`postgresql`, `mysql`) hand the model the statement text and three ways to
4//! answer it: a result set, a command tag / OK packet, or an error. The real-model eval showed
5//! a small model answering `SELECT current_user` with the command tag `INSERT 0 1` — copied from
6//! the ok action's example — so the client got "INSERT 0 1" for a query that asked for a value.
7//! Both servers know the statement exactly, so each tells the model which shape answers it
8//! (`answer_with` in the event) instead of leaving it to map SQL to the wire in its head.
9//!
10//! This is a classifier over the first keyword, not a parser: it reads past leading whitespace,
11//! `--` and `/* */` comments and opening parentheses, and nothing else. A statement it cannot
12//! place is [`StatementShape::Unknown`], and the caller says nothing rather than guess.
13
14/// The answer a statement's leading keyword calls for.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum StatementShape {
17 /// Returns rows: `SELECT`, `SHOW`, `VALUES`, `TABLE`, `WITH`, `EXPLAIN`, `DESCRIBE`/`DESC`,
18 /// `FETCH`. Answered with a result set, even when the answer is one value or no rows.
19 Rows,
20 /// Changes or configures something and returns no rows: `INSERT`, `UPDATE`, `DELETE`,
21 /// `CREATE`, `SET`, `BEGIN`, …. Answered with a command tag / OK.
22 NoRows,
23 /// Empty, or a keyword this does not classify.
24 Unknown,
25}
26
27/// Classify `sql` by its first keyword. See the module docs for what "first" means.
28pub fn statement_shape(sql: &str) -> StatementShape {
29 let keyword = first_keyword(sql).to_ascii_uppercase();
30 match keyword.as_str() {
31 "SELECT" | "SHOW" | "VALUES" | "TABLE" | "WITH" | "EXPLAIN" | "DESCRIBE" | "DESC"
32 | "FETCH" => StatementShape::Rows,
33 "INSERT" | "UPDATE" | "DELETE" | "MERGE" | "CREATE" | "DROP" | "ALTER" | "TRUNCATE"
34 | "SET" | "BEGIN" | "START" | "COMMIT" | "ROLLBACK" | "SAVEPOINT" | "RELEASE" | "USE"
35 | "GRANT" | "REVOKE" | "LOCK" | "UNLOCK" | "DISCARD" | "RESET" | "DEALLOCATE"
36 | "PREPARE" | "COPY" | "VACUUM" | "ANALYZE" | "COMMENT" | "LISTEN" | "NOTIFY"
37 | "UNLISTEN" => StatementShape::NoRows,
38 _ => StatementShape::Unknown,
39 }
40}
41
42/// The first SQL keyword of `sql`, skipping whitespace, comments and opening parentheses.
43fn first_keyword(sql: &str) -> &str {
44 let mut rest = sql;
45 loop {
46 let trimmed = rest.trim_start_matches(|c: char| c.is_whitespace() || c == '(');
47 if let Some(after) = trimmed.strip_prefix("--") {
48 rest = after.split_once('\n').map(|(_, tail)| tail).unwrap_or("");
49 } else if let Some(after) = trimmed.strip_prefix("/*") {
50 rest = after.split_once("*/").map(|(_, tail)| tail).unwrap_or("");
51 } else {
52 rest = trimmed;
53 break;
54 }
55 }
56 let end = rest
57 .find(|c: char| !(c.is_ascii_alphanumeric() || c == '_'))
58 .unwrap_or(rest.len());
59 &rest[..end]
60}