From eb01545ce277cdb9035e49319e6506f2147cc0ef Mon Sep 17 00:00:00 2001
From: Alex Dukhno <alex.dukhno@icloud.com>
Date: Thu, 1 Oct 2020 21:03:45 +0300
Subject: [PATCH 1/2] create drop schema conditionally

---
 src/ast/mod.rs              | 19 +++++++++++++++++--
 src/parser.rs               |  6 +++++-
 tests/sqlparser_common.rs   |  2 +-
 tests/sqlparser_postgres.rs | 27 +++++++++++++++++++++++++++
 4 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 31171db7d..5f5d58ab2 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -558,7 +558,10 @@ pub enum Statement {
     /// `ROLLBACK [ TRANSACTION | WORK ] [ AND [ NO ] CHAIN ]`
     Rollback { chain: bool },
     /// CREATE SCHEMA
-    CreateSchema { schema_name: ObjectName },
+    CreateSchema {
+        schema_name: ObjectName,
+        if_not_exists: bool,
+    },
     /// `ASSERT <condition> [AS <message>]`
     Assert {
         condition: Expr,
@@ -829,7 +832,19 @@ impl fmt::Display for Statement {
             Statement::Rollback { chain } => {
                 write!(f, "ROLLBACK{}", if *chain { " AND CHAIN" } else { "" },)
             }
-            Statement::CreateSchema { schema_name } => write!(f, "CREATE SCHEMA {}", schema_name),
+            Statement::CreateSchema {
+                schema_name,
+                if_not_exists,
+            } => write!(
+                f,
+                "CREATE SCHEMA{if_not_exists}{name}",
+                if_not_exists = if *if_not_exists {
+                    " IF NOT EXISTS "
+                } else {
+                    " "
+                },
+                name = schema_name
+            ),
             Statement::Assert { condition, message } => {
                 write!(f, "ASSERT {}", condition)?;
                 if let Some(m) = message {
diff --git a/src/parser.rs b/src/parser.rs
index f7630b037..438c9f1ef 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -1078,8 +1078,12 @@ impl<'a> Parser<'a> {
     }
 
     pub fn parse_create_schema(&mut self) -> Result<Statement, ParserError> {
+        let if_not_exists = self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
         let schema_name = self.parse_object_name()?;
-        Ok(Statement::CreateSchema { schema_name })
+        Ok(Statement::CreateSchema {
+            schema_name,
+            if_not_exists,
+        })
     }
 
     pub fn parse_create_external_table(
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index 4050cdf0b..6b032d926 100644
--- a/tests/sqlparser_common.rs
+++ b/tests/sqlparser_common.rs
@@ -1204,7 +1204,7 @@ fn parse_create_schema() {
     let sql = "CREATE SCHEMA X";
 
     match verified_stmt(sql) {
-        Statement::CreateSchema { schema_name } => {
+        Statement::CreateSchema { schema_name, .. } => {
             assert_eq!(schema_name.to_string(), "X".to_owned())
         }
         _ => unreachable!(),
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
index 1083e9971..94fc96042 100644
--- a/tests/sqlparser_postgres.rs
+++ b/tests/sqlparser_postgres.rs
@@ -300,6 +300,33 @@ fn parse_bad_if_not_exists() {
     );
 }
 
+#[test]
+fn parse_create_schema_if_not_exists() {
+    let sql = "CREATE SCHEMA IF NOT EXISTS schema_name";
+    let ast = pg().verified_stmt(sql);
+    match ast {
+        Statement::CreateSchema {
+            if_not_exists: true,
+            schema_name,
+        } => assert_eq!("schema_name", schema_name.to_string()),
+        _ => unreachable!(),
+    }
+}
+
+#[test]
+fn parse_drop_schema_if_exists() {
+    let sql = "DROP SCHEMA IF EXISTS schema_name";
+    let ast = pg().verified_stmt(sql);
+    match ast {
+        Statement::Drop {
+            object_type,
+            if_exists: true,
+            ..
+        } => assert_eq!(object_type, ObjectType::Schema),
+        _ => unreachable!(),
+    }
+}
+
 #[test]
 fn parse_copy_example() {
     let sql = r#"COPY public.actor (actor_id, first_name, last_name, last_update, value) FROM stdin;

From 36b881248a391411399a08bce35ed454340cec5a Mon Sep 17 00:00:00 2001
From: Alex Dukhno <alex.dukhno@icloud.com>
Date: Fri, 2 Oct 2020 11:38:23 +0300
Subject: [PATCH 2/2] address review comments

---
 src/ast/mod.rs              | 8 ++------
 tests/sqlparser_postgres.rs | 2 +-
 2 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 5f5d58ab2..8a58207d7 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -837,12 +837,8 @@ impl fmt::Display for Statement {
                 if_not_exists,
             } => write!(
                 f,
-                "CREATE SCHEMA{if_not_exists}{name}",
-                if_not_exists = if *if_not_exists {
-                    " IF NOT EXISTS "
-                } else {
-                    " "
-                },
+                "CREATE SCHEMA {if_not_exists}{name}",
+                if_not_exists = if *if_not_exists { "IF NOT EXISTS " } else { "" },
                 name = schema_name
             ),
             Statement::Assert { condition, message } => {
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
index 94fc96042..ecc0e77e3 100644
--- a/tests/sqlparser_postgres.rs
+++ b/tests/sqlparser_postgres.rs
@@ -303,7 +303,7 @@ fn parse_bad_if_not_exists() {
 #[test]
 fn parse_create_schema_if_not_exists() {
     let sql = "CREATE SCHEMA IF NOT EXISTS schema_name";
-    let ast = pg().verified_stmt(sql);
+    let ast = pg_and_generic().verified_stmt(sql);
     match ast {
         Statement::CreateSchema {
             if_not_exists: true,