Skip to content

Commit 4beea9a

Browse files
LucaCappelletti94iffyio
LucaCappelletti94
and
iffyio
authored
Support PostgreSQL C Functions with Multiple AS Parameters (#2095)
Co-authored-by: Ifeanyi Ubah <[email protected]>
1 parent e380494 commit 4beea9a

File tree

5 files changed

+147
-50
lines changed

5 files changed

+147
-50
lines changed

‎src/ast/ddl.rs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3099,8 +3099,12 @@ impl fmt::Display for CreateFunction {
30993099
if let Some(remote_connection) = &self.remote_connection {
31003100
write!(f, " REMOTE WITH CONNECTION {remote_connection}")?;
31013101
}
3102-
if let Some(CreateFunctionBody::AsBeforeOptions(function_body)) = &self.function_body {
3103-
write!(f, " AS {function_body}")?;
3102+
if let Some(CreateFunctionBody::AsBeforeOptions { body, link_symbol }) = &self.function_body
3103+
{
3104+
write!(f, " AS {body}")?;
3105+
if let Some(link_symbol) = link_symbol {
3106+
write!(f, ", {link_symbol}")?;
3107+
}
31043108
}
31053109
if let Some(CreateFunctionBody::Return(function_body)) = &self.function_body {
31063110
write!(f, " RETURN {function_body}")?;

‎src/ast/mod.rs‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9108,7 +9108,20 @@ pub enum CreateFunctionBody {
91089108
/// ```
91099109
///
91109110
/// [BigQuery]: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#syntax_11
9111-
AsBeforeOptions(Expr),
9111+
/// [PostgreSQL]: https://www.postgresql.org/docs/current/sql-createfunction.html
9112+
AsBeforeOptions {
9113+
/// The primary expression.
9114+
body: Expr,
9115+
/// Link symbol if the primary expression contains the name of shared library file.
9116+
///
9117+
/// Example:
9118+
/// ```sql
9119+
/// CREATE FUNCTION cas_in(input cstring) RETURNS cas
9120+
/// AS 'MODULE_PATHNAME', 'cas_in_wrapper'
9121+
/// ```
9122+
/// [PostgreSQL]: https://www.postgresql.org/docs/current/sql-createfunction.html
9123+
link_symbol: Option<Expr>,
9124+
},
91129125
/// A function body expression using the 'AS' keyword and shows up
91139126
/// after any `OPTIONS` clause.
91149127
///

‎src/parser/mod.rs‎

Lines changed: 30 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -5204,9 +5204,7 @@ impl<'a> Parser<'a> {
52045204
}
52055205
if self.parse_keyword(Keyword::AS) {
52065206
ensure_not_set(&body.function_body, "AS")?;
5207-
body.function_body = Some(CreateFunctionBody::AsBeforeOptions(
5208-
self.parse_create_function_body_string()?,
5209-
));
5207+
body.function_body = Some(self.parse_create_function_body_string()?);
52105208
} else if self.parse_keyword(Keyword::LANGUAGE) {
52115209
ensure_not_set(&body.language, "LANGUAGE")?;
52125210
body.language = Some(self.parse_identifier()?);
@@ -5298,15 +5296,15 @@ impl<'a> Parser<'a> {
52985296
let name = self.parse_object_name(false)?;
52995297
self.expect_keyword_is(Keyword::AS)?;
53005298

5301-
let as_ = self.parse_create_function_body_string()?;
5299+
let body = self.parse_create_function_body_string()?;
53025300
let using = self.parse_optional_create_function_using()?;
53035301

53045302
Ok(Statement::CreateFunction(CreateFunction {
53055303
or_alter: false,
53065304
or_replace,
53075305
temporary,
53085306
name,
5309-
function_body: Some(CreateFunctionBody::AsBeforeOptions(as_)),
5307+
function_body: Some(body),
53105308
using,
53115309
if_not_exists: false,
53125310
args: None,
@@ -5368,7 +5366,10 @@ impl<'a> Parser<'a> {
53685366
let expr = self.parse_expr()?;
53695367
if options.is_none() {
53705368
options = self.maybe_parse_options(Keyword::OPTIONS)?;
5371-
Some(CreateFunctionBody::AsBeforeOptions(expr))
5369+
Some(CreateFunctionBody::AsBeforeOptions {
5370+
body: expr,
5371+
link_symbol: None,
5372+
})
53725373
} else {
53735374
Some(CreateFunctionBody::AsAfterOptions(expr))
53745375
}
@@ -10574,19 +10575,30 @@ impl<'a> Parser<'a> {
1057410575

1057510576
/// Parse the body of a `CREATE FUNCTION` specified as a string.
1057610577
/// e.g. `CREATE FUNCTION ... AS $$ body $$`.
10577-
fn parse_create_function_body_string(&mut self) -> Result<Expr, ParserError> {
10578-
let peek_token = self.peek_token();
10579-
let span = peek_token.span;
10580-
match peek_token.token {
10581-
Token::DollarQuotedString(s) if dialect_of!(self is PostgreSqlDialect | GenericDialect) =>
10582-
{
10583-
self.next_token();
10584-
Ok(Expr::Value(Value::DollarQuotedString(s).with_span(span)))
10578+
fn parse_create_function_body_string(&mut self) -> Result<CreateFunctionBody, ParserError> {
10579+
let parse_string_expr = |parser: &mut Parser| -> Result<Expr, ParserError> {
10580+
let peek_token = parser.peek_token();
10581+
let span = peek_token.span;
10582+
match peek_token.token {
10583+
Token::DollarQuotedString(s) if dialect_of!(parser is PostgreSqlDialect | GenericDialect) =>
10584+
{
10585+
parser.next_token();
10586+
Ok(Expr::Value(Value::DollarQuotedString(s).with_span(span)))
10587+
}
10588+
_ => Ok(Expr::Value(
10589+
Value::SingleQuotedString(parser.parse_literal_string()?).with_span(span),
10590+
)),
1058510591
}
10586-
_ => Ok(Expr::Value(
10587-
Value::SingleQuotedString(self.parse_literal_string()?).with_span(span),
10588-
)),
10589-
}
10592+
};
10593+
10594+
Ok(CreateFunctionBody::AsBeforeOptions {
10595+
body: parse_string_expr(self)?,
10596+
link_symbol: if self.consume_token(&Token::Comma) {
10597+
Some(parse_string_expr(self)?)
10598+
} else {
10599+
None
10600+
},
10601+
})
1059010602
}
1059110603

1059210604
/// Parse a literal string

‎tests/sqlparser_hive.rs‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -408,10 +408,13 @@ fn parse_create_function() {
408408
assert_eq!(name.to_string(), "mydb.myfunc");
409409
assert_eq!(
410410
function_body,
411-
Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
412-
(Value::SingleQuotedString("org.random.class.Name".to_string()))
413-
.with_empty_span()
414-
)))
411+
Some(CreateFunctionBody::AsBeforeOptions {
412+
body: Expr::Value(
413+
(Value::SingleQuotedString("org.random.class.Name".to_string()))
414+
.with_empty_span()
415+
),
416+
link_symbol: None,
417+
})
415418
);
416419
assert_eq!(
417420
using,

‎tests/sqlparser_postgres.rs‎

Lines changed: 90 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -4260,9 +4260,12 @@ $$"#;
42604260
behavior: None,
42614261
called_on_null: None,
42624262
parallel: None,
4263-
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
4264-
(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()
4265-
))),
4263+
function_body: Some(CreateFunctionBody::AsBeforeOptions {
4264+
body: Expr::Value(
4265+
(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()
4266+
),
4267+
link_symbol: None,
4268+
}),
42664269
if_not_exists: false,
42674270
using: None,
42684271
determinism_specifier: None,
@@ -4298,9 +4301,12 @@ $$"#;
42984301
behavior: None,
42994302
called_on_null: None,
43004303
parallel: None,
4301-
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
4302-
(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()
4303-
))),
4304+
function_body: Some(CreateFunctionBody::AsBeforeOptions {
4305+
body: Expr::Value(
4306+
(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()
4307+
),
4308+
link_symbol: None,
4309+
}),
43044310
if_not_exists: false,
43054311
using: None,
43064312
determinism_specifier: None,
@@ -4340,9 +4346,12 @@ $$"#;
43404346
behavior: None,
43414347
called_on_null: None,
43424348
parallel: None,
4343-
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
4344-
(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()
4345-
))),
4349+
function_body: Some(CreateFunctionBody::AsBeforeOptions {
4350+
body: Expr::Value(
4351+
(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()
4352+
),
4353+
link_symbol: None,
4354+
}),
43464355
if_not_exists: false,
43474356
using: None,
43484357
determinism_specifier: None,
@@ -4382,9 +4391,12 @@ $$"#;
43824391
behavior: None,
43834392
called_on_null: None,
43844393
parallel: None,
4385-
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
4386-
(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()
4387-
))),
4394+
function_body: Some(CreateFunctionBody::AsBeforeOptions {
4395+
body: Expr::Value(
4396+
(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()
4397+
),
4398+
link_symbol: None,
4399+
}),
43884400
if_not_exists: false,
43894401
using: None,
43904402
determinism_specifier: None,
@@ -4417,13 +4429,16 @@ $$"#;
44174429
behavior: None,
44184430
called_on_null: None,
44194431
parallel: None,
4420-
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
4421-
(Value::DollarQuotedString(DollarQuotedString {
4422-
value: "\n BEGIN\n RETURN TRUE;\n END;\n ".to_owned(),
4423-
tag: None
4424-
}))
4425-
.with_empty_span()
4426-
))),
4432+
function_body: Some(CreateFunctionBody::AsBeforeOptions {
4433+
body: Expr::Value(
4434+
(Value::DollarQuotedString(DollarQuotedString {
4435+
value: "\n BEGIN\n RETURN TRUE;\n END;\n ".to_owned(),
4436+
tag: None
4437+
}))
4438+
.with_empty_span()
4439+
),
4440+
link_symbol: None,
4441+
}),
44274442
if_not_exists: false,
44284443
using: None,
44294444
determinism_specifier: None,
@@ -4455,9 +4470,12 @@ fn parse_create_function() {
44554470
behavior: Some(FunctionBehavior::Immutable),
44564471
called_on_null: Some(FunctionCalledOnNull::Strict),
44574472
parallel: Some(FunctionParallel::Safe),
4458-
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
4459-
(Value::SingleQuotedString("select $1 + $2;".into())).with_empty_span()
4460-
))),
4473+
function_body: Some(CreateFunctionBody::AsBeforeOptions {
4474+
body: Expr::Value(
4475+
(Value::SingleQuotedString("select $1 + $2;".into())).with_empty_span()
4476+
),
4477+
link_symbol: None,
4478+
}),
44614479
if_not_exists: false,
44624480
using: None,
44634481
determinism_specifier: None,
@@ -4488,6 +4506,52 @@ fn parse_incorrect_create_function_parallel() {
44884506
assert!(pg().parse_sql_statements(sql).is_err());
44894507
}
44904508

4509+
#[test]
4510+
fn parse_create_function_c_with_module_pathname() {
4511+
let sql = "CREATE FUNCTION cas_in(input cstring) RETURNS cas LANGUAGE c IMMUTABLE PARALLEL SAFE AS 'MODULE_PATHNAME', 'cas_in_wrapper'";
4512+
assert_eq!(
4513+
pg_and_generic().verified_stmt(sql),
4514+
Statement::CreateFunction(CreateFunction {
4515+
or_alter: false,
4516+
or_replace: false,
4517+
temporary: false,
4518+
name: ObjectName::from(vec![Ident::new("cas_in")]),
4519+
args: Some(vec![OperateFunctionArg::with_name(
4520+
"input",
4521+
DataType::Custom(ObjectName::from(vec![Ident::new("cstring")]), vec![]),
4522+
),]),
4523+
return_type: Some(DataType::Custom(
4524+
ObjectName::from(vec![Ident::new("cas")]),
4525+
vec![]
4526+
)),
4527+
language: Some("c".into()),
4528+
behavior: Some(FunctionBehavior::Immutable),
4529+
called_on_null: None,
4530+
parallel: Some(FunctionParallel::Safe),
4531+
function_body: Some(CreateFunctionBody::AsBeforeOptions {
4532+
body: Expr::Value(
4533+
(Value::SingleQuotedString("MODULE_PATHNAME".into())).with_empty_span()
4534+
),
4535+
link_symbol: Some(Expr::Value(
4536+
(Value::SingleQuotedString("cas_in_wrapper".into())).with_empty_span()
4537+
)),
4538+
}),
4539+
if_not_exists: false,
4540+
using: None,
4541+
determinism_specifier: None,
4542+
options: None,
4543+
remote_connection: None,
4544+
})
4545+
);
4546+
4547+
// Test that attribute order flexibility works (IMMUTABLE before LANGUAGE)
4548+
let sql_alt_order = "CREATE FUNCTION cas_in(input cstring) RETURNS cas IMMUTABLE PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', 'cas_in_wrapper'";
4549+
pg_and_generic().one_statement_parses_to(
4550+
sql_alt_order,
4551+
"CREATE FUNCTION cas_in(input cstring) RETURNS cas LANGUAGE c IMMUTABLE PARALLEL SAFE AS 'MODULE_PATHNAME', 'cas_in_wrapper'"
4552+
);
4553+
}
4554+
44914555
#[test]
44924556
fn parse_drop_function() {
44934557
let sql = "DROP FUNCTION IF EXISTS test_func";
@@ -6070,8 +6134,8 @@ fn parse_trigger_related_functions() {
60706134
args: Some(vec![]),
60716135
return_type: Some(DataType::Trigger),
60726136
function_body: Some(
6073-
CreateFunctionBody::AsBeforeOptions(
6074-
Expr::Value((
6137+
CreateFunctionBody::AsBeforeOptions {
6138+
body: Expr::Value((
60756139
Value::DollarQuotedString(
60766140
DollarQuotedString {
60776141
value: "\n BEGIN\n -- Check that empname and salary are given\n IF NEW.empname IS NULL THEN\n RAISE EXCEPTION 'empname cannot be null';\n END IF;\n IF NEW.salary IS NULL THEN\n RAISE EXCEPTION '% cannot have null salary', NEW.empname;\n END IF;\n\n -- Who works for us when they must pay for it?\n IF NEW.salary < 0 THEN\n RAISE EXCEPTION '% cannot have a negative salary', NEW.empname;\n END IF;\n\n -- Remember who changed the payroll when\n NEW.last_date := current_timestamp;\n NEW.last_user := current_user;\n RETURN NEW;\n END;\n ".to_owned(),
@@ -6081,7 +6145,8 @@ fn parse_trigger_related_functions() {
60816145
},
60826146
)
60836147
).with_empty_span()),
6084-
),
6148+
link_symbol: None,
6149+
},
60856150
),
60866151
behavior: None,
60876152
called_on_null: None,

0 commit comments

Comments
 (0)