1use crate::cli::DatabaseKind;
2
3const SQL_RESERVED: &[&str] = &[
7 "ABORT",
8 "ALL",
9 "ALTER",
10 "AND",
11 "ANY",
12 "AS",
13 "ASC",
14 "AUTHORIZATION",
15 "BEFORE",
16 "BEGIN",
17 "BETWEEN",
18 "BOTH",
19 "BY",
20 "CASE",
21 "CAST",
22 "CHECK",
23 "COLLATE",
24 "COLUMN",
25 "COMMIT",
26 "CONSTRAINT",
27 "CREATE",
28 "CROSS",
29 "CURRENT",
30 "CURRENT_DATE",
31 "CURRENT_ROLE",
32 "CURRENT_SCHEMA",
33 "CURRENT_TIME",
34 "CURRENT_TIMESTAMP",
35 "CURRENT_USER",
36 "DEFAULT",
37 "DEFERRABLE",
38 "DELETE",
39 "DESC",
40 "DISTINCT",
41 "DO",
42 "DROP",
43 "ELSE",
44 "END",
45 "EXCEPT",
46 "EXISTS",
47 "FALSE",
48 "FETCH",
49 "FOR",
50 "FOREIGN",
51 "FROM",
52 "FULL",
53 "GRANT",
54 "GROUP",
55 "HAVING",
56 "IF",
57 "IN",
58 "INDEX",
59 "INNER",
60 "INSERT",
61 "INTERSECT",
62 "INTO",
63 "IS",
64 "JOIN",
65 "KEY",
66 "LATERAL",
67 "LEADING",
68 "LEFT",
69 "LIKE",
70 "LIMIT",
71 "LOCALTIME",
72 "LOCALTIMESTAMP",
73 "NATURAL",
74 "NOT",
75 "NULL",
76 "OF",
77 "OFFSET",
78 "ON",
79 "ONLY",
80 "OR",
81 "ORDER",
82 "OUTER",
83 "OVERLAPS",
84 "PLACING",
85 "PRIMARY",
86 "REFERENCES",
87 "RETURNING",
88 "RIGHT",
89 "ROLLBACK",
90 "SCHEMA",
91 "SELECT",
92 "SESSION_USER",
93 "SET",
94 "SIMILAR",
95 "SOME",
96 "SYMMETRIC",
97 "TABLE",
98 "THEN",
99 "TO",
100 "TRAILING",
101 "TRIGGER",
102 "TRUE",
103 "UNION",
104 "UNIQUE",
105 "UPDATE",
106 "USER",
107 "USING",
108 "VALUES",
109 "VARIADIC",
110 "VIEW",
111 "WHEN",
112 "WHERE",
113 "WINDOW",
114 "WITH",
115];
116
117fn is_reserved_keyword(name: &str) -> bool {
120 let upper = name.to_ascii_uppercase();
121 SQL_RESERVED.binary_search(&upper.as_str()).is_ok()
122}
123
124pub fn is_safe_unquoted(name: &str, _db: DatabaseKind) -> bool {
130 if name.is_empty() {
131 return false;
132 }
133 let bytes = name.as_bytes();
134 let first = bytes[0];
135 let first_ok = first == b'_' || first.is_ascii_lowercase();
136 if !first_ok {
137 return false;
138 }
139 for &b in &bytes[1..] {
140 let ok = b == b'_' || b.is_ascii_lowercase() || b.is_ascii_digit();
141 if !ok {
142 return false;
143 }
144 }
145 !is_reserved_keyword(name)
146}
147
148pub fn quote_ident(name: &str, db: DatabaseKind) -> String {
152 if is_safe_unquoted(name, db) {
153 name.to_string()
154 } else {
155 quote_ident_always(name, db)
156 }
157}
158
159pub fn quote_ident_always(name: &str, db: DatabaseKind) -> String {
161 match db {
162 DatabaseKind::Mysql => format!("`{}`", name.replace('`', "``")),
163 DatabaseKind::Postgres | DatabaseKind::Sqlite => {
164 format!("\"{}\"", name.replace('"', "\"\""))
165 }
166 }
167}
168
169pub fn quote_qualified(schema: Option<&str>, table: &str, db: DatabaseKind) -> String {
172 match schema {
173 Some(s) => format!("{}.{}", quote_ident(s, db), quote_ident(table, db)),
174 None => quote_ident(table, db),
175 }
176}
177
178pub fn is_safe_ident(name: &str) -> bool {
184 !name.is_empty()
185 && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
186 && !name.starts_with(|c: char| c.is_ascii_digit())
187}
188
189#[cfg(test)]
190mod tests {
191 use super::*;
192
193 #[test]
196 fn safe_identifier_pg_not_quoted() {
197 assert_eq!(quote_ident("users", DatabaseKind::Postgres), "users");
198 assert_eq!(quote_ident("agent_id", DatabaseKind::Postgres), "agent_id");
199 assert_eq!(
200 quote_ident("agent__connector", DatabaseKind::Postgres),
201 "agent__connector"
202 );
203 }
204
205 #[test]
206 fn safe_identifier_mysql_not_quoted() {
207 assert_eq!(quote_ident("users", DatabaseKind::Mysql), "users");
208 }
209
210 #[test]
211 fn uppercase_identifier_pg_quoted() {
212 assert_eq!(quote_ident("Users", DatabaseKind::Postgres), "\"Users\"");
215 }
216
217 #[test]
218 fn reserved_word_quoted_in_pg() {
219 assert_eq!(quote_ident("select", DatabaseKind::Postgres), "\"select\"");
220 assert_eq!(quote_ident("user", DatabaseKind::Postgres), "\"user\"");
221 assert_eq!(quote_ident("order", DatabaseKind::Postgres), "\"order\"");
222 }
223
224 #[test]
225 fn reserved_word_quoted_in_mysql() {
226 assert_eq!(quote_ident("select", DatabaseKind::Mysql), "`select`");
227 }
228
229 #[test]
230 fn identifier_with_dash_quoted() {
231 assert_eq!(
232 quote_ident("user-id", DatabaseKind::Postgres),
233 "\"user-id\""
234 );
235 }
236
237 #[test]
238 fn identifier_starting_with_digit_quoted() {
239 assert_eq!(quote_ident("123abc", DatabaseKind::Postgres), "\"123abc\"");
240 }
241
242 #[test]
243 fn empty_identifier_quoted() {
244 assert_eq!(quote_ident("", DatabaseKind::Postgres), "\"\"");
247 }
248
249 #[test]
250 fn injection_attempt_quoted_and_escaped() {
251 assert_eq!(
252 quote_ident("user\"; DROP TABLE x; --", DatabaseKind::Postgres),
253 "\"user\"\"; DROP TABLE x; --\""
254 );
255 }
256
257 #[test]
260 fn always_quote_safe_identifier_pg() {
261 assert_eq!(
262 quote_ident_always("users", DatabaseKind::Postgres),
263 "\"users\""
264 );
265 }
266
267 #[test]
268 fn always_quote_safe_identifier_mysql() {
269 assert_eq!(quote_ident_always("users", DatabaseKind::Mysql), "`users`");
270 }
271
272 #[test]
273 fn always_quote_escapes_internal_backtick() {
274 assert_eq!(quote_ident_always("ev`il", DatabaseKind::Mysql), "`ev``il`");
275 }
276
277 #[test]
280 fn qualified_safe_idents_not_quoted() {
281 assert_eq!(
282 quote_qualified(Some("agent"), "agent_connector", DatabaseKind::Postgres),
283 "agent.agent_connector"
284 );
285 assert_eq!(
286 quote_qualified(Some("app"), "users", DatabaseKind::Mysql),
287 "app.users"
288 );
289 }
290
291 #[test]
292 fn qualified_with_reserved_schema_quoted() {
293 assert_eq!(
294 quote_qualified(Some("user"), "items", DatabaseKind::Postgres),
295 "\"user\".items"
296 );
297 }
298
299 #[test]
300 fn qualified_without_schema() {
301 assert_eq!(quote_qualified(None, "users", DatabaseKind::Mysql), "users");
302 }
303
304 #[test]
307 fn safe_ident_rejects_dash() {
308 assert!(!is_safe_ident("user-id"));
309 }
310
311 #[test]
312 fn safe_ident_rejects_leading_digit() {
313 assert!(!is_safe_ident("123abc"));
314 }
315
316 #[test]
317 fn safe_ident_rejects_empty() {
318 assert!(!is_safe_ident(""));
319 }
320
321 #[test]
322 fn safe_ident_accepts_underscore_prefix() {
323 assert!(is_safe_ident("_private"));
324 }
325
326 #[test]
327 fn safe_ident_accepts_mixed_case() {
328 assert!(is_safe_ident("UserAccount2"));
329 }
330
331 #[test]
334 fn reserved_list_is_sorted() {
335 for pair in SQL_RESERVED.windows(2) {
336 assert!(
337 pair[0] < pair[1],
338 "SQL_RESERVED must be sorted; '{}' >= '{}'",
339 pair[0],
340 pair[1]
341 );
342 }
343 }
344}