From 9b3fc1330bcab3c6652561a4e5ba84275ab49096 Mon Sep 17 00:00:00 2001
From: Ilya <marhoily@users.noreply.github.com>
Date: Tue, 24 Oct 2023 01:07:39 +0300
Subject: [PATCH] parse SQLite pragma statement (#969)

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
---
 src/ast/mod.rs            | 18 ++++++++++++++++
 src/keywords.rs           |  1 +
 src/parser/mod.rs         | 28 ++++++++++++++++++++++++
 tests/sqlparser_sqlite.rs | 45 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 92 insertions(+)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index d323db8a0..81ee5a343 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -2037,6 +2037,12 @@ pub enum Statement {
         name: ObjectName,
         representation: UserDefinedTypeRepresentation,
     },
+    // PRAGMA <schema-name>.<pragma-name> = <pragma-value>
+    Pragma {
+        name: ObjectName,
+        value: Option<Value>,
+        is_eq: bool,
+    },
 
     //////////////////////////////////
     /// PeerDB Specific Statements ///
@@ -3464,6 +3470,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(())
+            }
 
             //////////////////////////////////////////
             // PeerDB Specific Statements
diff --git a/src/keywords.rs b/src/keywords.rs
index 76c94801c..6002354ec 100644
--- a/src/keywords.rs
+++ b/src/keywords.rs
@@ -481,6 +481,7 @@ define_keywords!(
     POSITION_REGEX,
     POSTGRES,
     POWER,
+    PRAGMA,
     PRECEDES,
     PRECEDING,
     PRECISION,
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index c257c12e5..52bf4e3cc 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -495,6 +495,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 => {
@@ -7550,6 +7552,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>
     /// ```
diff --git a/tests/sqlparser_sqlite.rs b/tests/sqlparser_sqlite.rs
index 8d7ccf315..2fdd4e3de 100644
--- a/tests/sqlparser_sqlite.rs
+++ b/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";