Add Postgres-specific PREPARE, EXECUTE and DEALLOCATE (#243)

Adds top-statements PREPARE, EXECUTE and DEALLOCATE for Postgres-specific feature prepared statement.

src/ast/mod.rs

@@ -553,12 +553,27 @@ pub enum Statement {
     Rollback { chain: bool },
     /// CREATE SCHEMA
     CreateSchema { schema_name: ObjectName },
-
-    /// ASSERT <condition> [AS <message>]
+    /// `ASSERT <condition> [AS <message>]`
     Assert {
         condition: Expr,
         message: Option<Expr>,
     },
+    /// `DEALLOCATE [ PREPARE ] { name | ALL }`
+    ///
+    /// Note: this is a PostgreSQL-specific statement.
+    Deallocate { name: Ident, prepare: bool },
+    /// `EXECUTE name [ ( parameter [, ...] ) ]`
+    ///
+    /// Note: this is a PostgreSQL-specific statement.
+    Execute { name: Ident, parameters: Vec<Expr> },
+    /// `PREPARE name [ ( data_type [, ...] ) ] AS statement`
+    ///
+    /// Note: this is a PostgreSQL-specific statement.
+    Prepare {
+        name: Ident,
+        data_types: Vec<DataType>,
+        statement: Box<Statement>,
+    },
 }
 
 impl fmt::Display for Statement {
@@ -834,6 +849,30 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::Deallocate { name, prepare } => write!(
+                f,
+                "DEALLOCATE {prepare}{name}",
+                prepare = if *prepare { "PREPARE " } else { "" },
+                name = name,
+            ),
+            Statement::Execute { name, parameters } => {
+                write!(f, "EXECUTE {}", name)?;
+                if !parameters.is_empty() {
+                    write!(f, "({})", display_comma_separated(parameters))?;
+                }
+                Ok(())
+            }
+            Statement::Prepare {
+                name,
+                data_types,
+                statement,
+            } => {
+                write!(f, "PREPARE {} ", name)?;
+                if !data_types.is_empty() {
+                    write!(f, "({}) ", display_comma_separated(data_types))?;
+                }
+                write!(f, "AS {}", statement)
+            }
         }
     }
 }

src/parser.rs

@@ -149,6 +149,11 @@ impl Parser {
                 Keyword::COMMIT => Ok(self.parse_commit()?),
                 Keyword::ROLLBACK => Ok(self.parse_rollback()?),
                 Keyword::ASSERT => Ok(self.parse_assert()?),
+                // `PREPARE`, `EXECUTE` and `DEALLOCATE` are Postgres-specific
+                // syntaxes. They are used for Postgres prepared statement.
+                Keyword::DEALLOCATE => Ok(self.parse_deallocate()?),
+                Keyword::EXECUTE => Ok(self.parse_execute()?),
+                Keyword::PREPARE => Ok(self.parse_prepare()?),
                 _ => self.expected("an SQL statement", Token::Word(w)),
             },
             Token::LParen => {
@@ -2398,6 +2403,42 @@ impl Parser {
             Ok(false)
         }
     }
+
+    fn parse_deallocate(&mut self) -> Result<Statement, ParserError> {
+        let prepare = self.parse_keyword(Keyword::PREPARE);
+        let name = self.parse_identifier()?;
+        Ok(Statement::Deallocate { name, prepare })
+    }
+
+    fn parse_execute(&mut self) -> Result<Statement, ParserError> {
+        let name = self.parse_identifier()?;
+
+        let mut parameters = vec![];
+        if self.consume_token(&Token::LParen) {
+            parameters = self.parse_comma_separated(Parser::parse_expr)?;
+            self.expect_token(&Token::RParen)?;
+        }
+
+        Ok(Statement::Execute { name, parameters })
+    }
+
+    fn parse_prepare(&mut self) -> Result<Statement, ParserError> {
+        let name = self.parse_identifier()?;
+
+        let mut data_types = vec![];
+        if self.consume_token(&Token::LParen) {
+            data_types = self.parse_comma_separated(Parser::parse_data_type)?;
+            self.expect_token(&Token::RParen)?;
+        }
+
+        self.expect_keyword(Keyword::AS)?;
+        let statement = Box::new(self.parse_statement()?);
+        Ok(Statement::Prepare {
+            name,
+            data_types,
+            statement,
+        })
+    }
 }
 
 impl Word {

tests/sqlparser_postgres.rs

@@ -423,6 +423,136 @@ fn parse_show() {
     )
 }
 
