tursotui_sql/
query_kind.rs1#[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
20pub 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
39pub 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
49pub(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 s = s.find('\n').map_or("", |i| &s[i + 1..]).trim_start();
56 } else if s.starts_with("/*") {
57 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}