Skip to content

Commit f489dfc

Browse files
TCeasonTCeason
authored
feat(query): extend tag support to view, udf, and procedure objects (#19447)
1 parent 503eb55 commit f489dfc

File tree

18 files changed

+850
-25
lines changed

18 files changed

+850
-25
lines changed

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

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@ use std::collections::HashSet;
1616

1717
use databend_common_meta_app::KeyWithTenant;
1818
use databend_common_meta_app::id_generator::IdGenerator;
19+
use databend_common_meta_app::principal::ProcedureIdentity;
20+
use databend_common_meta_app::principal::ProcedureNameIdent;
21+
use databend_common_meta_app::principal::UdfIdent;
1922
use databend_common_meta_app::principal::connection_ident::ConnectionIdent;
2023
use databend_common_meta_app::principal::user_stage_ident::StageIdent;
2124
use databend_common_meta_app::schema::CreateTagReply;
@@ -87,6 +90,11 @@ fn get_object_key(tenant: &Tenant, object: &TaggableObject) -> String {
8790
TaggableObject::Stage { name } => StageIdent::new(tenant.clone(), name).to_string_key(),
8891
TaggableObject::Database { db_id } => DatabaseId::new(*db_id).to_string_key(),
8992
TaggableObject::Table { table_id } => TableId::new(*table_id).to_string_key(),
93+
TaggableObject::UDF { name } => UdfIdent::new(tenant.clone(), name).to_string_key(),
94+
TaggableObject::Procedure { name, args } => {
95+
ProcedureNameIdent::new(tenant.clone(), ProcedureIdentity::new(name, args))
96+
.to_string_key()
97+
}
9098
}
9199
}
92100

@@ -402,8 +410,11 @@ where
402410
}
403411
}
404412
}
405-
// Stage and Connection use name-based keys, no soft delete concept
406-
TaggableObject::Stage { .. } | TaggableObject::Connection { .. } => {}
413+
// Stage, Connection, UDF, and Procedure use name-based keys, no soft delete concept
414+
TaggableObject::Stage { .. }
415+
| TaggableObject::Connection { .. }
416+
| TaggableObject::UDF { .. }
417+
| TaggableObject::Procedure { .. } => {}
407418
}
408419

409420
debug!(req :? =(&req); "set_object_tags retry due to concurrent modification");

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,8 @@ impl From<TagMetaError> for ErrorCode {
129129
TaggableObject::Stage { .. } => ErrorCode::UnknownStage(s),
130130
TaggableObject::Database { .. } => ErrorCode::UnknownDatabase(s),
131131
TaggableObject::Table { .. } => ErrorCode::UnknownTable(s),
132+
TaggableObject::UDF { .. } => ErrorCode::UnknownFunction(s),
133+
TaggableObject::Procedure { .. } => ErrorCode::UnknownProcedure(s),
132134
},
133135
}
134136
}

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,8 @@ pub enum TaggableObject {
101101
Table { table_id: u64 },
102102
Stage { name: String },
103103
Connection { name: String },
104+
UDF { name: String },
105+
Procedure { name: String, args: String },
104106
}
105107

106108
impl std::fmt::Display for TaggableObject {
@@ -110,6 +112,10 @@ impl std::fmt::Display for TaggableObject {
110112
TaggableObject::Table { table_id } => write!(f, "table(id={})", table_id),
111113
TaggableObject::Stage { name } => write!(f, "stage({})", name),
112114
TaggableObject::Connection { name } => write!(f, "connection({})", name),
115+
TaggableObject::UDF { name } => write!(f, "udf({})", name),
116+
TaggableObject::Procedure { name, args } => {
117+
write!(f, "procedure({}({}))", name, args)
118+
}
113119
}
114120
}
115121
}
@@ -121,6 +127,8 @@ impl TaggableObject {
121127
TaggableObject::Table { .. } => "table",
122128
TaggableObject::Stage { .. } => "stage",
123129
TaggableObject::Connection { .. } => "connection",
130+
TaggableObject::UDF { .. } => "udf",
131+
TaggableObject::Procedure { .. } => "procedure",
124132
}
125133
}
126134

@@ -131,6 +139,8 @@ impl TaggableObject {
131139
TaggableObject::Table { table_id } => b.push_u64(*table_id),
132140
TaggableObject::Stage { name } => b.push_str(name),
133141
TaggableObject::Connection { name } => b.push_str(name),
142+
TaggableObject::UDF { name } => b.push_str(name),
143+
TaggableObject::Procedure { name, args } => b.push_str(name).push_str(args),
134144
}
135145
}
136146

