1use std::ops::Deref;
4
5use crate::error::Error;
6use crate::ffi;
7use crate::types::{ToSql, ToSqlOutput, ValueRef};
8use crate::{Connection, DatabaseName, Result, Row};
9
10pub struct Sql {
11 buf: String,
12}
13
14impl Sql {
15 pub fn new() -> Sql {
16 Sql { buf: String::new() }
17 }
18
19 pub fn push_pragma(
20 &mut self,
21 schema_name: Option<DatabaseName<'_>>,
22 pragma_name: &str,
23 ) -> Result<()> {
24 self.push_keyword("PRAGMA")?;
25 self.push_space();
26 if let Some(schema_name) = schema_name {
27 self.push_schema_name(schema_name);
28 self.push_dot();
29 }
30 self.push_keyword(pragma_name)
31 }
32
33 pub fn push_keyword(&mut self, keyword: &str) -> Result<()> {
34 if !keyword.is_empty() && is_identifier(keyword) {
35 self.buf.push_str(keyword);
36 Ok(())
37 } else {
38 Err(Error::SqliteFailure(
39 ffi::Error::new(ffi::SQLITE_MISUSE),
40 Some(format!("Invalid keyword \"{}\"", keyword)),
41 ))
42 }
43 }
44
45 pub fn push_schema_name(&mut self, schema_name: DatabaseName<'_>) {
46 match schema_name {
47 DatabaseName::Main => self.buf.push_str("main"),
48 DatabaseName::Temp => self.buf.push_str("temp"),
49 DatabaseName::Attached(s) => self.push_identifier(s),
50 };
51 }
52
53 pub fn push_identifier(&mut self, s: &str) {
54 if is_identifier(s) {
55 self.buf.push_str(s);
56 } else {
57 self.wrap_and_escape(s, '"');
58 }
59 }
60
61 pub fn push_value(&mut self, value: &dyn ToSql) -> Result<()> {
62 let value = value.to_sql()?;
63 let value = match value {
64 ToSqlOutput::Borrowed(v) => v,
65 ToSqlOutput::Owned(ref v) => ValueRef::from(v),
66 #[cfg(feature = "blob")]
67 ToSqlOutput::ZeroBlob(_) => {
68 return Err(Error::SqliteFailure(
69 ffi::Error::new(ffi::SQLITE_MISUSE),
70 Some(format!("Unsupported value \"{:?}\"", value)),
71 ));
72 }
73 #[cfg(feature = "array")]
74 ToSqlOutput::Array(_) => {
75 return Err(Error::SqliteFailure(
76 ffi::Error::new(ffi::SQLITE_MISUSE),
77 Some(format!("Unsupported value \"{:?}\"", value)),
78 ));
79 }
80 };
81 match value {
82 ValueRef::Integer(i) => {
83 self.push_int(i);
84 }
85 ValueRef::Real(r) => {
86 self.push_real(r);
87 }
88 ValueRef::Text(s) => {
89 let s = std::str::from_utf8(s)?;
90 self.push_string_literal(s);
91 }
92 _ => {
93 return Err(Error::SqliteFailure(
94 ffi::Error::new(ffi::SQLITE_MISUSE),
95 Some(format!("Unsupported value \"{:?}\"", value)),
96 ));
97 }
98 };
99 Ok(())
100 }
101
102 pub fn push_string_literal(&mut self, s: &str) {
103 self.wrap_and_escape(s, '\'');
104 }
105
106 pub fn push_int(&mut self, i: i64) {
107 self.buf.push_str(&i.to_string());
108 }
109
110 pub fn push_real(&mut self, f: f64) {
111 self.buf.push_str(&f.to_string());
112 }
113
114 pub fn push_space(&mut self) {
115 self.buf.push(' ');
116 }
117
118 pub fn push_dot(&mut self) {
119 self.buf.push('.');
120 }
121
122 pub fn push_equal_sign(&mut self) {
123 self.buf.push('=');
124 }
125
126 pub fn open_brace(&mut self) {
127 self.buf.push('(');
128 }
129
130 pub fn close_brace(&mut self) {
131 self.buf.push(')');
132 }
133
134 pub fn as_str(&self) -> &str {
135 &self.buf
136 }
137
138 fn wrap_and_escape(&mut self, s: &str, quote: char) {
139 self.buf.push(quote);
140 let chars = s.chars();
141 for ch in chars {
142 if ch == quote {
144 self.buf.push(ch);
145 }
146 self.buf.push(ch)
147 }
148 self.buf.push(quote);
149 }
150}
151
152impl Deref for Sql {
153 type Target = str;
154
155 fn deref(&self) -> &str {
156 self.as_str()
157 }
158}
159
160impl Connection {
161 pub fn pragma_query_value<T, F>(
169 &self,
170 schema_name: Option<DatabaseName<'_>>,
171 pragma_name: &str,
172 f: F,
173 ) -> Result<T>
174 where
175 F: FnOnce(&Row<'_>) -> Result<T>,
176 {
177 let mut query = Sql::new();
178 query.push_pragma(schema_name, pragma_name)?;
179 self.query_row(&query, [], f)
180 }
181
182 pub fn pragma_query<F>(
187 &self,
188 schema_name: Option<DatabaseName<'_>>,
189 pragma_name: &str,
190 mut f: F,
191 ) -> Result<()>
192 where
193 F: FnMut(&Row<'_>) -> Result<()>,
194 {
195 let mut query = Sql::new();
196 query.push_pragma(schema_name, pragma_name)?;
197 let mut stmt = self.prepare(&query)?;
198 let mut rows = stmt.query([])?;
199 while let Some(result_row) = rows.next()? {
200 let row = result_row;
201 f(&row)?;
202 }
203 Ok(())
204 }
205
206 pub fn pragma<F>(
216 &self,
217 schema_name: Option<DatabaseName<'_>>,
218 pragma_name: &str,
219 pragma_value: &dyn ToSql,
220 mut f: F,
221 ) -> Result<()>
222 where
223 F: FnMut(&Row<'_>) -> Result<()>,
224 {
225 let mut sql = Sql::new();
226 sql.push_pragma(schema_name, pragma_name)?;
227 sql.open_brace();
231 sql.push_value(pragma_value)?;
232 sql.close_brace();
233 let mut stmt = self.prepare(&sql)?;
234 let mut rows = stmt.query([])?;
235 while let Some(result_row) = rows.next()? {
236 let row = result_row;
237 f(&row)?;
238 }
239 Ok(())
240 }
241
242 pub fn pragma_update(
247 &self,
248 schema_name: Option<DatabaseName<'_>>,
249 pragma_name: &str,
250 pragma_value: &dyn ToSql,
251 ) -> Result<()> {
252 let mut sql = Sql::new();
253 sql.push_pragma(schema_name, pragma_name)?;
254 sql.push_equal_sign();
258 sql.push_value(pragma_value)?;
259 self.execute_batch(&sql)
260 }
261
262 pub fn pragma_update_and_check<F, T>(
266 &self,
267 schema_name: Option<DatabaseName<'_>>,
268 pragma_name: &str,
269 pragma_value: &dyn ToSql,
270 f: F,
271 ) -> Result<T>
272 where
273 F: FnOnce(&Row<'_>) -> Result<T>,
274 {
275 let mut sql = Sql::new();
276 sql.push_pragma(schema_name, pragma_name)?;
277 sql.push_equal_sign();
281 sql.push_value(pragma_value)?;
282 self.query_row(&sql, [], f)
283 }
284}
285
286fn is_identifier(s: &str) -> bool {
287 let chars = s.char_indices();
288 for (i, ch) in chars {
289 if i == 0 {
290 if !is_identifier_start(ch) {
291 return false;
292 }
293 } else if !is_identifier_continue(ch) {
294 return false;
295 }
296 }
297 true
298}
299
300fn is_identifier_start(c: char) -> bool {
301 (c >= 'A' && c <= 'Z') || c == '_' || (c >= 'a' && c <= 'z') || c > '\x7F'
302}
303
304fn is_identifier_continue(c: char) -> bool {
305 c == '$'
306 || (c >= '0' && c <= '9')
307 || (c >= 'A' && c <= 'Z')
308 || c == '_'
309 || (c >= 'a' && c <= 'z')
310 || c > '\x7F'
311}
312
313#[cfg(test)]
314mod test {
315 use super::Sql;
316 use crate::pragma;
317 use crate::{Connection, DatabaseName, Result};
318
319 #[test]
320 fn pragma_query_value() -> Result<()> {
321 let db = Connection::open_in_memory()?;
322 let user_version: i32 = db.pragma_query_value(None, "user_version", |row| row.get(0))?;
323 assert_eq!(0, user_version);
324 Ok(())
325 }
326
327 #[test]
328 #[cfg(feature = "modern_sqlite")]
329 fn pragma_func_query_value() -> Result<()> {
330 let db = Connection::open_in_memory()?;
331 let user_version: i32 =
332 db.query_row("SELECT user_version FROM pragma_user_version", [], |row| {
333 row.get(0)
334 })?;
335 assert_eq!(0, user_version);
336 Ok(())
337 }
338
339 #[test]
340 fn pragma_query_no_schema() -> Result<()> {
341 let db = Connection::open_in_memory()?;
342 let mut user_version = -1;
343 db.pragma_query(None, "user_version", |row| {
344 user_version = row.get(0)?;
345 Ok(())
346 })?;
347 assert_eq!(0, user_version);
348 Ok(())
349 }
350
351 #[test]
352 fn pragma_query_with_schema() -> Result<()> {
353 let db = Connection::open_in_memory()?;
354 let mut user_version = -1;
355 db.pragma_query(Some(DatabaseName::Main), "user_version", |row| {
356 user_version = row.get(0)?;
357 Ok(())
358 })?;
359 assert_eq!(0, user_version);
360 Ok(())
361 }
362
363 #[test]
364 fn pragma() -> Result<()> {
365 let db = Connection::open_in_memory()?;
366 let mut columns = Vec::new();
367 db.pragma(None, "table_info", &"sqlite_master", |row| {
368 let column: String = row.get(1)?;
369 columns.push(column);
370 Ok(())
371 })?;
372 assert_eq!(5, columns.len());
373 Ok(())
374 }
375
376 #[test]
377 #[cfg(feature = "modern_sqlite")]
378 fn pragma_func() -> Result<()> {
379 let db = Connection::open_in_memory()?;
380 let mut table_info = db.prepare("SELECT * FROM pragma_table_info(?)")?;
381 let mut columns = Vec::new();
382 let mut rows = table_info.query(&["sqlite_master"])?;
383
384 while let Some(row) = rows.next()? {
385 let row = row;
386 let column: String = row.get(1)?;
387 columns.push(column);
388 }
389 assert_eq!(5, columns.len());
390 Ok(())
391 }
392
393 #[test]
394 fn pragma_update() -> Result<()> {
395 let db = Connection::open_in_memory()?;
396 db.pragma_update(None, "user_version", &1)
397 }
398
399 #[test]
400 fn pragma_update_and_check() -> Result<()> {
401 let db = Connection::open_in_memory()?;
402 let journal_mode: String =
403 db.pragma_update_and_check(None, "journal_mode", &"OFF", |row| row.get(0))?;
404 assert_eq!("off", &journal_mode);
405 Ok(())
406 }
407
408 #[test]
409 fn is_identifier() {
410 assert!(pragma::is_identifier("full"));
411 assert!(pragma::is_identifier("r2d2"));
412 assert!(!pragma::is_identifier("sp ce"));
413 assert!(!pragma::is_identifier("semi;colon"));
414 }
415
416 #[test]
417 fn double_quote() {
418 let mut sql = Sql::new();
419 sql.push_schema_name(DatabaseName::Attached(r#"schema";--"#));
420 assert_eq!(r#""schema"";--""#, sql.as_str());
421 }
422
423 #[test]
424 fn wrap_and_escape() {
425 let mut sql = Sql::new();
426 sql.push_string_literal("value'; --");
427 assert_eq!("'value''; --'", sql.as_str());
428 }
429
430 #[test]
431 fn locking_mode() -> Result<()> {
432 let db = Connection::open_in_memory()?;
433 let r = db.pragma_update(None, "locking_mode", &"exclusive");
434 if cfg!(feature = "extra_check") {
435 r.unwrap_err();
436 } else {
437 r?;
438 }
439 Ok(())
440 }
441}