parse SQLite pragma statement (apache#969)

Co-authored-by: Andrew Lamb <[email protected]>

Author: marhoily

src/ast/mod.rs

@@ -1914,6 +1914,12 @@ pub enum Statement {
         name: ObjectName,
         representation: UserDefinedTypeRepresentation,
     },
+    // PRAGMA <schema-name>.<pragma-name> = <pragma-value>
+    Pragma {
+        name: ObjectName,
+        value: Option<Value>,
+        is_eq: bool,
+    },
 }
 
 impl fmt::Display for Statement {
@@ -3276,6 +3282,18 @@ impl fmt::Display for Statement {
             } => {
                 write!(f, "CREATE TYPE {name} AS {representation}")
             }
+            Statement::Pragma { name, value, is_eq } => {
+                write!(f, "PRAGMA {name}")?;
+                if value.is_some() {
+                    let val = value.as_ref().unwrap();
+                    if *is_eq {
+                        write!(f, " = {val}")?;
+                    } else {
+                        write!(f, "({val})")?;
+                    }
+                }
+                Ok(())
+            }
         }
     }
 }

src/keywords.rs

@@ -472,6 +472,7 @@ define_keywords!(
     POSITION,
     POSITION_REGEX,
     POWER,
+    PRAGMA,
     PRECEDES,
     PRECEDING,
     PRECISION,

src/parser/mod.rs

@@ -491,6 +491,8 @@ impl<'a> Parser<'a> {
                 Keyword::EXECUTE => Ok(self.parse_execute()?),
                 Keyword::PREPARE => Ok(self.parse_prepare()?),
                 Keyword::MERGE => Ok(self.parse_merge()?),
+                // `PRAGMA` is sqlite specific https://www.sqlite.org/pragma.html
+                Keyword::PRAGMA => Ok(self.parse_pragma()?),
                 _ => self.expected("an SQL statement", next_token),
             },
             Token::LParen => {
@@ -7502,6 +7504,32 @@ impl<'a> Parser<'a> {
         })
     }
 
+    // PRAGMA [schema-name '.'] pragma-name [('=' pragma-value) | '(' pragma-value ')']
+    pub fn parse_pragma(&mut self) -> Result<Statement, ParserError> {
+        let name = self.parse_object_name()?;
+        if self.consume_token(&Token::LParen) {
+            let value = self.parse_number_value()?;
+            self.expect_token(&Token::RParen)?;
+            Ok(Statement::Pragma {
+                name,
+                value: Some(value),
+                is_eq: false,
+            })
+        } else if self.consume_token(&Token::Eq) {
+            Ok(Statement::Pragma {
+                name,
+                value: Some(self.parse_number_value()?),
+                is_eq: true,
+            })
+        } else {
+            Ok(Statement::Pragma {
+                name,
+                value: None,
+                is_eq: false,
+            })
+        }
+    }
+
     /// ```sql
     /// CREATE [ { TEMPORARY | TEMP } ] SEQUENCE [ IF NOT EXISTS ] <sequence_name>
     /// ```

tests/sqlparser_sqlite.rs

@@ -24,6 +24,51 @@ use sqlparser::ast::*;
 use sqlparser::dialect::{GenericDialect, SQLiteDialect};
 use sqlparser::tokenizer::Token;
 
+#[test]
+fn pragma_no_value() {
+    let sql = "PRAGMA cache_size";
+    match sqlite_and_generic().verified_stmt(sql) {
+        Statement::Pragma {
+            name,
+            value: None,
+            is_eq: false,
+        } => {
+            assert_eq!("cache_size", name.to_string());
+        }
+        _ => unreachable!(),
+    }
+}
+#[test]
+fn pragma_eq_style() {
+    let sql = "PRAGMA cache_size = 10";
+    match sqlite_and_generic().verified_stmt(sql) {
+        Statement::Pragma {
+            name,
+            value: Some(val),
+            is_eq: true,
+        } => {
+            assert_eq!("cache_size", name.to_string());
+            assert_eq!("10", val.to_string());
+        }
+        _ => unreachable!(),
+    }
+}
+#[test]
+fn pragma_funciton_style() {
+    let sql = "PRAGMA cache_size(10)";
+    match sqlite_and_generic().verified_stmt(sql) {
+        Statement::Pragma {
+            name,
+            value: Some(val),
+            is_eq: false,
+        } => {
+            assert_eq!("cache_size", name.to_string());
+            assert_eq!("10", val.to_string());
+        }
+        _ => unreachable!(),
+    }
+}
+
 #[test]
 fn parse_create_table_without_rowid() {
     let sql = "CREATE TABLE t (a INT) WITHOUT ROWID";