@@ -149,9 +159,16 @@ impl TaggableObject {
149159
"connection" => Ok(TaggableObject::Connection {
150160
name: parser.next_str()?,
151161
}),
162+
"udf" => Ok(TaggableObject::UDF {
163+
name: parser.next_str()?,
164+
}),
165+
"procedure" => Ok(TaggableObject::Procedure {
166+
name: parser.next_str()?,
167+
args: parser.next_str()?,
168+
}),
152169
_ => Err(KeyError::InvalidSegment {
153170
i: parser.index(),
154-
expect: "database|table|stage|connection".to_string(),
171+
expect: "database|table|stage|connection|udf|procedure".to_string(),
155172
got: type_str.to_string(),
156173
}),
157174
}

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use derive_visitor::DriveMut;
2121
use crate::ast::CreateOption;
2222
use crate::ast::Identifier;
2323
use crate::ast::Literal;
24+
use crate::ast::TypeName;
2425
use crate::ast::quote::QuotedString;
2526
use crate::ast::statements::show::ShowLimit;
2627
use crate::ast::write_comma_separated_list;
@@ -146,6 +147,21 @@ pub enum AlterObjectTagTarget {
146147
if_exists: bool,
147148
connection_name: Identifier,
148149
},
150+
View {
151+
if_exists: bool,
152+
catalog: Option<Identifier>,
153+
database: Option<Identifier>,
154+
view: Identifier,
155+
},
156+
Function {
157+
if_exists: bool,
158+
udf_name: Identifier,
159+
},
160+
Procedure {
161+
if_exists: bool,
162+
name: Identifier,
163+
arg_types: Vec<TypeName>,
164+
},
149165
}
150166

151167
impl Display for AlterObjectTagTarget {
@@ -197,6 +213,44 @@ impl Display for AlterObjectTagTarget {
197213
}
198214
write!(f, "{connection_name}")?;
199215
}
216+
AlterObjectTagTarget::View {
217+
if_exists,
218+
catalog,
219+
database,
220+
view,
221+
} => {
222+
write!(f, "VIEW ")?;
223+
if *if_exists {
224+
write!(f, "IF EXISTS ")?;
225+
}
226+
write_dot_separated_list(
227+
f,
228+
catalog.iter().chain(database.iter()).chain(Some(view)),
229+
)?;
230+
}
231+
AlterObjectTagTarget::Function {
232+
if_exists,
233+
udf_name,
234+
} => {
235+
write!(f, "FUNCTION ")?;
236+
if *if_exists {
237+
write!(f, "IF EXISTS ")?;
238+
}
239+
write!(f, "{udf_name}")?;
240+
}
241+
AlterObjectTagTarget::Procedure {
242+
if_exists,
243+
name,
244+
arg_types,
245+
} => {
246+
write!(f, "PROCEDURE ")?;
247+
if *if_exists {
248+
write!(f, "IF EXISTS ")?;
249+
}
250+
write!(f, "{name}(")?;
251+
write_comma_separated_list(f, arg_types)?;
252+
write!(f, ")")?;
253+
}
200254
}
201255
Ok(())
202256
}

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use crate::ast::ExplainKind;
2525
use crate::ast::Expr;
2626
use crate::ast::Identifier;
2727
use crate::ast::Literal;
28+
use crate::ast::ProcedureIdentity;
2829
use crate::ast::SelectTarget;
2930
use crate::ast::Statement;
3031
use crate::ast::StatementWithFormat;
@@ -44,6 +45,7 @@ use crate::parser::input::Dialect;
4445
use crate::parser::input::Input;
4546
use crate::parser::input::ParseMode;
4647
use crate::parser::statement::insert_stmt;
48+
use crate::parser::statement::procedure_type_name;
4749
use crate::parser::statement::replace_stmt;
4850
use crate::parser::statement::statement;
4951
use crate::parser::token::Token;
@@ -84,6 +86,29 @@ pub fn parse_database_ref(sql: &str, dialect: Dialect) -> Result<DatabaseRef> {
8486
run_parser(&tokens, dialect, ParseMode::Default, false, database_ref)
8587
}
8688

