Skip to main content

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}