Skip to main content

rusqlite/
pragma.rs

1//! Pragma helpers
2
3use std::ops::Deref;
4
5use crate::ffi;
6use crate::types::{FromSql, ToSql, ToSqlOutput, ValueRef};
7use crate::{Connection, Result, Row};
8
9pub struct Sql {
10    buf: String,
11}
12
13impl Sql {
14    pub fn new() -> Self {
15        Self { buf: String::new() }
16    }
17
18    pub fn push_pragma(
19        &mut self,
20        _conn: &Connection,
21        schema_name: Option<&str>,
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        cfg_select! {
31            any(
32                feature = "bundled-sqlcipher",
33                feature = "bundled-sqlcipher-vendored-openssl",
34                not(feature = "extra_check")
35            ) => {
36                if !pragma_name.is_empty() && is_identifier(pragma_name) {
37                    self.buf.push_str(pragma_name);
38                    Ok(())
39                } else {
40                    Err(err!(ffi::SQLITE_MISUSE, "Invalid pragma \"{pragma_name}\""))
41                }
42            }
43            _ => {
44                const SQL: &str = "SELECT 1 FROM pragma_pragma_list WHERE name = ?";
45                let mut stmt = cfg_select! {
46                    feature = "cache" => {
47                        // no_fmt
48                        _conn.prepare_cached(SQL)
49                    }
50                    _ => {
51                        // no_fmt
52                        _conn.prepare(SQL)
53                    }
54                }?;
55                if stmt.exists([pragma_name])? {
56                    self.buf.push_str(pragma_name);
57                    Ok(())
58                } else {
59                    Err(err!(ffi::SQLITE_MISUSE, "Invalid pragma \"{pragma_name}\""))
60                }
61            }
62        }
63    }
64
65    pub fn push_keyword(&mut self, keyword: &str) -> Result<()> {
66        if !keyword.is_empty() && is_keyword(keyword) {
67            self.buf.push_str(keyword);
68            Ok(())
69        } else {
70            Err(err!(ffi::SQLITE_MISUSE, "Invalid keyword \"{keyword}\""))
71        }
72    }
73
74    pub fn push_schema_name(&mut self, schema_name: &str) {
75        self.push_identifier(schema_name);
76    }
77
78    pub fn push_identifier(&mut self, s: &str) {
79        if is_identifier(s) {
80            self.buf.push_str(s);
81        } else {
82            self.wrap_and_escape(s, '"');
83        }
84    }
85
86    pub fn push_value(&mut self, value: &dyn ToSql) -> Result<()> {
87        let value = value.to_sql()?;
88        let value = match value {
89            ToSqlOutput::Borrowed(v) => v,
90            ToSqlOutput::Owned(ref v) => ValueRef::from(v),
91            #[cfg(any(feature = "blob", feature = "functions", feature = "pointer"))]
92            _ => {
93                return Err(err!(ffi::SQLITE_MISUSE, "Unsupported value \"{value:?}\""));
94            }
95        };
96        match value {
97            ValueRef::Integer(i) => {
98                self.push_int(i);
99            }
100            ValueRef::Real(r) => {
101                self.push_real(r);
102            }
103            ValueRef::Text(s) => {
104                let s = std::str::from_utf8(s)?;
105                self.push_string_literal(s);
106            }
107            _ => {
108                return Err(err!(ffi::SQLITE_MISUSE, "Unsupported value \"{value:?}\""));
109            }
110        }
111        Ok(())
112    }
113
114    pub fn push_string_literal(&mut self, s: &str) {
115        self.wrap_and_escape(s, '\'');
116    }
117
118    pub fn push_int(&mut self, i: i64) {
119        self.buf.push_str(&i.to_string());
120    }
121
122    pub fn push_real(&mut self, f: f64) {
123        self.buf.push_str(&f.to_string());
124    }
125
126    pub fn push_space(&mut self) {
127        self.buf.push(' ');
128    }
129
130    pub fn push_dot(&mut self) {
131        self.buf.push('.');
132    }
133
134    pub fn push_equal_sign(&mut self) {
135        self.buf.push('=');
136    }
137
138    pub fn open_brace(&mut self) {
139        self.buf.push('(');
140    }
141
142    pub fn close_brace(&mut self) {
143        self.buf.push(')');
144    }
145
146    pub fn as_str(&self) -> &str {
147        &self.buf
148    }
149
150    fn wrap_and_escape(&mut self, s: &str, quote: char) {
151        self.buf.push(quote);
152        let chars = s.chars();
153        for ch in chars {
154            // escape `quote` by doubling it
155            if ch == quote {
156                self.buf.push(ch);
157            }
158            self.buf.push(ch);
159        }
160        self.buf.push(quote);
161    }
162}
163
164impl Deref for Sql {
165    type Target = str;
166
167    fn deref(&self) -> &str {
168        self.as_str()
169    }
170}
171
172impl Connection {
173    /// Query the current value of `pragma_name`.
174    pub fn pragma_query_one<T: FromSql>(
175        &self,
176        schema_name: Option<&str>,
177        pragma_name: &str,
178    ) -> Result<T> {
179        let mut query = Sql::new();
180        query.push_pragma(self, schema_name, pragma_name)?;
181        self.one_column(&query, [])
182    }
183    /// Query the current value of `pragma_name`.
184    ///
185    /// Some pragmas will return multiple rows/values which cannot be retrieved
186    /// with this method.
187    ///
188    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
189    /// `SELECT user_version FROM pragma_user_version;`
190    pub fn pragma_query_value<T, F>(
191        &self,
192        schema_name: Option<&str>,
193        pragma_name: &str,
194        f: F,
195    ) -> Result<T>
196    where
197        F: FnOnce(&Row<'_>) -> Result<T>,
198    {
199        let mut query = Sql::new();
200        query.push_pragma(self, schema_name, pragma_name)?;
201        self.query_row(&query, [], f)
202    }
203
204    /// Query the current rows/values of `pragma_name`.
205    ///
206    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
207    /// `SELECT * FROM pragma_collation_list;`
208    pub fn pragma_query<F>(
209        &self,
210        schema_name: Option<&str>,
211        pragma_name: &str,
212        mut f: F,
213    ) -> Result<()>
214    where
215        F: FnMut(&Row<'_>) -> Result<()>,
216    {
217        let mut query = Sql::new();
218        query.push_pragma(self, schema_name, pragma_name)?;
219        let mut stmt = self.prepare(&query)?;
220        let mut rows = stmt.query([])?;
221        while let Some(result_row) = rows.next()? {
222            let row = result_row;
223            f(row)?;
224        }
225        Ok(())
226    }
227
228    /// Query the current value(s) of `pragma_name` associated to
229    /// `pragma_value`.
230    ///
231    /// This method can be used with query-only pragmas which need an argument
232    /// (e.g. `table_info('one_tbl')`) or pragmas which returns value(s)
233    /// (e.g. `integrity_check`).
234    ///
235    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
236    /// `SELECT * FROM pragma_table_info(?1);`
237    pub fn pragma<F, V>(
238        &self,
239        schema_name: Option<&str>,
240        pragma_name: &str,
241        pragma_value: V,
242        mut f: F,
243    ) -> Result<()>
244    where
245        F: FnMut(&Row<'_>) -> Result<()>,
246        V: ToSql,
247    {
248        let mut sql = Sql::new();
249        sql.push_pragma(self, schema_name, pragma_name)?;
250        // The argument may be either in parentheses
251        // or it may be separated from the pragma name by an equal sign.
252        // The two syntaxes yield identical results.
253        sql.open_brace();
254        sql.push_value(&pragma_value)?;
255        sql.close_brace();
256        let mut stmt = self.prepare(&sql)?;
257        let mut rows = stmt.query([])?;
258        while let Some(result_row) = rows.next()? {
259            let row = result_row;
260            f(row)?;
261        }
262        Ok(())
263    }
264
265    /// Set a new value to `pragma_name`.
266    ///
267    /// Some pragmas will return the updated value which cannot be retrieved
268    /// with this method.
269    pub fn pragma_update<V>(
270        &self,
271        schema_name: Option<&str>,
272        pragma_name: &str,
273        pragma_value: V,
274    ) -> Result<()>
275    where
276        V: ToSql,
277    {
278        let mut sql = Sql::new();
279        sql.push_pragma(self, schema_name, pragma_name)?;
280        // The argument may be either in parentheses
281        // or it may be separated from the pragma name by an equal sign.
282        // The two syntaxes yield identical results.
283        sql.push_equal_sign();
284        sql.push_value(&pragma_value)?;
285        self.execute_batch(&sql)
286    }
287
288    /// Set a new value to `pragma_name` and return the updated value.
289    ///
290    /// Only few pragmas automatically return the updated value.
291    pub fn pragma_update_and_check<F, T, V>(
292        &self,
293        schema_name: Option<&str>,
294        pragma_name: &str,
295        pragma_value: V,
296        f: F,
297    ) -> Result<T>
298    where
299        F: FnOnce(&Row<'_>) -> Result<T>,
300        V: ToSql,
301    {
302        let mut sql = Sql::new();
303        sql.push_pragma(self, schema_name, pragma_name)?;
304        // The argument may be either in parentheses
305        // or it may be separated from the pragma name by an equal sign.
306        // The two syntaxes yield identical results.
307        sql.push_equal_sign();
308        sql.push_value(&pragma_value)?;
309        self.query_row(&sql, [], f)
310    }
311}
312
313fn is_identifier(s: &str) -> bool {
314    let chars = s.char_indices();
315    for (i, ch) in chars {
316        if i == 0 {
317            if !is_identifier_start(ch) {
318                return false;
319            }
320        } else if !is_identifier_continue(ch) {
321            return false;
322        }
323    }
324    true
325}
326
327fn is_identifier_start(c: char) -> bool {
328    c.is_ascii_uppercase() || c == '_' || c.is_ascii_lowercase() || c > '\x7F'
329}
330
331fn is_identifier_continue(c: char) -> bool {
332    c == '$'
333        || c.is_ascii_digit()
334        || c.is_ascii_uppercase()
335        || c == '_'
336        || c.is_ascii_lowercase()
337        || c > '\x7F'
338}
339
340fn is_keyword(s: &str) -> bool {
341    unsafe {
342        ffi::sqlite3_keyword_check(
343            s.as_ptr().cast::<std::ffi::c_char>(),
344            s.len().try_into().unwrap(),
345        ) != 0
346    }
347}
348
349#[cfg(test)]
350mod test {
351    #[cfg(all(target_family = "wasm", target_os = "unknown"))]
352    use wasm_bindgen_test::wasm_bindgen_test as test;
353
354    use super::Sql;
355    use crate::{Connection, Result, pragma};
356
357    #[test]
358    #[cfg_attr(miri, ignore)]
359    fn pragma_query_value() -> Result<()> {
360        let db = Connection::open_in_memory()?;
361        let user_version: i32 = db.pragma_query_value(None, "user_version", |row| row.get(0))?;
362        assert_eq!(0, user_version);
363        Ok(())
364    }
365
366    #[test]
367    #[cfg_attr(miri, ignore)]
368    fn pragma_func_query_value() -> Result<()> {
369        let db = Connection::open_in_memory()?;
370        let user_version: i32 =
371            db.one_column("SELECT user_version FROM pragma_user_version", [])?;
372        assert_eq!(0, user_version);
373        Ok(())
374    }
375
376    #[test]
377    #[cfg_attr(miri, ignore)]
378    fn pragma_query_no_schema() -> Result<()> {
379        let db = Connection::open_in_memory()?;
380        let mut user_version = -1;
381        db.pragma_query(None, "user_version", |row| {
382            user_version = row.get(0)?;
383            Ok(())
384        })?;
385        assert_eq!(0, user_version);
386        Ok(())
387    }
388
389    #[test]
390    #[cfg_attr(miri, ignore)]
391    fn pragma_query_with_schema() -> Result<()> {
392        let db = Connection::open_in_memory()?;
393        let mut user_version = -1;
394        db.pragma_query(Some("main"), "user_version", |row| {
395            user_version = row.get(0)?;
396            Ok(())
397        })?;
398        assert_eq!(0, user_version);
399        Ok(())
400    }
401
402    #[test]
403    #[cfg_attr(miri, ignore)]
404    fn pragma() -> Result<()> {
405        let db = Connection::open_in_memory()?;
406        let mut columns = Vec::new();
407        db.pragma(None, "table_info", "sqlite_master", |row| {
408            let column: String = row.get(1)?;
409            columns.push(column);
410            Ok(())
411        })?;
412        assert_eq!(5, columns.len());
413        Ok(())
414    }
415
416    #[test]
417    #[cfg_attr(miri, ignore)]
418    fn pragma_func() -> Result<()> {
419        let db = Connection::open_in_memory()?;
420        let mut table_info = db.prepare("SELECT * FROM pragma_table_info(?1)")?;
421        let mut columns = Vec::new();
422        let mut rows = table_info.query(["sqlite_master"])?;
423
424        while let Some(row) = rows.next()? {
425            let column: String = row.get(1)?;
426            columns.push(column);
427        }
428        assert_eq!(5, columns.len());
429        Ok(())
430    }
431
432    #[test]
433    #[cfg_attr(miri, ignore)]
434    fn pragma_update() -> Result<()> {
435        let db = Connection::open_in_memory()?;
436        db.pragma_update(None, "user_version", 1)
437    }
438
439    #[test]
440    #[cfg_attr(miri, ignore)]
441    fn pragma_update_and_check() -> Result<()> {
442        let db = Connection::open_in_memory()?;
443        let journal_mode: String =
444            db.pragma_update_and_check(None, "journal_mode", "OFF", |row| row.get(0))?;
445        assert!(
446            journal_mode == "off" || journal_mode == "memory",
447            "mode: {journal_mode:?}"
448        );
449        // Sanity checks to ensure the move to a generic `ToSql` wasn't breaking
450        let mode =
451            db.pragma_update_and_check(None, "journal_mode", "OFF", |row| row.get::<_, String>(0))?;
452        assert!(mode == "off" || mode == "memory", "mode: {mode:?}");
453
454        let param: &dyn crate::ToSql = &"OFF";
455        let mode =
456            db.pragma_update_and_check(None, "journal_mode", param, |row| row.get::<_, String>(0))?;
457        assert!(mode == "off" || mode == "memory", "mode: {mode:?}");
458        Ok(())
459    }
460
461    #[test]
462    fn is_identifier() {
463        assert!(pragma::is_identifier("full"));
464        assert!(pragma::is_identifier("r2d2"));
465        assert!(!pragma::is_identifier("sp ce"));
466        assert!(!pragma::is_identifier("semi;colon"));
467    }
468
469    #[test]
470    fn double_quote() {
471        let mut sql = Sql::new();
472        sql.push_schema_name(r#"schema";--"#);
473        assert_eq!(r#""schema"";--""#, sql.as_str());
474    }
475
476    #[test]
477    fn wrap_and_escape() {
478        let mut sql = Sql::new();
479        sql.push_string_literal("value'; --");
480        assert_eq!("'value''; --'", sql.as_str());
481    }
482
483    #[test]
484    #[cfg_attr(miri, ignore)]
485    fn locking_mode() -> Result<()> {
486        let db = Connection::open_in_memory()?;
487        db.pragma_update(None, "locking_mode", "exclusive")?;
488        Ok(())
489    }
490
491    #[test]
492    #[cfg_attr(miri, ignore)]
493    #[cfg(all(
494        not(feature = "bundled-sqlcipher"),
495        not(feature = "bundled-sqlcipher-vendored-openssl"),
496        feature = "extra_check"
497    ))]
498    fn invalid_pragma() -> Result<()> {
499        use crate::{error::error_from_sqlite_code, ffi};
500        let db = Connection::open_in_memory()?;
501        let err = db.pragma_update(None, "boom", "1").unwrap_err();
502        assert_eq!(
503            err,
504            error_from_sqlite_code(
505                ffi::SQLITE_MISUSE,
506                Some("Invalid pragma \"boom\"".to_owned())
507            )
508        );
509        Ok(())
510    }
511}