+#[test]
+fn parse_deallocate() {
+    let stmt = pg_and_generic().verified_stmt("DEALLOCATE a");
+    assert_eq!(
+        stmt,
+        Statement::Deallocate {
+            name: "a".into(),
+            prepare: false,
+        }
+    );
+
+    let stmt = pg_and_generic().verified_stmt("DEALLOCATE ALL");
+    assert_eq!(
+        stmt,
+        Statement::Deallocate {
+            name: "ALL".into(),
+            prepare: false,
+        }
+    );
+
+    let stmt = pg_and_generic().verified_stmt("DEALLOCATE PREPARE a");
+    assert_eq!(
+        stmt,
+        Statement::Deallocate {
+            name: "a".into(),
+            prepare: true,
+        }
+    );
+
+    let stmt = pg_and_generic().verified_stmt("DEALLOCATE PREPARE ALL");
+    assert_eq!(
+        stmt,
+        Statement::Deallocate {
+            name: "ALL".into(),
+            prepare: true,
+        }
+    );
+}
+
+#[test]
+fn parse_execute() {
+    let stmt = pg_and_generic().verified_stmt("EXECUTE a");
+    assert_eq!(
+        stmt,
+        Statement::Execute {
+            name: "a".into(),
+            parameters: vec![],
+        }
+    );
+
+    let stmt = pg_and_generic().verified_stmt("EXECUTE a(1, 't')");
+
+    #[cfg(feature = "bigdecimal")]
+    assert_eq!(
+        stmt,
+        Statement::Execute {
+            name: "a".into(),
+            parameters: vec![
+                Expr::Value(Value::Number(bigdecimal::BigDecimal::from(1))),
+                Expr::Value(Value::SingleQuotedString("t".to_string()))
+            ],
+        }
+    );
+}
+
+#[test]
+fn parse_prepare() {
+    let stmt =
+        pg_and_generic().verified_stmt("PREPARE a AS INSERT INTO customers VALUES (a1, a2, a3)");
+    let sub_stmt = match stmt {
+        Statement::Prepare {
+            name,
+            data_types,
+            statement,
+            ..
+        } => {
+            assert_eq!(name, "a".into());
+            assert!(data_types.is_empty());
+
+            statement
+        }
+        _ => unreachable!(),
+    };
+    match sub_stmt.as_ref() {
+        Statement::Insert {
+            table_name,
+            columns,
+            source,
+            ..
+        } => {
+            assert_eq!(table_name.to_string(), "customers");
+            assert!(columns.is_empty());
+
+            let expected_values = [vec![
+                Expr::Identifier("a1".into()),
+                Expr::Identifier("a2".into()),
+                Expr::Identifier("a3".into()),
+            ]];
+            match &source.body {
+                SetExpr::Values(Values(values)) => assert_eq!(values.as_slice(), &expected_values),
+                _ => unreachable!(),
+            }
+        }
+        _ => unreachable!(),
+    };
+
+    let stmt = pg_and_generic()
+        .verified_stmt("PREPARE a (INT, TEXT) AS SELECT * FROM customers WHERE customers.id = a1");
+    let sub_stmt = match stmt {
+        Statement::Prepare {
+            name,
+            data_types,
+            statement,
+            ..
+        } => {
+            assert_eq!(name, "a".into());
+            assert_eq!(data_types, vec![DataType::Int, DataType::Text]);
+
+            statement
+        }
+        _ => unreachable!(),
+    };
+    assert_eq!(
+        sub_stmt,
+        Box::new(Statement::Query(Box::new(pg_and_generic().verified_query(
+            "SELECT * FROM customers WHERE customers.id = a1"
+        ))))
+    );
+}
+
 fn pg() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(PostgreSqlDialect {})],