Skip to content

Commit 690091c

Browse files
fmguerreirofmguerreiro
committed
Add SECURITY DEFINER/INVOKER parsing for PostgreSQL functions.
1 parent 1c00c8a commit 690091c

6 files changed

Lines changed: 106 additions & 1 deletion

File tree

‎src/ast/ddl.rs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ use crate::ast::{
4444
ArgMode, AttachedToken, CommentDef, ConditionalStatements, CreateFunctionBody,
4545
CreateFunctionUsing, CreateTableLikeKind, CreateTableOptions, CreateViewParams, DataType, Expr,
4646
FileFormat, FunctionBehavior, FunctionCalledOnNull, FunctionDesc, FunctionDeterminismSpecifier,
47-
FunctionParallel, HiveDistributionStyle, HiveFormat, HiveIOFormat, HiveRowFormat,
47+
FunctionParallel, FunctionSecurity, HiveDistributionStyle, HiveFormat, HiveIOFormat, HiveRowFormat,
4848
HiveSetLocation, Ident, InitializeKind, MySQLColumnPosition, ObjectName, OnCommit,
4949
OneOrManyWithParens, OperateFunctionArg, OrderByExpr, ProjectionSelect, Query, RefreshModeKind,
5050
RowAccessPolicy, SequenceOptions, Spanned, SqlOption, StorageSerializationPolicy, TableVersion,
@@ -3226,6 +3226,10 @@ pub struct CreateFunction {
32263226
///
32273227
/// [PostgreSQL](https://www.postgresql.org/docs/current/sql-createfunction.html)
32283228
pub parallel: Option<FunctionParallel>,
3229+
/// SECURITY DEFINER | SECURITY INVOKER
3230+
///
3231+
/// [PostgreSQL](https://www.postgresql.org/docs/current/sql-createfunction.html)
3232+
pub security: Option<FunctionSecurity>,
32293233
/// USING ... (Hive only)
32303234
pub using: Option<CreateFunctionUsing>,
32313235
/// Language used in a UDF definition.
@@ -3292,6 +3296,9 @@ impl fmt::Display for CreateFunction {
32923296
if let Some(parallel) = &self.parallel {
32933297
write!(f, " {parallel}")?;
32943298
}
3299+
if let Some(security) = &self.security {
3300+
write!(f, " {security}")?;
3301+
}
32953302
if let Some(remote_connection) = &self.remote_connection {
32963303
write!(f, " REMOTE WITH CONNECTION {remote_connection}")?;
32973304
}

‎src/ast/mod.rs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8822,6 +8822,26 @@ impl fmt::Display for FunctionParallel {
88228822
}
88238823
}
88248824

8825+
/// SECURITY DEFINER | SECURITY INVOKER
8826+
///
8827+
/// [PostgreSQL](https://www.postgresql.org/docs/current/sql-createfunction.html)
8828+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
8829+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
8830+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
8831+
pub enum FunctionSecurity {
8832+
Definer,
8833+
Invoker,
8834+
}
8835+
8836+
impl fmt::Display for FunctionSecurity {
8837+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
8838+
match self {
8839+
FunctionSecurity::Definer => write!(f, "SECURITY DEFINER"),
8840+
FunctionSecurity::Invoker => write!(f, "SECURITY INVOKER"),
8841+
}
8842+
}
8843+
}
8844+
88258845
/// [BigQuery] Determinism specifier used in a UDF definition.
88268846
///
88278847
/// [BigQuery]: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#syntax_11

‎src/parser/mod.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5260,6 +5260,7 @@ impl<'a> Parser<'a> {
52605260
function_body: Option<CreateFunctionBody>,
52615261
called_on_null: Option<FunctionCalledOnNull>,
52625262
parallel: Option<FunctionParallel>,
5263+
security: Option<FunctionSecurity>,
52635264
options: Vec<SqlOption>,
52645265
}
52655266
let mut body = Body::default();
@@ -5339,6 +5340,15 @@ impl<'a> Parser<'a> {
53395340
first_value
53405341
};
53415342
body.options.push(SqlOption::KeyValue { key: name, value });
5343+
} else if self.parse_keyword(Keyword::SECURITY) {
5344+
ensure_not_set(&body.security, "SECURITY DEFINER | SECURITY INVOKER")?;
5345+
if self.parse_keyword(Keyword::DEFINER) {
5346+
body.security = Some(FunctionSecurity::Definer);
5347+
} else if self.parse_keyword(Keyword::INVOKER) {
5348+
body.security = Some(FunctionSecurity::Invoker);
5349+
} else {
5350+
return self.expected("DEFINER or INVOKER", self.peek_token());
5351+
}
53425352
} else if self.parse_keyword(Keyword::RETURN) {
53435353
ensure_not_set(&body.function_body, "RETURN")?;
53445354
body.function_body = Some(CreateFunctionBody::Return(self.parse_expr()?));
@@ -5357,6 +5367,7 @@ impl<'a> Parser<'a> {
53575367
behavior: body.behavior,
53585368
called_on_null: body.called_on_null,
53595369
parallel: body.parallel,
5370+
security: body.security,
53605371
language: body.language,
53615372
function_body: body.function_body,
53625373
if_not_exists: false,
@@ -5398,6 +5409,7 @@ impl<'a> Parser<'a> {
53985409
behavior: None,
53995410
called_on_null: None,
54005411
parallel: None,
5412+
security: None,
54015413
language: None,
54025414
determinism_specifier: None,
54035415
options: None,
@@ -5480,6 +5492,7 @@ impl<'a> Parser<'a> {
54805492
behavior: None,
54815493
called_on_null: None,
54825494
parallel: None,
5495+
security: None,
54835496
}))
54845497
}
54855498

