PostgreSQL: Support ALTER EXTENSION ... SET SCHEMA

fuzz/fuzz_seeds/postgres_alter_extension.sql

@@ -0,0 +1 @@
+ALTER EXTENSION citext SET SCHEMA other;

src/ast/ddl.rs

@@ -4867,6 +4867,61 @@ impl Spanned for DropExtension {
     }
 }
 
+/// ALTER EXTENSION statement.
+/// Note: this is a PostgreSQL-specific statement.
+///
+/// # References
+///
+/// PostgreSQL Documentation:
+/// <https://www.postgresql.org/docs/current/sql-alterextension.html>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct AlterExtension {
+    /// Name of the extension being altered.
+    pub name: Ident,
+    /// The operation to perform on the extension.
+    pub operation: AlterExtensionOperation,
+}
+
+/// Operations supported by `ALTER EXTENSION`.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum AlterExtensionOperation {
+    /// Move the extension to another schema.
+    ///
+    /// ```sql
+    /// ALTER EXTENSION name SET SCHEMA new_schema
+    /// ```
+    SetSchema {
+        /// Target schema name.
+        schema_name: ObjectName,
+    },
+}
+
+impl fmt::Display for AlterExtensionOperation {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            AlterExtensionOperation::SetSchema { schema_name } => {
+                write!(f, "SET SCHEMA {schema_name}")
+            }
+        }
+    }
+}
+
+impl fmt::Display for AlterExtension {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "ALTER EXTENSION {} {}", self.name, self.operation)
+    }
+}
+
+impl Spanned for AlterExtension {
+    fn span(&self) -> Span {
+        Span::empty()
+    }
+}
+
 /// CREATE COLLATION statement.
 /// Note: this is a PostgreSQL-specific statement.
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]

src/ast/mod.rs

