feat(query): extend tag support to view, udf, and procedure objects

src/meta/api/src/api_impl/tag_api.rs

@@ -16,6 +16,9 @@ use std::collections::HashSet;
 
 use databend_common_meta_app::KeyWithTenant;
 use databend_common_meta_app::id_generator::IdGenerator;
+use databend_common_meta_app::principal::ProcedureIdentity;
+use databend_common_meta_app::principal::ProcedureNameIdent;
+use databend_common_meta_app::principal::UdfIdent;
 use databend_common_meta_app::principal::connection_ident::ConnectionIdent;
 use databend_common_meta_app::principal::user_stage_ident::StageIdent;
 use databend_common_meta_app::schema::CreateTagReply;
@@ -87,6 +90,11 @@ fn get_object_key(tenant: &Tenant, object: &TaggableObject) -> String {
         TaggableObject::Stage { name } => StageIdent::new(tenant.clone(), name).to_string_key(),
         TaggableObject::Database { db_id } => DatabaseId::new(*db_id).to_string_key(),
         TaggableObject::Table { table_id } => TableId::new(*table_id).to_string_key(),
+        TaggableObject::UDF { name } => UdfIdent::new(tenant.clone(), name).to_string_key(),
+        TaggableObject::Procedure { name, args } => {
+            ProcedureNameIdent::new(tenant.clone(), ProcedureIdentity::new(name, args))
+                .to_string_key()
+        }
     }
 }
 
@@ -402,8 +410,11 @@ where
                         }
                     }
                 }
-                // Stage and Connection use name-based keys, no soft delete concept
-                TaggableObject::Stage { .. } | TaggableObject::Connection { .. } => {}
+                // Stage, Connection, UDF, and Procedure use name-based keys, no soft delete concept
+                TaggableObject::Stage { .. }
+                | TaggableObject::Connection { .. }
+                | TaggableObject::UDF { .. }
+                | TaggableObject::Procedure { .. } => {}
             }
 
             debug!(req :? =(&req); "set_object_tags retry due to concurrent modification");

src/meta/app/src/schema/tag/error.rs

@@ -129,6 +129,8 @@ impl From<TagMetaError> for ErrorCode {
                 TaggableObject::Stage { .. } => ErrorCode::UnknownStage(s),
                 TaggableObject::Database { .. } => ErrorCode::UnknownDatabase(s),
                 TaggableObject::Table { .. } => ErrorCode::UnknownTable(s),
+                TaggableObject::UDF { .. } => ErrorCode::UnknownFunction(s),
+                TaggableObject::Procedure { .. } => ErrorCode::UnknownProcedure(s),
             },
         }
     }

src/meta/app/src/schema/tag/mod.rs

@@ -101,6 +101,8 @@ pub enum TaggableObject {
     Table { table_id: u64 },
     Stage { name: String },
     Connection { name: String },
+    UDF { name: String },
+    Procedure { name: String, args: String },
 }
 
 impl std::fmt::Display for TaggableObject {
@@ -110,6 +112,10 @@ impl std::fmt::Display for TaggableObject {
             TaggableObject::Table { table_id } => write!(f, "table(id={})", table_id),
             TaggableObject::Stage { name } => write!(f, "stage({})", name),
             TaggableObject::Connection { name } => write!(f, "connection({})", name),
+            TaggableObject::UDF { name } => write!(f, "udf({})", name),
+            TaggableObject::Procedure { name, args } => {
+                write!(f, "procedure({}({}))", name, args)
+            }
         }
     }
 }
@@ -121,6 +127,8 @@ impl TaggableObject {
             TaggableObject::Table { .. } => "table",
             TaggableObject::Stage { .. } => "stage",
             TaggableObject::Connection { .. } => "connection",
+            TaggableObject::UDF { .. } => "udf",
+            TaggableObject::Procedure { .. } => "procedure",
         }
     }
 
@@ -131,6 +139,8 @@ impl TaggableObject {
             TaggableObject::Table { table_id } => b.push_u64(*table_id),
             TaggableObject::Stage { name } => b.push_str(name),
             TaggableObject::Connection { name } => b.push_str(name),
+            TaggableObject::UDF { name } => b.push_str(name),
+            TaggableObject::Procedure { name, args } => b.push_str(name).push_str(args),
         }
     }
 