@@ -5569,6 +5582,7 @@ impl<'a> Parser<'a> {
55695582
behavior: None,
55705583
called_on_null: None,
55715584
parallel: None,
5585+
security: None,
55725586
}))
55735587
}
55745588

‎tests/sqlparser_bigquery.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2294,6 +2294,7 @@ fn test_bigquery_create_function() {
22942294
remote_connection: None,
22952295
called_on_null: None,
22962296
parallel: None,
2297+
security: None,
22972298
})
22982299
);
22992300

‎tests/sqlparser_mssql.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,7 @@ fn parse_create_function() {
266266
behavior: None,
267267
called_on_null: None,
268268
parallel: None,
269+
security: None,
269270
using: None,
270271
language: None,
271272
determinism_specifier: None,
@@ -439,6 +440,7 @@ fn parse_create_function_parameter_default_values() {
439440
behavior: None,
440441
called_on_null: None,
441442
parallel: None,
443+
security: None,
442444
using: None,
443445
language: None,
444446
determinism_specifier: None,

‎tests/sqlparser_postgres.rs‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4294,6 +4294,7 @@ $$"#;
42944294
behavior: None,
42954295
called_on_null: None,
42964296
parallel: None,
4297+
security: None,
42974298
function_body: Some(CreateFunctionBody::AsBeforeOptions {
42984299
body: Expr::Value(
42994300
(Value::DollarQuotedString(DollarQuotedString {value: "\nBEGIN\n IF str1 <> str2 THEN\n RETURN TRUE;\n ELSE\n RETURN FALSE;\n END IF;\nEND;\n".to_owned(), tag: None})).with_empty_span()
@@ -4335,6 +4336,7 @@ $$"#;
43354336
behavior: None,
43364337
called_on_null: None,
43374338
parallel: None,
4339+
security: None,
43384340
function_body: Some(CreateFunctionBody::AsBeforeOptions {
43394341
body: Expr::Value(
43404342
(Value::DollarQuotedString(DollarQuotedString {value: "\nBEGIN\n IF int1 <> 0 THEN\n RETURN TRUE;\n ELSE\n RETURN FALSE;\n END IF;\nEND;\n".to_owned(), tag: None})).with_empty_span()
@@ -4380,6 +4382,7 @@ $$"#;
43804382
behavior: None,
43814383
called_on_null: None,
43824384
parallel: None,
4385+
security: None,
43834386
function_body: Some(CreateFunctionBody::AsBeforeOptions {
43844387
body: Expr::Value(
43854388
(Value::DollarQuotedString(DollarQuotedString {value: "\nBEGIN\n IF a <> b THEN\n RETURN TRUE;\n ELSE\n RETURN FALSE;\n END IF;\nEND;\n".to_owned(), tag: None})).with_empty_span()
@@ -4425,6 +4428,7 @@ $$"#;
44254428
behavior: None,
44264429
called_on_null: None,
44274430
parallel: None,
4431+
security: None,
44284432
function_body: Some(CreateFunctionBody::AsBeforeOptions {
44294433
body: Expr::Value(
44304434
(Value::DollarQuotedString(DollarQuotedString {value: "\nBEGIN\n IF int1 <> int2 THEN\n RETURN TRUE;\n ELSE\n RETURN FALSE;\n END IF;\nEND;\n".to_owned(), tag: None})).with_empty_span()
@@ -4463,6 +4467,7 @@ $$"#;
44634467
behavior: None,
44644468
called_on_null: None,
44654469
parallel: None,
4470+
security: None,
44664471
function_body: Some(CreateFunctionBody::AsBeforeOptions {
44674472
body: Expr::Value(
44684473
(Value::DollarQuotedString(DollarQuotedString {
@@ -4504,6 +4509,7 @@ fn parse_create_function() {
45044509
behavior: Some(FunctionBehavior::Immutable),
45054510
called_on_null: Some(FunctionCalledOnNull::Strict),
45064511
parallel: Some(FunctionParallel::Safe),
4512+
security: None,
45074513
function_body: Some(CreateFunctionBody::AsBeforeOptions {
45084514
body: Expr::Value(
45094515
(Value::SingleQuotedString("select $1 + $2;".into())).with_empty_span()
@@ -4562,6 +4568,7 @@ fn parse_create_function_c_with_module_pathname() {
45624568
behavior: Some(FunctionBehavior::Immutable),
45634569
called_on_null: None,
45644570
parallel: Some(FunctionParallel::Safe),
4571+
security: None,
45654572
function_body: Some(CreateFunctionBody::AsBeforeOptions {
45664573
body: Expr::Value(
45674574
(Value::SingleQuotedString("MODULE_PATHNAME".into())).with_empty_span()
@@ -4674,6 +4681,59 @@ AS $$ BEGIN RETURN event; END; $$"#;
46744681
}
46754682
}
46764683

4684+
#[test]
4685+
fn parse_create_function_with_security_definer() {
4686+
let sql = r#"CREATE FUNCTION public.my_func() RETURNS void LANGUAGE sql SECURITY DEFINER AS $$ SELECT 1 $$"#;
4687+
4688+
let stmt = pg().verified_stmt(sql);
4689+
match stmt {
4690+
Statement::CreateFunction(CreateFunction { name, security, .. }) => {
4691+
assert_eq!(name.to_string(), "public.my_func");
4692+
assert_eq!(security, Some(FunctionSecurity::Definer));
4693+
}
4694+
_ => panic!("Expected CreateFunction"),
4695+
}
4696+
}
4697+
4698+
#[test]
4699+
fn parse_create_function_with_security_invoker() {
4700+
let sql = r#"CREATE FUNCTION public.my_func() RETURNS void LANGUAGE sql SECURITY INVOKER AS $$ SELECT 1 $$"#;
4701+
4702+
let stmt = pg().verified_stmt(sql);
4703+
match stmt {
4704+
Statement::CreateFunction(CreateFunction { name, security, .. }) => {
4705+
assert_eq!(name.to_string(), "public.my_func");
4706+
assert_eq!(security, Some(FunctionSecurity::Invoker));
4707+
}
4708+
_ => panic!("Expected CreateFunction"),
4709+
}
4710+
}
4711+
4712+
#[test]
4713+
fn parse_create_function_with_security_and_other_attributes() {
4714+
pg().one_statement_parses_to(
4715+
r#"CREATE FUNCTION test_func() RETURNS integer LANGUAGE plpgsql IMMUTABLE SECURITY DEFINER AS $$ BEGIN RETURN 42; END; $$"#,
4716+
r#"CREATE FUNCTION test_func() RETURNS INTEGER LANGUAGE plpgsql IMMUTABLE SECURITY DEFINER AS $$ BEGIN RETURN 42; END; $$"#,
4717+
);
4718+
4719+
let stmt = pg().verified_stmt(r#"CREATE FUNCTION test_func() RETURNS INTEGER LANGUAGE plpgsql IMMUTABLE SECURITY DEFINER AS $$ BEGIN RETURN 42; END; $$"#);
4720+
match stmt {
4721+
Statement::CreateFunction(CreateFunction {
4722+
name,
4723+
security,
4724+
behavior,
4725+
language,
4726+
..
4727+
}) => {
4728+
assert_eq!(name.to_string(), "test_func");
4729+
assert_eq!(security, Some(FunctionSecurity::Definer));
4730+
assert_eq!(behavior, Some(FunctionBehavior::Immutable));
4731+
assert_eq!(language.as_ref().map(|i| i.value.as_str()), Some("plpgsql"));
4732+
}
4733+
_ => panic!("Expected CreateFunction"),
4734+
}
4735+
}
4736+
46774737
#[test]
46784738
fn parse_drop_function() {
46794739
let sql = "DROP FUNCTION IF EXISTS test_func";
@@ -6275,6 +6335,7 @@ fn parse_trigger_related_functions() {
62756335
behavior: None,
62766336
called_on_null: None,
62776337
parallel: None,
6338+
security: None,
62786339
using: None,
62796340
language: Some(Ident::new("plpgsql")),
62806341
determinism_specifier: None,

0 commit comments

Comments
 (0)