From f82ea6567389b2382d492985ae7c60b302fef140 Mon Sep 17 00:00:00 2001
From: George <g.andronchik@gmail.com>
Date: Wed, 18 May 2022 18:16:52 +0800
Subject: [PATCH] feat: support DISCARD [ALL | PLANS | SEQUENCES | TEMPORARY |
 TEMP]

---
 src/ast/mod.rs            | 29 +++++++++++++++++++++++++++++
 src/keywords.rs           |  2 ++
 src/parser.rs             | 19 +++++++++++++++++++
 tests/sqlparser_common.rs | 27 +++++++++++++++++++++++++++
 4 files changed, 77 insertions(+)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index d9052916d..91df33f87 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -911,6 +911,11 @@ pub enum Statement {
         /// deleted along with the dropped table
         purge: bool,
     },
+    /// DISCARD [ ALL | PLANS | SEQUENCES | TEMPORARY | TEMP ]
+    ///
+    /// Note: this is a PostgreSQL-specific statement,
+    /// but may also compatible with other SQL.
+    Discard { object_type: DiscardObject },
     /// SET [ SESSION | LOCAL ] ROLE role_name
     ///
     /// Note: this is a PostgreSQL-specific statement,
@@ -1575,6 +1580,10 @@ impl fmt::Display for Statement {
                 if *cascade { " CASCADE" } else { "" },
                 if *purge { " PURGE" } else { "" }
             ),
+            Statement::Discard { object_type } => {
+                write!(f, "DISCARD {object_type}", object_type = object_type)?;
+                Ok(())
+            }
             Statement::SetRole {
                 local,
                 session,
@@ -2547,6 +2556,26 @@ impl fmt::Display for MergeClause {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum DiscardObject {
+    ALL,
+    PLANS,
+    SEQUENCES,
+    TEMP,
+}
+
+impl fmt::Display for DiscardObject {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            DiscardObject::ALL => f.write_str("ALL"),
+            DiscardObject::PLANS => f.write_str("PLANS"),
+            DiscardObject::SEQUENCES => f.write_str("SEQUENCES"),
+            DiscardObject::TEMP => f.write_str("TEMP"),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/keywords.rs b/src/keywords.rs
index 709b60bf8..b4ddb68dd 100644
--- a/src/keywords.rs
+++ b/src/keywords.rs
@@ -184,6 +184,7 @@ define_keywords!(
     DESCRIBE,
     DETERMINISTIC,
     DIRECTORY,
+    DISCARD,
     DISCONNECT,
     DISTINCT,
     DISTRIBUTE,
@@ -371,6 +372,7 @@ define_keywords!(
     PERCENTILE_DISC,
     PERCENT_RANK,
     PERIOD,
+    PLANS,
     PORTION,
     POSITION,
     POSITION_REGEX,
diff --git a/src/parser.rs b/src/parser.rs
index 5ee3d5cb6..c24877048 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -166,6 +166,7 @@ impl<'a> Parser<'a> {
                 Keyword::MSCK => Ok(self.parse_msck()?),
                 Keyword::CREATE => Ok(self.parse_create()?),
                 Keyword::DROP => Ok(self.parse_drop()?),
+                Keyword::DISCARD => Ok(self.parse_discard()?),
                 Keyword::DELETE => Ok(self.parse_delete()?),
                 Keyword::INSERT => Ok(self.parse_insert()?),
                 Keyword::UPDATE => Ok(self.parse_update()?),
@@ -1786,6 +1787,24 @@ impl<'a> Parser<'a> {
         })
     }
 
+    pub fn parse_discard(&mut self) -> Result<Statement, ParserError> {
+        let object_type = if self.parse_keyword(Keyword::ALL) {
+            DiscardObject::ALL
+        } else if self.parse_keyword(Keyword::PLANS) {
+            DiscardObject::PLANS
+        } else if self.parse_keyword(Keyword::SEQUENCES) {
+            DiscardObject::SEQUENCES
+        } else if self.parse_keyword(Keyword::TEMP) || self.parse_keyword(Keyword::TEMPORARY) {
+            DiscardObject::TEMP
+        } else {
+            return self.expected(
+                "ALL, PLANS, SEQUENCES, TEMP or TEMPORARY after DISCARD",
+                self.peek_token(),
+            );
+        };
+        Ok(Statement::Discard { object_type })
+    }
+
     pub fn parse_create_index(&mut self, unique: bool) -> Result<Statement, ParserError> {
         let if_not_exists = self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
         let index_name = self.parse_object_name()?;
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index 0986e407e..c80a3a0c8 100644
--- a/tests/sqlparser_common.rs
+++ b/tests/sqlparser_common.rs
@@ -4740,3 +4740,30 @@ fn parse_is_boolean() {
         res.unwrap_err()
     );
 }
+
+#[test]
+fn parse_discard() {
+    let sql = "DISCARD ALL";
+    match verified_stmt(sql) {
+        Statement::Discard { object_type, .. } => assert_eq!(object_type, DiscardObject::ALL),
+        _ => unreachable!(),
+    }
+
+    let sql = "DISCARD PLANS";
+    match verified_stmt(sql) {
+        Statement::Discard { object_type, .. } => assert_eq!(object_type, DiscardObject::PLANS),
+        _ => unreachable!(),
+    }
+
+    let sql = "DISCARD SEQUENCES";
+    match verified_stmt(sql) {
+        Statement::Discard { object_type, .. } => assert_eq!(object_type, DiscardObject::SEQUENCES),
+        _ => unreachable!(),
+    }
+
+    let sql = "DISCARD TEMP";
+    match verified_stmt(sql) {
+        Statement::Discard { object_type, .. } => assert_eq!(object_type, DiscardObject::TEMP),
+        _ => unreachable!(),
+    }
+}