89+
/// Parse a procedure reference string like "my_proc(INT, STRING)" or "my_proc()".
90+
/// Returns a `ProcedureIdentity` with the procedure name and argument types.
91+
pub fn parse_procedure_ref(sql: &str, dialect: Dialect) -> Result<ProcedureIdentity> {
92+
let tokens = tokenize_sql(sql)?;
93+
run_parser(&tokens, dialect, ParseMode::Default, false, |i| {
94+
nom::combinator::map(
95+
nom::sequence::pair(ident, procedure_type_name),
96+
|(name, args_type): (Identifier, _)| ProcedureIdentity {
97+
name: name.to_string(),
98+
args_type,
99+
},
100+
)
101+
.parse(i)
102+
})
103+
}
104+
105+
/// Parse a UDF name string (a single identifier).
106+
/// Rejects trailing tokens to avoid silently ignoring malformed input.
107+
pub fn parse_udf_ref(sql: &str, dialect: Dialect) -> Result<Identifier> {
108+
let tokens = tokenize_sql(sql)?;
109+
run_parser(&tokens, dialect, ParseMode::Default, false, ident)
110+
}
111+
87112
pub fn parse_comma_separated_exprs(tokens: &[Token], dialect: Dialect) -> Result<Vec<Expr>> {
88113
run_parser(tokens, dialect, ParseMode::Default, true, |i| {
89114
comma_separated_list0(expr)(i)

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

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ pub enum CreateDatabaseOption {
5656
Options(Vec<SQLProperty>),
5757
}
5858

59-
fn procedure_type_name(i: Input) -> IResult<Vec<TypeName>> {
59+
pub fn procedure_type_name(i: Input) -> IResult<Vec<TypeName>> {
6060
let procedure_type_names = map(
6161
rule! {
6262
"(" ~ #comma_separated_list1(type_name) ~ ")"
@@ -4431,6 +4431,36 @@ fn alter_object_tag_target(i: Input) -> IResult<AlterObjectTagTarget> {
44314431
connection_name,
44324432
},
44334433
),
4434+
map(
4435+
rule! {
4436+
VIEW ~ ( IF ~ ^EXISTS )? ~ #dot_separated_idents_1_to_3
4437+
},
4438+
|(_, opt_if_exists, (catalog, database, view))| AlterObjectTagTarget::View {
4439+
if_exists: opt_if_exists.is_some(),
4440+
catalog,
4441+
database,
4442+
view,
4443+
},
4444+
),
4445+
map(
4446+
rule! {
4447+
FUNCTION ~ ( IF ~ ^EXISTS )? ~ #ident
4448+
},
4449+
|(_, opt_if_exists, udf_name)| AlterObjectTagTarget::Function {
4450+
if_exists: opt_if_exists.is_some(),
4451+
udf_name,
4452+
},
4453+
),
4454+
map(
4455+
rule! {
4456+
PROCEDURE ~ ( IF ~ ^EXISTS )? ~ #ident ~ #procedure_type_name
4457+
},
4458+
|(_, opt_if_exists, name, arg_types)| AlterObjectTagTarget::Procedure {
4459+
if_exists: opt_if_exists.is_some(),
4460+
name,
4461+
arg_types,
4462+
},
4463+
),
44344464
))
44354465
.parse(i)
44364466
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,12 @@ SELECT * from s;"#,
342342
r#"ALTER DATABASE IF EXISTS ctl.c RENAME TO a;"#,
343343
r#"ALTER DATABASE c RENAME TO a;"#,
344344
r#"ALTER DATABASE c set tag tag1='a';"#,
345+
r#"ALTER VIEW v SET TAG tag1 = 'val1';"#,
346+
r#"ALTER VIEW IF EXISTS db.v UNSET TAG tag1;"#,
347+
r#"ALTER FUNCTION my_udf SET TAG tag1 = 'val1';"#,
348+
r#"ALTER FUNCTION IF EXISTS my_udf UNSET TAG tag1;"#,
349+
r#"ALTER PROCEDURE my_proc(INT, STRING) SET TAG tag1 = 'val1';"#,
350+
r#"ALTER PROCEDURE IF EXISTS my_proc() UNSET TAG tag1;"#,
345351
r#"ALTER DATABASE ctl.c RENAME TO a;"#,
346352
r#"ALTER DATABASE ctl.c refresh cache;"#,
347353
r#"VACUUM TABLE t;"#,

0 commit comments

Comments
 (0)