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

Author: silathdiir

src/ast/mod.rs

@@ -551,12 +551,34 @@ pub enum Statement {
     Rollback { chain: bool },
     /// CREATE SCHEMA
     CreateSchema { schema_name: ObjectName },
-
     /// ASSERT <condition> [AS <message>]
     Assert {
         condition: Expr,
         message: Option<Expr>,
     },
+    /// DEALLOCATE [ PREPARE ] { name | ALL }
+    ///
+    /// Note: this is a PostgreSQL-specific statement.
+    Deallocate {
+        name: Option<ObjectName>,
+        all: bool,
+        prepare: bool,
+    },
+    /// EXECUTE name [ ( parameter [, ...] ) ]
+    ///
+    /// Note: this is a PostgreSQL-specific statement.
+    Execute {
+        name: ObjectName,
+        parameters: Vec<Expr>,
+    },
+    /// PREPARE name [ ( data_type [, ...] ) ] AS statement
+    ///
+    /// Note: this is a PostgreSQL-specific statement.
+    Prepare {
+        name: ObjectName,
+        data_types: Vec<DataType>,
+        statement: Box<Statement>,
+    },
 }
 
 impl fmt::Display for Statement {
@@ -824,6 +846,36 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::Deallocate { name, all, prepare } => {
+                f.write_str("DEALLOCATE ")?;
+                if *prepare {
+                    f.write_str("PREPARE ")?;
+                }
+                if *all {
+                    f.write_str("ALL")
+                } else {
+                    write!(f, "{}", name.as_ref().unwrap())
+                }
+            }
+            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/dialect/postgresql.rs

@@ -20,7 +20,7 @@ impl Dialect for PostgreSqlDialect {
         // See https://www.postgresql.org/docs/11/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
         // We don't yet support identifiers beginning with "letters with
         // diacritical marks and non-Latin letters"
-        (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || ch == '_'
+        (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || ch == '_' || ch == '$'
     }
 
     fn is_identifier_part(&self, ch: char) -> bool {

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 => {
@@ -2386,6 +2391,49 @@ impl Parser {
             Ok(false)
         }
     }
+
+    fn parse_deallocate(&mut self) -> Result<Statement, ParserError> {
+        let prepare = self.parse_keyword(Keyword::PREPARE);
+        let all = self.parse_keyword(Keyword::ALL);
+
+        let name = if all {
+            None
+        } else {
+            Some(self.parse_object_name()?)
+        };
+
+        Ok(Statement::Deallocate { name, all, prepare })
+    }
+
+    fn parse_execute(&mut self) -> Result<Statement, ParserError> {
+        let name = self.parse_object_name()?;
+
+        let mut parameters = vec![];
+        if self.expect_token(&Token::LParen).is_ok() {
+            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_object_name()?;
+
+        let mut data_types = vec![];
+        if self.expect_token(&Token::LParen).is_ok() {
+            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,137 @@ fn parse_show() {
     )
 }
 
+#[test]
+fn parse_deallocate() {
+    let stmt = pg().verified_stmt("DEALLOCATE a");
+    assert_eq!(
+        stmt,
+        Statement::Deallocate {
+            name: Some(ObjectName(vec!["a".into()])),
+            all: false,
+            prepare: false,
+        }
+    );
+
+    let stmt = pg().verified_stmt("DEALLOCATE ALL");
+    assert_eq!(
+        stmt,
+        Statement::Deallocate {
+            name: None,
+            all: true,
+            prepare: false,
+        }
+    );
+
+    let stmt = pg().verified_stmt("DEALLOCATE PREPARE a");
+    assert_eq!(
+        stmt,
+        Statement::Deallocate {
+            name: Some(ObjectName(vec!["a".into()])),
+            all: false,
+            prepare: true,
+        }
+    );
+
+    let stmt = pg().verified_stmt("DEALLOCATE PREPARE ALL");
+    assert_eq!(
+        stmt,
+        Statement::Deallocate {
+            name: None,
+            all: true,
+            prepare: true,
+        }
+    );
+}
+
+#[test]
+fn parse_execute() {
+    let stmt = pg().verified_stmt("EXECUTE a");
+    assert_eq!(
+        stmt,
+        Statement::Execute {
+            name: ObjectName(vec!["a".into()]),
+            parameters: vec![],
+        }
+    );
+
+    let stmt = pg().verified_stmt("EXECUTE a(1, 't')");
+    assert_eq!(
+        stmt,
+        Statement::Execute {
+            name: ObjectName(vec!["a".into()]),
+            parameters: vec![
+                Expr::Value(Value::Number("1".to_string())),
+                Expr::Value(Value::SingleQuotedString("t".to_string()))
+            ],
+        }
+    );
+}
+
+#[test]
+fn parse_prepare() {
+    let stmt = pg().verified_stmt("PREPARE a AS INSERT INTO customers VALUES ($1, $2, $3)");
+    let sub_stmt = match stmt {
+        Statement::Prepare {
+            name,
+            data_types,
+            statement,
+            ..
+        } => {
+            assert_eq!(name, ObjectName(vec!["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("$1".into()),
+                Expr::Identifier("$2".into()),
+                Expr::Identifier("$3".into()),
+            ]];
+            match &source.body {
+                SetExpr::Values(Values(values)) => assert_eq!(values.as_slice(), &expected_values),
+                _ => unreachable!(),
+            }
+        }
+        _ => unreachable!(),
+    };
+
+    let stmt = pg()
+        .verified_stmt("PREPARE a (INT, TEXT) AS SELECT * FROM customers WHERE customers.id = $1");
+    let sub_stmt = match stmt {
+        Statement::Prepare {
+            name,
+            data_types,
+            statement,
+            ..
+        } => {
+            assert_eq!(name, ObjectName(vec!["a".into()]));
+            assert_eq!(data_types, vec![DataType::Int, DataType::Text]);
+
+            statement
+        }
+        _ => unreachable!(),
+    };
+    assert_eq!(
+        sub_stmt,
+        Box::new(Statement::Query(Box::new(pg().verified_query(
+            "SELECT * FROM customers WHERE customers.id = $1"
+        ))))
+    );
+}
+
 fn pg() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(PostgreSqlDialect {})],