rusqlite/
pragma.rs

1//! Pragma helpers
2
3use 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 = "functions")]
74            ToSqlOutput::Arg(_) => {
75                return Err(Error::SqliteFailure(
76                    ffi::Error::new(ffi::SQLITE_MISUSE),
77                    Some(format!("Unsupported value \"{value:?}\"")),
78                ));
79            }
80            #[cfg(feature = "array")]
81            ToSqlOutput::Array(_) => {
82                return Err(Error::SqliteFailure(
83                    ffi::Error::new(ffi::SQLITE_MISUSE),
84                    Some(format!("Unsupported value \"{value:?}\"")),
85                ));
86            }
87        };
88        match value {
89            ValueRef::Integer(i) => {
90                self.push_int(i);
91            }
92            ValueRef::Real(r) => {
93                self.push_real(r);
94            }
95            ValueRef::Text(s) => {
96                let s = std::str::from_utf8(s)?;
97                self.push_string_literal(s);
98            }
99            _ => {
100                return Err(Error::SqliteFailure(
101                    ffi::Error::new(ffi::SQLITE_MISUSE),
102                    Some(format!("Unsupported value \"{value:?}\"")),
103                ));
104            }
105        };
106        Ok(())
107    }
108
109    pub fn push_string_literal(&mut self, s: &str) {
110        self.wrap_and_escape(s, '\'');
111    }
112
113    pub fn push_int(&mut self, i: i64) {
114        self.buf.push_str(&i.to_string());
115    }
116
117    pub fn push_real(&mut self, f: f64) {
118        self.buf.push_str(&f.to_string());
119    }
120
121    pub fn push_space(&mut self) {
122        self.buf.push(' ');
123    }
124
125    pub fn push_dot(&mut self) {
126        self.buf.push('.');
127    }
128
129    pub fn push_equal_sign(&mut self) {
130        self.buf.push('=');
131    }
132
133    pub fn open_brace(&mut self) {
134        self.buf.push('(');
135    }
136
137    pub fn close_brace(&mut self) {
138        self.buf.push(')');
139    }
140
141    pub fn as_str(&self) -> &str {
142        &self.buf
143    }
144
145    fn wrap_and_escape(&mut self, s: &str, quote: char) {
146        self.buf.push(quote);
147        let chars = s.chars();
148        for ch in chars {
149            // escape `quote` by doubling it
150            if ch == quote {
151                self.buf.push(ch);
152            }
153            self.buf.push(ch);
154        }
155        self.buf.push(quote);
156    }
157}
158
159impl Deref for Sql {
160    type Target = str;
161
162    fn deref(&self) -> &str {
163        self.as_str()
164    }
165}
166
167impl Connection {
168    /// Query the current value of `pragma_name`.
169    ///
170    /// Some pragmas will return multiple rows/values which cannot be retrieved
171    /// with this method.
172    ///
173    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
174    /// `SELECT user_version FROM pragma_user_version;`
175    pub fn pragma_query_value<T, F>(
176        &self,
177        schema_name: Option<DatabaseName<'_>>,
178        pragma_name: &str,
179        f: F,
180    ) -> Result<T>
181    where
182        F: FnOnce(&Row<'_>) -> Result<T>,
183    {
184        let mut query = Sql::new();
185        query.push_pragma(schema_name, pragma_name)?;
186        self.query_row(&query, [], f)
187    }
188
189    /// Query the current rows/values of `pragma_name`.
190    ///
191    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
192    /// `SELECT * FROM pragma_collation_list;`
193    pub fn pragma_query<F>(
194        &self,
195        schema_name: Option<DatabaseName<'_>>,
196        pragma_name: &str,
197        mut f: F,
198    ) -> Result<()>
199    where
200        F: FnMut(&Row<'_>) -> Result<()>,
201    {
202        let mut query = Sql::new();
203        query.push_pragma(schema_name, pragma_name)?;
204        let mut stmt = self.prepare(&query)?;
205        let mut rows = stmt.query([])?;
206        while let Some(result_row) = rows.next()? {
207            let row = result_row;
208            f(row)?;
209        }
210        Ok(())
211    }
212
213    /// Query the current value(s) of `pragma_name` associated to
214    /// `pragma_value`.
215    ///
216    /// This method can be used with query-only pragmas which need an argument
217    /// (e.g. `table_info('one_tbl')`) or pragmas which returns value(s)
218    /// (e.g. `integrity_check`).
219    ///
220    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
221    /// `SELECT * FROM pragma_table_info(?1);`
222    pub fn pragma<F, V>(
223        &self,
224        schema_name: Option<DatabaseName<'_>>,
225        pragma_name: &str,
226        pragma_value: V,
227        mut f: F,
228    ) -> Result<()>
229    where
230        F: FnMut(&Row<'_>) -> Result<()>,
231        V: ToSql,
232    {
233        let mut sql = Sql::new();
234        sql.push_pragma(schema_name, pragma_name)?;
235        // The argument may be either in parentheses
236        // or it may be separated from the pragma name by an equal sign.
237        // The two syntaxes yield identical results.
238        sql.open_brace();
239        sql.push_value(&pragma_value)?;
240        sql.close_brace();
241        let mut stmt = self.prepare(&sql)?;
242        let mut rows = stmt.query([])?;
243        while let Some(result_row) = rows.next()? {
244            let row = result_row;
245            f(row)?;
246        }
247        Ok(())
248    }
249
250    /// Set a new value to `pragma_name`.
251    ///
252    /// Some pragmas will return the updated value which cannot be retrieved
253    /// with this method.
254    pub fn pragma_update<V>(
255        &self,
256        schema_name: Option<DatabaseName<'_>>,
257        pragma_name: &str,
258        pragma_value: V,
259    ) -> Result<()>
260    where
261        V: ToSql,
262    {
263        let mut sql = Sql::new();
264        sql.push_pragma(schema_name, pragma_name)?;
265        // The argument may be either in parentheses
266        // or it may be separated from the pragma name by an equal sign.
267        // The two syntaxes yield identical results.
268        sql.push_equal_sign();
269        sql.push_value(&pragma_value)?;
270        self.execute_batch(&sql)
271    }
272
273    /// Set a new value to `pragma_name` and return the updated value.
274    ///
275    /// Only few pragmas automatically return the updated value.
276    pub fn pragma_update_and_check<F, T, V>(
277        &self,
278        schema_name: Option<DatabaseName<'_>>,
279        pragma_name: &str,
280        pragma_value: V,
281        f: F,
282    ) -> Result<T>
283    where
284        F: FnOnce(&Row<'_>) -> Result<T>,
285        V: ToSql,
286    {
287        let mut sql = Sql::new();
288        sql.push_pragma(schema_name, pragma_name)?;
289        // The argument may be either in parentheses
290        // or it may be separated from the pragma name by an equal sign.
291        // The two syntaxes yield identical results.
292        sql.push_equal_sign();
293        sql.push_value(&pragma_value)?;
294        self.query_row(&sql, [], f)
295    }
296}
297
298fn is_identifier(s: &str) -> bool {
299    let chars = s.char_indices();
300    for (i, ch) in chars {
301        if i == 0 {
302            if !is_identifier_start(ch) {
303                return false;
304            }
305        } else if !is_identifier_continue(ch) {
306            return false;
307        }
308    }
309    true
310}
311
312fn is_identifier_start(c: char) -> bool {
313    c.is_ascii_uppercase() || c == '_' || c.is_ascii_lowercase() || c > '\x7F'
314}
315
316fn is_identifier_continue(c: char) -> bool {
317    c == '$'
318        || c.is_ascii_digit()
319        || c.is_ascii_uppercase()
320        || c == '_'
321        || c.is_ascii_lowercase()
322        || c > '\x7F'
323}
324
325#[cfg(test)]
326mod test {
327    use super::Sql;
328    use crate::pragma;
329    use crate::{Connection, DatabaseName, Result};
330
331    #[test]
332    fn pragma_query_value() -> Result<()> {
333        let db = Connection::open_in_memory()?;
334        let user_version: i32 = db.pragma_query_value(None, "user_version", |row| row.get(0))?;
335        assert_eq!(0, user_version);
336        Ok(())
337    }
338
339    #[test]
340    #[cfg(feature = "modern_sqlite")]
341    fn pragma_func_query_value() -> Result<()> {
342        let db = Connection::open_in_memory()?;
343        let user_version: i32 = db.one_column("SELECT user_version FROM pragma_user_version")?;
344        assert_eq!(0, user_version);
345        Ok(())
346    }
347
348    #[test]
349    fn pragma_query_no_schema() -> Result<()> {
350        let db = Connection::open_in_memory()?;
351        let mut user_version = -1;
352        db.pragma_query(None, "user_version", |row| {
353            user_version = row.get(0)?;
354            Ok(())
355        })?;
356        assert_eq!(0, user_version);
357        Ok(())
358    }
359
360    #[test]
361    fn pragma_query_with_schema() -> Result<()> {
362        let db = Connection::open_in_memory()?;
363        let mut user_version = -1;
364        db.pragma_query(Some(DatabaseName::Main), "user_version", |row| {
365            user_version = row.get(0)?;
366            Ok(())
367        })?;
368        assert_eq!(0, user_version);
369        Ok(())
370    }
371
372    #[test]
373    fn pragma() -> Result<()> {
374        let db = Connection::open_in_memory()?;
375        let mut columns = Vec::new();
376        db.pragma(None, "table_info", "sqlite_master", |row| {
377            let column: String = row.get(1)?;
378            columns.push(column);
379            Ok(())
380        })?;
381        assert_eq!(5, columns.len());
382        Ok(())
383    }
384
385    #[test]
386    #[cfg(feature = "modern_sqlite")]
387    fn pragma_func() -> Result<()> {
388        let db = Connection::open_in_memory()?;
389        let mut table_info = db.prepare("SELECT * FROM pragma_table_info(?1)")?;
390        let mut columns = Vec::new();
391        let mut rows = table_info.query(["sqlite_master"])?;
392
393        while let Some(row) = rows.next()? {
394            let column: String = row.get(1)?;
395            columns.push(column);
396        }
397        assert_eq!(5, columns.len());
398        Ok(())
399    }
400
401    #[test]
402    fn pragma_update() -> Result<()> {
403        let db = Connection::open_in_memory()?;
404        db.pragma_update(None, "user_version", 1)
405    }
406
407    #[test]
408    fn pragma_update_and_check() -> Result<()> {
409        let db = Connection::open_in_memory()?;
410        let journal_mode: String =
411            db.pragma_update_and_check(None, "journal_mode", "OFF", |row| row.get(0))?;
412        assert!(
413            journal_mode == "off" || journal_mode == "memory",
414            "mode: {journal_mode:?}"
415        );
416        // Sanity checks to ensure the move to a generic `ToSql` wasn't breaking
417        let mode =
418            db.pragma_update_and_check(None, "journal_mode", "OFF", |row| row.get::<_, String>(0))?;
419        assert!(mode == "off" || mode == "memory", "mode: {mode:?}");
420
421        let param: &dyn crate::ToSql = &"OFF";
422        let mode =
423            db.pragma_update_and_check(None, "journal_mode", param, |row| row.get::<_, String>(0))?;
424        assert!(mode == "off" || mode == "memory", "mode: {mode:?}");
425        Ok(())
426    }
427
428    #[test]
429    fn is_identifier() {
430        assert!(pragma::is_identifier("full"));
431        assert!(pragma::is_identifier("r2d2"));
432        assert!(!pragma::is_identifier("sp ce"));
433        assert!(!pragma::is_identifier("semi;colon"));
434    }
435
436    #[test]
437    fn double_quote() {
438        let mut sql = Sql::new();
439        sql.push_schema_name(DatabaseName::Attached(r#"schema";--"#));
440        assert_eq!(r#""schema"";--""#, sql.as_str());
441    }
442
443    #[test]
444    fn wrap_and_escape() {
445        let mut sql = Sql::new();
446        sql.push_string_literal("value'; --");
447        assert_eq!("'value''; --'", sql.as_str());
448    }
449
450    #[test]
451    fn locking_mode() -> Result<()> {
452        let db = Connection::open_in_memory()?;
453        let r = db.pragma_update(None, "locking_mode", "exclusive");
454        if cfg!(feature = "extra_check") {
455            r.unwrap_err();
456        } else {
457            r?;
458        }
459        Ok(())
460    }
461}
rusqlite/pragma.rs

rusqlite/
pragma.rs
