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 = "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            // escape `quote` by doubling it
143            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    /// Query the current value of `pragma_name`.
162    ///
163    /// Some pragmas will return multiple rows/values which cannot be retrieved
164    /// with this method.
165    ///
166    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
167    /// `SELECT user_version FROM pragma_user_version;`
168    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    /// Query the current rows/values of `pragma_name`.
183    ///
184    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
185    /// `SELECT * FROM pragma_collation_list;`
186    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    /// Query the current value(s) of `pragma_name` associated to
207    /// `pragma_value`.
208    ///
209    /// This method can be used with query-only pragmas which need an argument
210    /// (e.g. `table_info('one_tbl')`) or pragmas which returns value(s)
211    /// (e.g. `integrity_check`).
212    ///
213    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
214    /// `SELECT * FROM pragma_table_info(?);`
215    pub fn pragma<F, V>(
216        &self,
217        schema_name: Option<DatabaseName<'_>>,
218        pragma_name: &str,
219        pragma_value: V,
220        mut f: F,
221    ) -> Result<()>
222    where
223        F: FnMut(&Row<'_>) -> Result<()>,
224        V: ToSql,
225    {
226        let mut sql = Sql::new();
227        sql.push_pragma(schema_name, pragma_name)?;
228        // The argument may be either in parentheses
229        // or it may be separated from the pragma name by an equal sign.
230        // The two syntaxes yield identical results.
231        sql.open_brace();
232        sql.push_value(&pragma_value)?;
233        sql.close_brace();
234        let mut stmt = self.prepare(&sql)?;
235        let mut rows = stmt.query([])?;
236        while let Some(result_row) = rows.next()? {
237            let row = result_row;
238            f(row)?;
239        }
240        Ok(())
241    }
242
243    /// Set a new value to `pragma_name`.
244    ///
245    /// Some pragmas will return the updated value which cannot be retrieved
246    /// with this method.
247    pub fn pragma_update<V>(
248        &self,
249        schema_name: Option<DatabaseName<'_>>,
250        pragma_name: &str,
251        pragma_value: V,
252    ) -> Result<()>
253    where
254        V: ToSql,
255    {
256        let mut sql = Sql::new();
257        sql.push_pragma(schema_name, pragma_name)?;
258        // The argument may be either in parentheses
259        // or it may be separated from the pragma name by an equal sign.
260        // The two syntaxes yield identical results.
261        sql.push_equal_sign();
262        sql.push_value(&pragma_value)?;
263        self.execute_batch(&sql)
264    }
265
266    /// Set a new value to `pragma_name` and return the updated value.
267    ///
268    /// Only few pragmas automatically return the updated value.
269    pub fn pragma_update_and_check<F, T, V>(
270        &self,
271        schema_name: Option<DatabaseName<'_>>,
272        pragma_name: &str,
273        pragma_value: V,
274        f: F,
275    ) -> Result<T>
276    where
277        F: FnOnce(&Row<'_>) -> Result<T>,
278        V: ToSql,
279    {
280        let mut sql = Sql::new();
281        sql.push_pragma(schema_name, pragma_name)?;
282        // The argument may be either in parentheses
283        // or it may be separated from the pragma name by an equal sign.
284        // The two syntaxes yield identical results.
285        sql.push_equal_sign();
286        sql.push_value(&pragma_value)?;
287        self.query_row(&sql, [], f)
288    }
289}
290
291fn is_identifier(s: &str) -> bool {
292    let chars = s.char_indices();
293    for (i, ch) in chars {
294        if i == 0 {
295            if !is_identifier_start(ch) {
296                return false;
297            }
298        } else if !is_identifier_continue(ch) {
299            return false;
300        }
301    }
302    true
303}
304
305fn is_identifier_start(c: char) -> bool {
306    ('A'..='Z').contains(&c) || c == '_' || ('a'..='z').contains(&c) || c > '\x7F'
307}
308
309fn is_identifier_continue(c: char) -> bool {
310    c == '$'
311        || ('0'..='9').contains(&c)
312        || ('A'..='Z').contains(&c)
313        || c == '_'
314        || ('a'..='z').contains(&c)
315        || c > '\x7F'
316}
317
318#[cfg(test)]
319mod test {
320    use super::Sql;
321    use crate::pragma;
322    use crate::{Connection, DatabaseName, Result};
323
324    #[test]
325    fn pragma_query_value() -> Result<()> {
326        let db = Connection::open_in_memory()?;
327        let user_version: i32 = db.pragma_query_value(None, "user_version", |row| row.get(0))?;
328        assert_eq!(0, user_version);
329        Ok(())
330    }
331
332    #[test]
333    #[cfg(feature = "modern_sqlite")]
334    fn pragma_func_query_value() -> Result<()> {
335        let db = Connection::open_in_memory()?;
336        let user_version: i32 =
337            db.query_row("SELECT user_version FROM pragma_user_version", [], |row| {
338                row.get(0)
339            })?;
340        assert_eq!(0, user_version);
341        Ok(())
342    }
343
344    #[test]
345    fn pragma_query_no_schema() -> Result<()> {
346        let db = Connection::open_in_memory()?;
347        let mut user_version = -1;
348        db.pragma_query(None, "user_version", |row| {
349            user_version = row.get(0)?;
350            Ok(())
351        })?;
352        assert_eq!(0, user_version);
353        Ok(())
354    }
355
356    #[test]
357    fn pragma_query_with_schema() -> Result<()> {
358        let db = Connection::open_in_memory()?;
359        let mut user_version = -1;
360        db.pragma_query(Some(DatabaseName::Main), "user_version", |row| {
361            user_version = row.get(0)?;
362            Ok(())
363        })?;
364        assert_eq!(0, user_version);
365        Ok(())
366    }
367
368    #[test]
369    fn pragma() -> Result<()> {
370        let db = Connection::open_in_memory()?;
371        let mut columns = Vec::new();
372        db.pragma(None, "table_info", &"sqlite_master", |row| {
373            let column: String = row.get(1)?;
374            columns.push(column);
375            Ok(())
376        })?;
377        assert_eq!(5, columns.len());
378        Ok(())
379    }
380
381    #[test]
382    #[cfg(feature = "modern_sqlite")]
383    fn pragma_func() -> Result<()> {
384        let db = Connection::open_in_memory()?;
385        let mut table_info = db.prepare("SELECT * FROM pragma_table_info(?)")?;
386        let mut columns = Vec::new();
387        let mut rows = table_info.query(["sqlite_master"])?;
388
389        while let Some(row) = rows.next()? {
390            let row = row;
391            let column: String = row.get(1)?;
392            columns.push(column);
393        }
394        assert_eq!(5, columns.len());
395        Ok(())
396    }
397
398    #[test]
399    fn pragma_update() -> Result<()> {
400        let db = Connection::open_in_memory()?;
401        db.pragma_update(None, "user_version", 1)
402    }
403
404    #[test]
405    fn pragma_update_and_check() -> Result<()> {
406        let db = Connection::open_in_memory()?;
407        let journal_mode: String =
408            db.pragma_update_and_check(None, "journal_mode", "OFF", |row| row.get(0))?;
409        assert!(
410            journal_mode == "off" || journal_mode == "memory",
411            "mode: {:?}",
412            journal_mode,
413        );
414        // Sanity checks to ensure the move to a generic `ToSql` wasn't breaking
415        let mode = db
416            .pragma_update_and_check(None, "journal_mode", &"OFF", |row| row.get::<_, String>(0))?;
417        assert!(mode == "off" || mode == "memory", "mode: {:?}", mode);
418
419        let param: &dyn crate::ToSql = &"OFF";
420        let mode =
421            db.pragma_update_and_check(None, "journal_mode", param, |row| row.get::<_, String>(0))?;
422        assert!(mode == "off" || mode == "memory", "mode: {:?}", mode);
423        Ok(())
424    }
425
426    #[test]
427    fn is_identifier() {
428        assert!(pragma::is_identifier("full"));
429        assert!(pragma::is_identifier("r2d2"));
430        assert!(!pragma::is_identifier("sp ce"));
431        assert!(!pragma::is_identifier("semi;colon"));
432    }
433
434    #[test]
435    fn double_quote() {
436        let mut sql = Sql::new();
437        sql.push_schema_name(DatabaseName::Attached(r#"schema";--"#));
438        assert_eq!(r#""schema"";--""#, sql.as_str());
439    }
440
441    #[test]
442    fn wrap_and_escape() {
443        let mut sql = Sql::new();
444        sql.push_string_literal("value'; --");
445        assert_eq!("'value''; --'", sql.as_str());
446    }
447
448    #[test]
449    fn locking_mode() -> Result<()> {
450        let db = Connection::open_in_memory()?;
451        let r = db.pragma_update(None, "locking_mode", &"exclusive");
452        if cfg!(feature = "extra_check") {
453            r.unwrap_err();
454        } else {
455            r?;
456        }
457        Ok(())
458    }
459}
rusqlite/pragma.rs

rusqlite/
pragma.rs