@@ -65,12 +65,12 @@ pub use self::dcl::{
 };
 pub use self::ddl::{
     AggregateModifyKind, Alignment, AlterCollation, AlterCollationOperation, AlterColumnOperation,
-    AlterConnectorOwner, AlterFunction, AlterFunctionAction, AlterFunctionKind,
-    AlterFunctionOperation, AlterGenericOption, AlterIndexOperation, AlterOperator,
-    AlterOperatorClass, AlterOperatorClassOperation, AlterOperatorFamily,
-    AlterOperatorFamilyOperation, AlterOperatorOperation, AlterPolicy, AlterPolicyOperation,
-    AlterSchema, AlterSchemaOperation, AlterTable, AlterTableAlgorithm, AlterTableLock,
-    AlterTableOperation, AlterTableType, AlterTextSearch, AlterTextSearchOperation,
+    AlterConnectorOwner, AlterExtension, AlterExtensionOperation, AlterFunction,
+    AlterFunctionAction, AlterFunctionKind, AlterFunctionOperation, AlterGenericOption,
+    AlterIndexOperation, AlterOperator, AlterOperatorClass, AlterOperatorClassOperation,
+    AlterOperatorFamily, AlterOperatorFamilyOperation, AlterOperatorOperation, AlterPolicy,
+    AlterPolicyOperation, AlterSchema, AlterSchemaOperation, AlterTable, AlterTableAlgorithm,
+    AlterTableLock, AlterTableOperation, AlterTableType, AlterTextSearch, AlterTextSearchOperation,
     AlterTextSearchOption, AlterTrigger, AlterTriggerOperation, AlterType, AlterTypeAddValue,
     AlterTypeAddValuePosition, AlterTypeOperation, AlterTypeRename, AlterTypeRenameValue,
     AlterView, AlterViewOperation, ClusteredBy, ColumnDef, ColumnOption, ColumnOptionDef,
@@ -3943,6 +3943,11 @@ pub enum Statement {
     /// See [PostgreSQL](https://www.postgresql.org/docs/current/sql-altercollation.html)
     AlterCollation(AlterCollation),
     /// ```sql
+    /// ALTER EXTENSION
+    /// ```
+    /// See [PostgreSQL](https://www.postgresql.org/docs/current/sql-alterextension.html)
+    AlterExtension(AlterExtension),
+    /// ```sql
     /// ALTER OPERATOR
     /// ```
     /// See [PostgreSQL](https://www.postgresql.org/docs/current/sql-alteroperator.html)
@@ -5785,6 +5790,7 @@ impl fmt::Display for Statement {
                 write!(f, "ALTER TYPE {name} {operation}")
             }
             Statement::AlterCollation(alter_collation) => write!(f, "{alter_collation}"),
+            Statement::AlterExtension(alter_extension) => write!(f, "{alter_extension}"),
             Statement::AlterOperator(alter_operator) => write!(f, "{alter_operator}"),
             Statement::AlterOperatorFamily(alter_operator_family) => {
                 write!(f, "{alter_operator_family}")
@@ -12767,6 +12773,12 @@ impl From<AlterCollation> for Statement {
     }
 }
 
+impl From<AlterExtension> for Statement {
+    fn from(a: AlterExtension) -> Self {
+        Self::AlterExtension(a)
+    }
+}
+
 impl From<AlterView> for Statement {
     fn from(a: AlterView) -> Self {
         Self::AlterView(a)

src/ast/spans.rs

@@ -276,6 +276,7 @@ impl Spanned for Values {
 /// - [Statement::CreateExtension]
 /// - [Statement::CreateCollation]
 /// - [Statement::AlterCollation]
+/// - [Statement::AlterExtension]
 /// - [Statement::Fetch]
 /// - [Statement::Flush]
 /// - [Statement::Discard]
@@ -416,6 +417,7 @@ impl Spanned for Statement {
             Statement::AlterFunction { .. } => Span::empty(),
             Statement::AlterType { .. } => Span::empty(),
             Statement::AlterCollation { .. } => Span::empty(),
+            Statement::AlterExtension { .. } => Span::empty(),
             Statement::AlterOperator { .. } => Span::empty(),
             Statement::AlterOperatorFamily { .. } => Span::empty(),
             Statement::AlterOperatorClass { .. } => Span::empty(),

src/parser/mod.rs

@@ -12021,6 +12021,7 @@ impl<'a> Parser<'a> {
             }
             Keyword::TYPE => self.parse_alter_type(),
             Keyword::COLLATION => self.parse_alter_collation().map(Into::into),
+            Keyword::EXTENSION => self.parse_alter_extension().map(Into::into),
             Keyword::TABLE => self.parse_alter_table(None),
             Keyword::ICEBERG => {
                 self.expect_keyword(Keyword::TABLE)?;
@@ -12070,7 +12071,7 @@ impl<'a> Parser<'a> {
             _ => {
                 self.index = index;
                 self.expected_ref(
-                    "one of VIEW or MATERIALIZED VIEW or TYPE or COLLATION or TABLE or INDEX or FUNCTION or AGGREGATE or ROLE or POLICY or TRIGGER or CONNECTOR or ICEBERG or FOREIGN TABLE or SCHEMA or USER or OPERATOR",
+                    "one of VIEW or MATERIALIZED VIEW or TYPE or COLLATION or TABLE or INDEX or FUNCTION or AGGREGATE or ROLE or POLICY or TRIGGER or CONNECTOR or ICEBERG or FOREIGN TABLE or SCHEMA or USER or OPERATOR or EXTENSION",
                     self.peek_token_ref(),
                 )
             }
@@ -12503,6 +12504,25 @@ impl<'a> Parser<'a> {
         Ok(AlterCollation { name, operation })
     }
 
+    /// Parse a [Statement::AlterExtension]
+    ///
+    /// [PostgreSQL Documentation](https://www.postgresql.org/docs/current/sql-alterextension.html)
+    pub fn parse_alter_extension(&mut self) -> Result<AlterExtension, ParserError> {
+        let name = self.parse_identifier()?;
+        let operation = if self.parse_keywords(&[Keyword::SET, Keyword::SCHEMA]) {
+            AlterExtensionOperation::SetSchema {
+                schema_name: self.parse_object_name(false)?,
+            }
+        } else {
+            return self.expected_ref(
+                "SET SCHEMA after ALTER EXTENSION name",
+                self.peek_token_ref(),
+            );
+        };
+
+        Ok(AlterExtension { name, operation })
+    }
+
     /// Parse a [Statement::AlterOperator]
     ///
     /// [PostgreSQL Documentation](https://www.postgresql.org/docs/current/sql-alteroperator.html)

tests/sqlparser_common.rs

@@ -16196,7 +16196,7 @@ fn parse_alter_trigger_spans() {
 
 #[test]
 fn parse_alter_rejects_missing_or_unknown_object_type() {
-    let objects = "one of VIEW or MATERIALIZED VIEW or TYPE or COLLATION or TABLE or INDEX or FUNCTION or AGGREGATE or ROLE or POLICY or TRIGGER or CONNECTOR or ICEBERG or FOREIGN TABLE or SCHEMA or USER or OPERATOR";
+    let objects = "one of VIEW or MATERIALIZED VIEW or TYPE or COLLATION or TABLE or INDEX or FUNCTION or AGGREGATE or ROLE or POLICY or TRIGGER or CONNECTOR or ICEBERG or FOREIGN TABLE or SCHEMA or USER or OPERATOR or EXTENSION";
     for (sql, found) in [
         ("ALTER", "EOF"),
         ("ALTER 1", "1"),

tests/sqlparser_postgres.rs

@@ -1068,6 +1068,13 @@ fn parse_drop_extension() {
             cascade_or_restrict: Some(ReferentialAction::Restrict),
         })
     );
+
+    let stmt = pg_and_generic().verified_stmt("DROP EXTENSION extension_name");
+    if let Statement::DropExtension(drop_extension) = stmt {
+        assert_eq!(Span::empty(), drop_extension.span());
+    } else {
+        unreachable!("expected DROP EXTENSION statement");
+    }
 }
 
 #[test]
@@ -13518,3 +13525,45 @@ fn parse_case_ending_in_end_keyword_is_not_silently_dropped() {
         ParserError::ParserError("Expected: end of statement, found: END".to_string())
     );
 }
+
+#[test]
+fn parse_alter_extension_set_schema() {
+    pg_and_generic().verified_stmt("ALTER EXTENSION ext_name SET SCHEMA other_schema");
+    pg_and_generic().verified_stmt("ALTER EXTENSION ext_name SET SCHEMA schema.other");
+    pg_and_generic().verified_stmt("ALTER EXTENSION ext_name SET SCHEMA \"other schema\"");
+
+    let stmt = pg_and_generic().verified_stmt("ALTER EXTENSION ext_name SET SCHEMA other_schema");
+    assert_eq!(Span::empty(), stmt.span());
+    match stmt {
+        Statement::AlterExtension(alter_extension) => {
+            assert_eq!(Span::empty(), alter_extension.span());
+            let AlterExtension { name, operation } = alter_extension;
+            assert_eq!(name, Ident::new("ext_name"));
+            assert_eq!(
+                operation,
+                AlterExtensionOperation::SetSchema {
+                    schema_name: ObjectName::from(Ident::new("other_schema"))
+                }
+            );
+        }
+        other => panic!("unexpected statement: {other:?}"),
+    }
+
+    for (sql, error) in [
+        (
+            "ALTER EXTENSION ext_name SET SCHEMA",
+            "Expected: identifier, found: EOF",
+        ),
+        (
+            "ALTER EXTENSION ext_name",
+            "Expected: SET SCHEMA after ALTER EXTENSION name, found: EOF",
+        ),
+        ("ALTER EXTENSION 1", "Expected: identifier, found: 1"),
+    ] {
+        assert_eq!(
+            pg_and_generic().parse_sql_statements(sql).unwrap_err(),
+            ParserError::ParserError(error.to_string()),
+            "{sql}"
+        );
+    }
+}