Support SHOW <var> and SET <var>

Author: benesch

src/ast/mod.rs

@@ -440,6 +440,20 @@ pub enum Statement {
         /// `RESTRICT` or no drop behavior at all was specified.
         cascade: bool,
     },
+    /// SET <variable>
+    ///
+    /// Note: this is not a standard SQL statement, but it is supported by at
+    /// least MySQL and PostgreSQL. Not all MySQL-specific syntatic forms are
+    /// supported yet.
+    SetVariable {
+        local: bool,
+        variable: Ident,
+        value: SetVariableValue,
+    },
+    /// SHOW <variable>
+    ///
+    /// Note: this is a PostgreSQL-specific statement.
+    ShowVariable { variable: Ident },
     /// SHOW COLUMNS
     ///
     /// Note: this is a MySQL-specific statement.
@@ -601,6 +615,18 @@ impl fmt::Display for Statement {
                 display_comma_separated(names),
                 if *cascade { " CASCADE" } else { "" },
             ),
+            Statement::SetVariable {
+                local,
+                variable,
+                value,
+            } => {
+                f.write_str("SET ")?;
+                if *local {
+                    f.write_str("LOCAL ")?;
+                }
+                write!(f, "{} = {}", variable, value)
+            }
+            Statement::ShowVariable { variable } => write!(f, "SHOW {}", variable),
             Statement::ShowColumns {
                 extended,
                 full,
@@ -827,3 +853,19 @@ impl fmt::Display for ShowStatementFilter {
         }
     }
 }
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum SetVariableValue {
+    Ident(Ident),
+    Literal(Value),
+}
+
+impl fmt::Display for SetVariableValue {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        use SetVariableValue::*;
+        match self {
+            Ident(ident) => f.write_str(ident),
+            Literal(literal) => write!(f, "{}", literal),
+        }
+    }
+}

src/dialect/keywords.rs

@@ -335,6 +335,7 @@ define_keywords!(
     SELECT,
     SENSITIVE,
     SERIALIZABLE,
+    SESSION,
     SESSION_USER,
     SET,
     SHOW,

src/parser.rs

@@ -125,9 +125,9 @@ impl Parser {
                     "UPDATE" => Ok(self.parse_update()?),
                     "ALTER" => Ok(self.parse_alter()?),
                     "COPY" => Ok(self.parse_copy()?),
+                    "SET" => Ok(self.parse_set()?),
                     "SHOW" => Ok(self.parse_show()?),
                     "START" => Ok(self.parse_start_transaction()?),
-                    "SET" => Ok(self.parse_set_transaction()?),
                     // `BEGIN` is a nonstandard but common alias for the
                     // standard `START TRANSACTION` statement. It is supported
                     // by at least PostgreSQL and MySQL.
@@ -1593,6 +1593,30 @@ impl Parser {
         })
     }
 
+    pub fn parse_set(&mut self) -> Result<Statement, ParserError> {
+        let modifier = self.parse_one_of_keywords(&["SESSION", "LOCAL"]);
+        let variable = self.parse_identifier()?;
+        if self.consume_token(&Token::Eq) || self.parse_keyword("TO") {
+            let token = self.peek_token();
+            let value = match (self.parse_value(), token) {
+                (Ok(value), _) => SetVariableValue::Literal(value),
+                (Err(_), Some(Token::Word(ident))) => SetVariableValue::Ident(ident.as_ident()),
+                (Err(_), other) => self.expected("variable value", other)?,
+            };
+            Ok(Statement::SetVariable {
+                local: modifier == Some("LOCAL"),
+                variable,
+                value,
+            })
+        } else if variable == "TRANSACTION" && modifier.is_none() {
+            Ok(Statement::SetTransaction {
+                modes: self.parse_transaction_modes()?,
+            })
+        } else {
+            self.expected("equals sign or TO", self.peek_token())
+        }
+    }
+
     pub fn parse_show(&mut self) -> Result<Statement, ParserError> {
         if self
             .parse_one_of_keywords(&["EXTENDED", "FULL", "COLUMNS", "FIELDS"])
@@ -1601,7 +1625,9 @@ impl Parser {
             self.prev_token();
             self.parse_show_columns()
         } else {
-            self.expected("EXTENDED, FULL, COLUMNS, or FIELDS", self.peek_token())
+            Ok(Statement::ShowVariable {
+                variable: self.parse_identifier()?,
+            })
         }
     }
 
