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