@@ -149,9 +159,16 @@ impl TaggableObject {
             "connection" => Ok(TaggableObject::Connection {
                 name: parser.next_str()?,
             }),
+            "udf" => Ok(TaggableObject::UDF {
+                name: parser.next_str()?,
+            }),
+            "procedure" => Ok(TaggableObject::Procedure {
+                name: parser.next_str()?,
+                args: parser.next_str()?,
+            }),
             _ => Err(KeyError::InvalidSegment {
                 i: parser.index(),
-                expect: "database|table|stage|connection".to_string(),
+                expect: "database|table|stage|connection|udf|procedure".to_string(),
                 got: type_str.to_string(),
             }),
         }

src/query/ast/src/ast/statements/tag.rs

@@ -21,6 +21,7 @@ use derive_visitor::DriveMut;
 use crate::ast::CreateOption;
 use crate::ast::Identifier;
 use crate::ast::Literal;
+use crate::ast::TypeName;
 use crate::ast::quote::QuotedString;
 use crate::ast::statements::show::ShowLimit;
 use crate::ast::write_comma_separated_list;
@@ -146,6 +147,21 @@ pub enum AlterObjectTagTarget {
         if_exists: bool,
         connection_name: Identifier,
     },
+    View {
+        if_exists: bool,
+        catalog: Option<Identifier>,
+        database: Option<Identifier>,
+        view: Identifier,
+    },
+    Function {
+        if_exists: bool,
+        udf_name: Identifier,
+    },
+    Procedure {
+        if_exists: bool,
+        name: Identifier,
+        arg_types: Vec<TypeName>,
+    },
 }
 
 impl Display for AlterObjectTagTarget {
@@ -197,6 +213,44 @@ impl Display for AlterObjectTagTarget {
                 }
                 write!(f, "{connection_name}")?;
             }
+            AlterObjectTagTarget::View {
+                if_exists,
+                catalog,
+                database,
+                view,
+            } => {
+                write!(f, "VIEW ")?;
+                if *if_exists {
+                    write!(f, "IF EXISTS ")?;
+                }
+                write_dot_separated_list(
+                    f,
+                    catalog.iter().chain(database.iter()).chain(Some(view)),
+                )?;
+            }
+            AlterObjectTagTarget::Function {
+                if_exists,
+                udf_name,
+            } => {
+                write!(f, "FUNCTION ")?;
+                if *if_exists {
+                    write!(f, "IF EXISTS ")?;
+                }
+                write!(f, "{udf_name}")?;
+            }
+            AlterObjectTagTarget::Procedure {
+                if_exists,
+                name,
+                arg_types,
+            } => {
+                write!(f, "PROCEDURE ")?;
+                if *if_exists {
+                    write!(f, "IF EXISTS ")?;
+                }
+                write!(f, "{name}(")?;
+                write_comma_separated_list(f, arg_types)?;
+                write!(f, ")")?;
+            }
         }
         Ok(())
     }

src/query/ast/src/parser/parser.rs

@@ -25,6 +25,7 @@ use crate::ast::ExplainKind;
 use crate::ast::Expr;
 use crate::ast::Identifier;
 use crate::ast::Literal;
+use crate::ast::ProcedureIdentity;
 use crate::ast::SelectTarget;
 use crate::ast::Statement;
 use crate::ast::StatementWithFormat;
@@ -44,6 +45,7 @@ use crate::parser::input::Dialect;
 use crate::parser::input::Input;
 use crate::parser::input::ParseMode;
 use crate::parser::statement::insert_stmt;
+use crate::parser::statement::procedure_type_name;
 use crate::parser::statement::replace_stmt;
 use crate::parser::statement::statement;
 use crate::parser::token::Token;
@@ -84,6 +86,29 @@ pub fn parse_database_ref(sql: &str, dialect: Dialect) -> Result<DatabaseRef> {
     run_parser(&tokens, dialect, ParseMode::Default, false, database_ref)
 }
 