@@ -1973,13 +1999,6 @@ impl Parser {
         })
     }
 
-    pub fn parse_set_transaction(&mut self) -> Result<Statement, ParserError> {
-        self.expect_keyword("TRANSACTION")?;
-        Ok(Statement::SetTransaction {
-            modes: self.parse_transaction_modes()?,
-        })
-    }
-
     pub fn parse_transaction_modes(&mut self) -> Result<Vec<TransactionMode>, ParserError> {
         let mut modes = vec![];
         let mut required = false;

tests/sqlparser_postgres.rs

@@ -16,6 +16,7 @@
 
 use sqlparser::ast::*;
 use sqlparser::dialect::{GenericDialect, PostgreSqlDialect};
+use sqlparser::parser::ParserError;
 use sqlparser::test_utils::*;
 
 #[test]
@@ -251,6 +252,102 @@ PHP	₱ USD $
     //assert_eq!(sql, ast.to_string());
 }
 
+#[test]
+fn parse_set() {
+    let stmt = pg_and_generic().verified_stmt("SET a = b");
+    assert_eq!(
+        stmt,
+        Statement::SetVariable {
+            local: false,
+            variable: "a".into(),
+            value: SetVariableValue::Ident("b".into()),
+        }
+    );
+
+    let stmt = pg_and_generic().verified_stmt("SET a = 'b'");
+    assert_eq!(
+        stmt,
+        Statement::SetVariable {
+            local: false,
+            variable: "a".into(),
+            value: SetVariableValue::Literal(Value::SingleQuotedString("b".into())),
+        }
+    );
+
+    let stmt = pg_and_generic().verified_stmt("SET a = 0");
+    assert_eq!(
+        stmt,
+        Statement::SetVariable {
+            local: false,
+            variable: "a".into(),
+            value: SetVariableValue::Literal(Value::Long(0)),
+        }
+    );
+
+    let stmt = pg_and_generic().verified_stmt("SET a = DEFAULT");
+    assert_eq!(
+        stmt,
+        Statement::SetVariable {
+            local: false,
+            variable: "a".into(),
+            value: SetVariableValue::Ident("DEFAULT".into()),
+        }
+    );
+
+    let stmt = pg_and_generic().verified_stmt("SET LOCAL a = b");
+    assert_eq!(
+        stmt,
+        Statement::SetVariable {
+            local: true,
+            variable: "a".into(),
+            value: SetVariableValue::Ident("b".into()),
+        }
+    );
+
+    pg_and_generic().one_statement_parses_to("SET a TO b", "SET a = b");
+    pg_and_generic().one_statement_parses_to("SET SESSION a = b", "SET a = b");
+
+    assert_eq!(
+        pg_and_generic().parse_sql_statements("SET"),
+        Err(ParserError::ParserError(
+            "Expected identifier, found: EOF".to_string()
+        )),
+    );
+
+    assert_eq!(
+        pg_and_generic().parse_sql_statements("SET a b"),
+        Err(ParserError::ParserError(
+            "Expected equals sign or TO, found: b".to_string()
+        )),
+    );
+
+    assert_eq!(
+        pg_and_generic().parse_sql_statements("SET a ="),
+        Err(ParserError::ParserError(
+            "Expected variable value, found: EOF".to_string()
+        )),
+    );
+}
+
+#[test]
+fn parse_show() {
+    let stmt = pg_and_generic().verified_stmt("SHOW a");
+    assert_eq!(
+        stmt,
+        Statement::ShowVariable {
+            variable: "a".into()
+        }
+    );
+
+    let stmt = pg_and_generic().verified_stmt("SHOW ALL");
+    assert_eq!(
+        stmt,
+        Statement::ShowVariable {
+            variable: "ALL".into()
+        }
+    )
+}
+
 fn pg() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(PostgreSqlDialect {})],