Adds support for pg CREATE EXTENSION (#1078)

apache · Jan 2, 2024 · 0be42ee · 0be42ee
1 parent 1d63466
commit 0be42ee
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 0 deletions.
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -1763,6 +1763,21 @@ pub enum Statement {
         query: Box<Query>,
     },
     /// ```sql
+    /// CREATE EXTENSION [ IF NOT EXISTS ] extension_name
+    ///     [ WITH ] [ SCHEMA schema_name ]
+    ///              [ VERSION version ]
+    ///              [ CASCADE ]
+    /// ```
+    ///
+    /// Note: this is a PostgreSQL-specific statement,
+    CreateExtension {
+        name: Ident,
+        if_not_exists: bool,
+        cascade: bool,
+        schema: Option<Ident>,
+        version: Option<Ident>,
+    },
+    /// ```sql
     /// FETCH
     /// ```
     /// Retrieve rows from a query using a cursor
@@ -3004,6 +3019,34 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::CreateExtension {
+                name,
+                if_not_exists,
+                cascade,
+                schema,
+                version,
+            } => {
+                write!(
+                    f,
+                    "CREATE EXTENSION {if_not_exists}{name}",
+                    if_not_exists = if *if_not_exists { "IF NOT EXISTS " } else { "" }
+                )?;
+                if *cascade || schema.is_some() || version.is_some() {
+                    write!(f, " WITH")?;
+
+                    if let Some(name) = schema {
+                        write!(f, " SCHEMA {name}")?;
+                    }
+                    if let Some(version) = version {
+                        write!(f, " VERSION {version}")?;
+                    }
+                    if *cascade {
+                        write!(f, " CASCADE")?;
+                    }
+                }
+
+                Ok(())
+            }
             Statement::CreateRole {
                 names,
                 if_not_exists,

diff --git a/src/keywords.rs b/src/keywords.rs
@@ -270,6 +270,7 @@ define_keywords!(
     EXPLICIT,
     EXPORT,
     EXTENDED,
+    EXTENSION,
     EXTERNAL,
     EXTRACT,
     FAIL,
@@ -706,6 +707,7 @@ define_keywords!(
     VAR_POP,
     VAR_SAMP,
     VERBOSE,
+    VERSION,
     VERSIONING,
     VIEW,
     VIRTUAL,

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -2970,6 +2970,8 @@ impl<'a> Parser<'a> {
                 "[EXTERNAL] TABLE or [MATERIALIZED] VIEW or FUNCTION after CREATE OR REPLACE",
                 self.peek_token(),
             )
+        } else if self.parse_keyword(Keyword::EXTENSION) {
+            self.parse_create_extension()
         } else if self.parse_keyword(Keyword::INDEX) {
             self.parse_create_index(false)
         } else if self.parse_keywords(&[Keyword::UNIQUE, Keyword::INDEX]) {
@@ -3988,6 +3990,39 @@ impl<'a> Parser<'a> {
         })
     }
 
+    pub fn parse_create_extension(&mut self) -> Result<Statement, ParserError> {
+        let if_not_exists = self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
+        let name = self.parse_identifier()?;
+
+        let (schema, version, cascade) = if self.parse_keyword(Keyword::WITH) {
+            let schema = if self.parse_keyword(Keyword::SCHEMA) {
+                Some(self.parse_identifier()?)
+            } else {
+                None
+            };
+
+            let version = if self.parse_keyword(Keyword::VERSION) {
+                Some(self.parse_identifier()?)
+            } else {
+                None
+            };
+
+            let cascade = self.parse_keyword(Keyword::CASCADE);
+
+            (schema, version, cascade)
+        } else {
+            (None, None, false)
+        };
+
+        Ok(Statement::CreateExtension {
+            name,
+            if_not_exists,
+            schema,
+            version,
+            cascade,
+        })
+    }
+
     //TODO: Implement parsing for Skewed and Clustered
     pub fn parse_hive_distribution(&mut self) -> Result<HiveDistributionStyle, ParserError> {
         if self.parse_keywords(&[Keyword::PARTITIONED, Keyword::BY]) {

diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -584,6 +584,23 @@ fn parse_alter_table_enable() {
     pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE TRIGGER USER");
     pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE TRIGGER trigger_name");
 }
+
+#[test]
+fn parse_create_extension() {
+    pg_and_generic().verified_stmt("CREATE EXTENSION extension_name");
+    pg_and_generic().verified_stmt("CREATE EXTENSION extension_name WITH SCHEMA schema_name");
+    pg_and_generic().verified_stmt("CREATE EXTENSION extension_name WITH VERSION version");
+    pg_and_generic().verified_stmt("CREATE EXTENSION extension_name WITH CASCADE");
+    pg_and_generic().verified_stmt(
+        "CREATE EXTENSION extension_name WITH SCHEMA schema_name VERSION version CASCADE",
+    );
+    pg_and_generic()
+        .verified_stmt("CREATE EXTENSION extension_name WITH SCHEMA schema_name CASCADE");
+    pg_and_generic().verified_stmt("CREATE EXTENSION extension_name WITH VERSION version CASCADE");
+    pg_and_generic()
+        .verified_stmt("CREATE EXTENSION extension_name WITH SCHEMA schema_name VERSION version");
+}
+
 #[test]
 fn parse_alter_table_alter_column() {
     pg().one_statement_parses_to(