+/// Parse a procedure reference string like "my_proc(INT, STRING)" or "my_proc()".
+/// Returns a `ProcedureIdentity` with the procedure name and argument types.
+pub fn parse_procedure_ref(sql: &str, dialect: Dialect) -> Result<ProcedureIdentity> {
+    let tokens = tokenize_sql(sql)?;
+    run_parser(&tokens, dialect, ParseMode::Default, false, |i| {
+        nom::combinator::map(
+            nom::sequence::pair(ident, procedure_type_name),
+            |(name, args_type): (Identifier, _)| ProcedureIdentity {
+                name: name.to_string(),
+                args_type,
+            },
+        )
+        .parse(i)
+    })
+}
+
+/// Parse a UDF name string (a single identifier).
+/// Rejects trailing tokens to avoid silently ignoring malformed input.
+pub fn parse_udf_ref(sql: &str, dialect: Dialect) -> Result<Identifier> {
+    let tokens = tokenize_sql(sql)?;
+    run_parser(&tokens, dialect, ParseMode::Default, false, ident)
+}
+
 pub fn parse_comma_separated_exprs(tokens: &[Token], dialect: Dialect) -> Result<Vec<Expr>> {
     run_parser(tokens, dialect, ParseMode::Default, true, |i| {
         comma_separated_list0(expr)(i)

src/query/ast/src/parser/statement.rs

@@ -56,7 +56,7 @@ pub enum CreateDatabaseOption {
     Options(Vec<SQLProperty>),
 }
 
-fn procedure_type_name(i: Input) -> IResult<Vec<TypeName>> {
+pub fn procedure_type_name(i: Input) -> IResult<Vec<TypeName>> {
     let procedure_type_names = map(
         rule! {
             "(" ~ #comma_separated_list1(type_name) ~ ")"
@@ -4431,6 +4431,36 @@ fn alter_object_tag_target(i: Input) -> IResult<AlterObjectTagTarget> {
                 connection_name,
             },
         ),
+        map(
+            rule! {
+                VIEW ~ ( IF ~ ^EXISTS )? ~ #dot_separated_idents_1_to_3
+            },
+            |(_, opt_if_exists, (catalog, database, view))| AlterObjectTagTarget::View {
+                if_exists: opt_if_exists.is_some(),
+                catalog,
+                database,
+                view,
+            },
+        ),
+        map(
+            rule! {
+                FUNCTION ~ ( IF ~ ^EXISTS )? ~ #ident
+            },
+            |(_, opt_if_exists, udf_name)| AlterObjectTagTarget::Function {
+                if_exists: opt_if_exists.is_some(),
+                udf_name,
+            },
+        ),
+        map(
+            rule! {
+                PROCEDURE ~ ( IF ~ ^EXISTS )? ~ #ident ~ #procedure_type_name
+            },
+            |(_, opt_if_exists, name, arg_types)| AlterObjectTagTarget::Procedure {
+                if_exists: opt_if_exists.is_some(),
+                name,
+                arg_types,
+            },
+        ),
     ))
     .parse(i)
 }

src/query/ast/tests/it/parser.rs

@@ -342,6 +342,12 @@ SELECT * from s;"#,
         r#"ALTER DATABASE IF EXISTS ctl.c RENAME TO a;"#,
         r#"ALTER DATABASE c RENAME TO a;"#,
         r#"ALTER DATABASE c set tag tag1='a';"#,
+        r#"ALTER VIEW v SET TAG tag1 = 'val1';"#,
+        r#"ALTER VIEW IF EXISTS db.v UNSET TAG tag1;"#,
+        r#"ALTER FUNCTION my_udf SET TAG tag1 = 'val1';"#,
+        r#"ALTER FUNCTION IF EXISTS my_udf UNSET TAG tag1;"#,
+        r#"ALTER PROCEDURE my_proc(INT, STRING) SET TAG tag1 = 'val1';"#,
+        r#"ALTER PROCEDURE IF EXISTS my_proc() UNSET TAG tag1;"#,
         r#"ALTER DATABASE ctl.c RENAME TO a;"#,
         r#"ALTER DATABASE ctl.c refresh cache;"#,
         r#"VACUUM TABLE t;"#,