Skip to main content

tursotui_sql/
query_kind.rs

1//! SQL statement classification by leading keyword.
2
3/// Detected query type — used for status bar messaging and execution routing.
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub enum QueryKind {
6    Select,
7    Explain,
8    Insert,
9    Update,
10    Delete,
11    Ddl,
12    Pragma,
13    Batch {
14        statement_count: usize,
15        has_trailing_select: bool,
16    },
17    Other,
18}
19
20/// Detect the query kind from the first non-whitespace, non-comment token.
21pub fn detect_query_kind(sql: &str) -> QueryKind {
22    let sql = skip_leading_whitespace_and_comments(sql);
23    let first_word = sql
24        .split(|c: char| !c.is_ascii_alphanumeric() && c != '_')
25        .next()
26        .unwrap_or("");
27    match first_word.to_uppercase().as_str() {
28        "SELECT" => QueryKind::Select,
29        "EXPLAIN" => QueryKind::Explain,
30        "INSERT" => QueryKind::Insert,
31        "UPDATE" => QueryKind::Update,
32        "DELETE" => QueryKind::Delete,
33        "CREATE" | "ALTER" | "DROP" => QueryKind::Ddl,
34        "PRAGMA" => QueryKind::Pragma,
35        _ => QueryKind::Other,
36    }
37}
38
39/// Returns true if `sql` begins a transaction-control statement
40/// (BEGIN/COMMIT/ROLLBACK/END), after stripping leading whitespace and comments.
41pub fn is_transaction_control(sql: &str) -> bool {
42    let upper = skip_leading_whitespace_and_comments(sql).to_uppercase();
43    upper.starts_with("BEGIN")
44        || upper.starts_with("COMMIT")
45        || upper.starts_with("ROLLBACK")
46        || upper.starts_with("END")
47}
48
49/// Skip leading whitespace and SQL comments (line and block) to find the first real token.
50pub(crate) fn skip_leading_whitespace_and_comments(sql: &str) -> &str {
51    let mut s = sql.trim_start();
52    loop {
53        if s.starts_with("--") {
54            // Skip to end of line
55            s = s.find('\n').map_or("", |i| &s[i + 1..]).trim_start();
56        } else if s.starts_with("/*") {
57            // Skip to end of block comment
58            s = match s[2..].find("*/") {
59                Some(i) => &s[2 + i + 2..],
60                None => "",
61            };
62            s = s.trim_start();
63        } else {
64            break;
65        }
66    }
67    s
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[test]
75    fn detect_query_kind_select() {
76        assert!(matches!(detect_query_kind("SELECT 1"), QueryKind::Select));
77    }
78
79    #[test]
80    fn detect_query_kind_case_insensitive() {
81        assert!(matches!(detect_query_kind("sElEcT 1"), QueryKind::Select));
82    }
83
84    #[test]
85    fn detect_query_kind_leading_whitespace() {
86        assert!(matches!(
87            detect_query_kind("  \n  INSERT INTO t VALUES (1)"),
88            QueryKind::Insert
89        ));
90    }
91
92    #[test]
93    fn detect_query_kind_leading_line_comment() {
94        assert!(matches!(
95            detect_query_kind("-- comment\nSELECT 1"),
96            QueryKind::Select
97        ));
98    }
99
100    #[test]
101    fn detect_query_kind_leading_block_comment() {
102        assert!(matches!(
103            detect_query_kind("/* block */  DELETE FROM t"),
104            QueryKind::Delete
105        ));
106    }
107
108    #[test]
109    fn detect_query_kind_explain() {
110        assert!(matches!(
111            detect_query_kind("EXPLAIN SELECT 1"),
112            QueryKind::Explain
113        ));
114    }
115
116    #[test]
117    fn detect_query_kind_ddl() {
118        assert!(matches!(
119            detect_query_kind("CREATE TABLE t (id INT)"),
120            QueryKind::Ddl
121        ));
122        assert!(matches!(
123            detect_query_kind("ALTER TABLE t ADD col INT"),
124            QueryKind::Ddl
125        ));
126        assert!(matches!(detect_query_kind("DROP TABLE t"), QueryKind::Ddl));
127    }
128
129    #[test]
130    fn detect_query_kind_pragma() {
131        assert!(matches!(
132            detect_query_kind("PRAGMA table_info(t)"),
133            QueryKind::Pragma
134        ));
135    }
136
137    #[test]
138    fn detect_query_kind_update() {
139        assert!(matches!(
140            detect_query_kind("UPDATE t SET x = 1"),
141            QueryKind::Update
142        ));
143    }
144
145    #[test]
146    fn detect_query_kind_unknown() {
147        assert!(matches!(detect_query_kind("VACUUM"), QueryKind::Other));
148        assert!(matches!(detect_query_kind(""), QueryKind::Other));
149    }
150
151    #[test]
152    fn detect_query_kind_begin_commit_rollback() {
153        assert!(matches!(detect_query_kind("BEGIN"), QueryKind::Other));
154        assert!(matches!(detect_query_kind("COMMIT"), QueryKind::Other));
155        assert!(matches!(detect_query_kind("ROLLBACK"), QueryKind::Other));
156        assert!(matches!(detect_query_kind("END"), QueryKind::Other));
157    }
158
159    #[test]
160    fn detect_query_kind_stacked_comments() {
161        assert!(matches!(
162            detect_query_kind("-- first\n-- second\nSELECT 1"),
163            QueryKind::Select
164        ));
165        assert!(matches!(
166            detect_query_kind("/* a */ /* b */ INSERT INTO t VALUES (1)"),
167            QueryKind::Insert
168        ));
169    }
170
171    #[test]
172    fn is_transaction_control_positive() {
173        assert!(is_transaction_control("BEGIN"));
174        assert!(is_transaction_control("COMMIT"));
175        assert!(is_transaction_control("ROLLBACK"));
176        assert!(is_transaction_control("END"));
177        assert!(is_transaction_control("  BEGIN TRANSACTION"));
178        assert!(is_transaction_control("-- comment\nBEGIN"));
179    }
180
181    #[test]
182    fn is_transaction_control_negative() {
183        assert!(!is_transaction_control("SELECT 1"));
184        assert!(!is_transaction_control("INSERT INTO t VALUES (1)"));
185        assert!(!is_transaction_control(""));
186    }
187}