Snowflake: parse ALTER TAG SET/UNSET MASKING POLICY

Author: sabir-akhadov-localstack

src/ast/ddl.rs

@@ -1413,6 +1413,61 @@ impl fmt::Display for AlterMaskingPolicyOperation {
     }
 }
 
+/// An operation in an `ALTER TAG` statement.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum AlterTagOperation {
+    /// `RENAME TO <new_name>`
+    RenameTo {
+        /// The new tag name.
+        new_name: ObjectName,
+    },
+    /// `SET MASKING POLICY <p> [, MASKING POLICY <p> ...] [FORCE]`
+    SetMaskingPolicy {
+        /// The masking policies to attach.
+        policies: Vec<ObjectName>,
+        /// `FORCE` flag.
+        force: bool,
+    },
+    /// `UNSET MASKING POLICY <p> [, MASKING POLICY <p> ...]`
+    UnsetMaskingPolicy {
+        /// The masking policies to detach.
+        policies: Vec<ObjectName>,
+    },
+}
+
+impl fmt::Display for AlterTagOperation {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            AlterTagOperation::RenameTo { new_name } => write!(f, "RENAME TO {new_name}"),
+            AlterTagOperation::SetMaskingPolicy { policies, force } => {
+                write!(f, "SET ")?;
+                for (i, p) in policies.iter().enumerate() {
+                    if i > 0 {
+                        write!(f, ", ")?;
+                    }
+                    write!(f, "MASKING POLICY {p}")?;
+                }
+                if *force {
+                    write!(f, " FORCE")?;
+                }
+                Ok(())
+            }
+            AlterTagOperation::UnsetMaskingPolicy { policies } => {
+                write!(f, "UNSET ")?;
+                for (i, p) in policies.iter().enumerate() {
+                    if i > 0 {
+                        write!(f, ", ")?;
+                    }
+                    write!(f, "MASKING POLICY {p}")?;
+                }
+                Ok(())
+            }
+        }
+    }
+}
+
 impl fmt::Display for AlterColumnOperation {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {

src/ast/mod.rs

@@ -64,7 +64,7 @@ pub use self::dcl::{
 pub use self::ddl::{
     Alignment, AlterCollation, AlterCollationOperation, AlterColumnOperation, AlterConnectorOwner,
     AlterFunction, AlterFunctionAction, AlterFunctionKind, AlterFunctionOperation,
-    AlterMaskingPolicyOperation,
+    AlterMaskingPolicyOperation, AlterTagOperation,
     AlterIndexOperation, AlterProcedure, AlterProcedureOperation,
     AlterOperator, AlterOperatorClass, AlterOperatorClassOperation,
     AlterOperatorFamily, AlterOperatorFamilyOperation, AlterOperatorOperation, AlterPolicy,
@@ -5276,14 +5276,16 @@ pub enum Statement {
     },
     /// ```sql
     /// ALTER TAG [IF EXISTS] <name> RENAME TO <new_name>
+    /// ALTER TAG [IF EXISTS] <name> SET MASKING POLICY <p> [, MASKING POLICY <p> ...] [FORCE]
+    /// ALTER TAG [IF EXISTS] <name> UNSET MASKING POLICY <p> [, MASKING POLICY <p> ...]
     /// ```
     AlterTag {
         /// `IF EXISTS` flag.
         if_exists: bool,
         /// Tag name.
         name: ObjectName,
-        /// New tag name.
-        new_name: ObjectName,
+        /// The operation to apply to the tag.
+        operation: AlterTagOperation,
     },
     /// ```sql
     /// DROP TAG [IF EXISTS] <name>
@@ -7810,11 +7812,11 @@ impl fmt::Display for Statement {
             Statement::AlterTag {
                 if_exists,
                 name,
-                new_name,
+                operation,
             } => {
                 write!(
                     f,
-                    "ALTER TAG {if_exists}{name} RENAME TO {new_name}",
+                    "ALTER TAG {if_exists}{name} {operation}",
                     if_exists = if *if_exists { "IF EXISTS " } else { "" },
                 )
             }

src/dialect/snowflake.rs

@@ -28,7 +28,7 @@ use crate::ast::helpers::stmt_data_loading::{
 };
 use crate::ast::{
     AlterExternalVolumeOperation, AlterFileFormatOperation, AlterMaskingPolicyOperation,
-    AlterProcedure, AlterProcedureOperation, AlterStageOperation, AlterTable,
+    AlterProcedure, AlterProcedureOperation, AlterStageOperation, AlterTable, AlterTagOperation,
     AlterTableOperation, AlterTableType, CatalogRestAuthentication, CatalogRestConfig,
     CatalogSource, CatalogSyncNamespaceMode, CatalogTableFormat, ColumnOption, ColumnPolicy,
     ColumnPolicyProperty, ContactEntry, CopyIntoSnowflakeKind, CreateTable, CreateTableLikeKind,
@@ -3252,16 +3252,42 @@ fn parse_create_tag(or_replace: bool, parser: &mut Parser) -> Result<Statement,
     })
 }
 
-/// Parse `ALTER TAG [IF EXISTS] <name> RENAME TO <new_name>`
+/// Parse a comma-separated list of `MASKING POLICY <name>` items, where the
+/// `MASKING POLICY` keywords are repeated before each policy name.
+fn parse_masking_policy_list(parser: &mut Parser) -> Result<Vec<ObjectName>, ParserError> {
+    let mut policies = Vec::new();
+    loop {
+        parser.expect_keywords(&[Keyword::MASKING, Keyword::POLICY])?;
+        policies.push(parser.parse_object_name(false)?);
+        if !parser.consume_token(&Token::Comma) {
+            break;
+        }
+    }
+    Ok(policies)
+}
+
+/// Parse `ALTER TAG [IF EXISTS] <name> { RENAME TO <new_name>
+///   | SET MASKING POLICY <p> [, MASKING POLICY <p> ...] [FORCE]
+///   | UNSET MASKING POLICY <p> [, MASKING POLICY <p> ...] }`
 fn parse_alter_tag(parser: &mut Parser) -> Result<Statement, ParserError> {
     let if_exists = parser.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
     let name = parser.parse_object_name(false)?;
-    parser.expect_keywords(&[Keyword::RENAME, Keyword::TO])?;
-    let new_name = parser.parse_object_name(false)?;
+    let operation = if parser.parse_keyword(Keyword::SET) {
+        let policies = parse_masking_policy_list(parser)?;
+        let force = parser.parse_keyword(Keyword::FORCE);
+        AlterTagOperation::SetMaskingPolicy { policies, force }
+    } else if parser.parse_keyword(Keyword::UNSET) {
+        let policies = parse_masking_policy_list(parser)?;
+        AlterTagOperation::UnsetMaskingPolicy { policies }
+    } else {
+        parser.expect_keywords(&[Keyword::RENAME, Keyword::TO])?;
+        let new_name = parser.parse_object_name(false)?;
+        AlterTagOperation::RenameTo { new_name }
+    };
     Ok(Statement::AlterTag {
         if_exists,
         name,
-        new_name,
+        operation,
     })
 }
 

tests/sqlparser_snowflake.rs

@@ -9164,3 +9164,84 @@ fn parse_snowflake_create_view_column_masking_policy() {
         other => panic!("expected CreateView, got {other:?}"),
     }
 }
+
+#[test]
+fn parse_snowflake_alter_tag_rename() {
+    match snowflake().verified_stmt("ALTER TAG IF EXISTS t RENAME TO u") {
+        Statement::AlterTag {
+            if_exists,
+            name,
+            operation,
+        } => {
+            assert!(if_exists);
+            assert_eq!("t", name.to_string());
+            assert_eq!(
+                AlterTagOperation::RenameTo {
+                    new_name: ObjectName::from(vec![Ident::new("u")]),
+                },
+                operation
+            );
+        }
+        other => panic!("expected AlterTag, got {other:?}"),
+    }
+}
+
+#[test]
+fn parse_snowflake_alter_tag_set_masking_policy() {
+    match snowflake().verified_stmt("ALTER TAG t SET MASKING POLICY p") {
+        Statement::AlterTag {
+            if_exists,
+            operation,
+            ..
+        } => {
+            assert!(!if_exists);
+            assert_eq!(
+                AlterTagOperation::SetMaskingPolicy {
+                    policies: vec![ObjectName::from(vec![Ident::new("p")])],
+                    force: false,
+                },
+                operation
+            );
+        }
+        other => panic!("expected AlterTag, got {other:?}"),
+    }
+
+    match snowflake()
+        .verified_stmt("ALTER TAG IF EXISTS t SET MASKING POLICY p, MASKING POLICY q FORCE")
+    {
+        Statement::AlterTag {
+            if_exists,
+            operation,
+            ..
+        } => {
+            assert!(if_exists);
+            assert_eq!(
+                AlterTagOperation::SetMaskingPolicy {
+                    policies: vec![
+                        ObjectName::from(vec![Ident::new("p")]),
+                        ObjectName::from(vec![Ident::new("q")]),
+                    ],
+                    force: true,
+                },
+                operation
+            );
+        }
+        other => panic!("expected AlterTag, got {other:?}"),
+    }
+}
+
+#[test]
+fn parse_snowflake_alter_tag_unset_masking_policy() {
+    match snowflake().verified_stmt("ALTER TAG t UNSET MASKING POLICY p, MASKING POLICY q") {
+        Statement::AlterTag { operation, .. } => assert_eq!(
+            AlterTagOperation::UnsetMaskingPolicy {
+                policies: vec![
+                    ObjectName::from(vec![Ident::new("p")]),
+                    ObjectName::from(vec![Ident::new("q")]),
+                ],
+            },
+            operation
+        ),
+        other => panic!("expected AlterTag, got {other:?}"),
+    }
+}