mirror of
https://github.com/apache/datafusion-sqlparser-rs.git
synced 2025-08-04 06:18:17 +00:00
Handle optional datatypes properly in CREATE FUNCTION statements (#1826)
Co-authored-by: Ifeanyi Ubah <ify1992@yahoo.com>
This commit is contained in:
Luca Cappelletti
GitHub
parent
3f4d5f96ee
commit
05d7ffb1d5
2 changed files with 225 additions and 5 deletions
|
|
@ -5273,12 +5273,21 @@ impl<'a> Parser<'a> {
|
||||||
// parse: [ argname ] argtype
|
// parse: [ argname ] argtype
|
||||||
let mut name = None;
|
let mut name = None;
|
||||||
let mut data_type = self.parse_data_type()?;
|
let mut data_type = self.parse_data_type()?;
|
||||||
if let DataType::Custom(n, _) = &data_type {
|
|
||||||
// the first token is actually a name
|
// To check whether the first token is a name or a type, we need to
|
||||||
match n.0[0].clone() {
|
// peek the next token, which if it is another type keyword, then the
|
||||||
ObjectNamePart::Identifier(ident) => name = Some(ident),
|
// first token is a name and not a type in itself.
|
||||||
|
let data_type_idx = self.get_current_index();
|
||||||
|
if let Some(next_data_type) = self.maybe_parse(|parser| parser.parse_data_type())? {
|
||||||
|
let token = self.token_at(data_type_idx);
|
||||||
|
|
||||||
|
// We ensure that the token is a `Word` token, and not other special tokens.
|
||||||
|
if !matches!(token.token, Token::Word(_)) {
|
||||||
|
return self.expected("a name or type", token.clone());
|
||||||
}
|
}
|
||||||
data_type = self.parse_data_type()?;
|
|
||||||
|
name = Some(Ident::new(token.to_string()));
|
||||||
|
data_type = next_data_type;
|
||||||
}
|
}
|
||||||
|
|
||||||
let default_expr = if self.parse_keyword(Keyword::DEFAULT) || self.consume_token(&Token::Eq)
|
let default_expr = if self.parse_keyword(Keyword::DEFAULT) || self.consume_token(&Token::Eq)
|
||||||
|
|
|
||||||
|
|
@ -21,6 +21,7 @@
|
||||||
|
|
||||||
#[macro_use]
|
#[macro_use]
|
||||||
mod test_utils;
|
mod test_utils;
|
||||||
|
|
||||||
use helpers::attached_token::AttachedToken;
|
use helpers::attached_token::AttachedToken;
|
||||||
use sqlparser::tokenizer::Span;
|
use sqlparser::tokenizer::Span;
|
||||||
use test_utils::*;
|
use test_utils::*;
|
||||||
|
|
@ -4105,6 +4106,216 @@ fn parse_update_in_with_subquery() {
|
||||||
pg_and_generic().verified_stmt(r#"WITH "result" AS (UPDATE "Hero" SET "name" = 'Captain America', "number_of_movies" = "number_of_movies" + 1 WHERE "secret_identity" = 'Sam Wilson' RETURNING "id", "name", "secret_identity", "number_of_movies") SELECT * FROM "result""#);
|
pg_and_generic().verified_stmt(r#"WITH "result" AS (UPDATE "Hero" SET "name" = 'Captain America', "number_of_movies" = "number_of_movies" + 1 WHERE "secret_identity" = 'Sam Wilson' RETURNING "id", "name", "secret_identity", "number_of_movies") SELECT * FROM "result""#);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parser_create_function_with_args() {
|
||||||
|
let sql1 = r#"CREATE OR REPLACE FUNCTION check_strings_different(str1 VARCHAR, str2 VARCHAR) RETURNS BOOLEAN LANGUAGE plpgsql AS $$
|
||||||
|
BEGIN
|
||||||
|
IF str1 <> str2 THEN
|
||||||
|
RETURN TRUE;
|
||||||
|
ELSE
|
||||||
|
RETURN FALSE;
|
||||||
|
END IF;
|
||||||
|
END;
|
||||||
|
$$"#;
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
pg_and_generic().verified_stmt(sql1),
|
||||||
|
Statement::CreateFunction(CreateFunction {
|
||||||
|
or_alter: false,
|
||||||
|
or_replace: true,
|
||||||
|
temporary: false,
|
||||||
|
name: ObjectName::from(vec![Ident::new("check_strings_different")]),
|
||||||
|
args: Some(vec![
|
||||||
|
OperateFunctionArg::with_name(
|
||||||
|
"str1",
|
||||||
|
DataType::Varchar(None),
|
||||||
|
),
|
||||||
|
OperateFunctionArg::with_name(
|
||||||
|
"str2",
|
||||||
|
DataType::Varchar(None),
|
||||||
|
),
|
||||||
|
]),
|
||||||
|
return_type: Some(DataType::Boolean),
|
||||||
|
language: Some("plpgsql".into()),
|
||||||
|
behavior: None,
|
||||||
|
called_on_null: None,
|
||||||
|
parallel: None,
|
||||||
|
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
|
||||||
|
(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()
|
||||||
|
))),
|
||||||
|
if_not_exists: false,
|
||||||
|
using: None,
|
||||||
|
determinism_specifier: None,
|
||||||
|
options: None,
|
||||||
|
remote_connection: None,
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
let sql2 = r#"CREATE OR REPLACE FUNCTION check_not_zero(int1 INT) RETURNS BOOLEAN LANGUAGE plpgsql AS $$
|
||||||
|
BEGIN
|
||||||
|
IF int1 <> 0 THEN
|
||||||
|
RETURN TRUE;
|
||||||
|
ELSE
|
||||||
|
RETURN FALSE;
|
||||||
|
END IF;
|
||||||
|
END;
|
||||||
|
$$"#;
|
||||||
|
assert_eq!(
|
||||||
|
pg_and_generic().verified_stmt(sql2),
|
||||||
|
Statement::CreateFunction(CreateFunction {
|
||||||
|
or_alter: false,
|
||||||
|
or_replace: true,
|
||||||
|
temporary: false,
|
||||||
|
name: ObjectName::from(vec![Ident::new("check_not_zero")]),
|
||||||
|
args: Some(vec![
|
||||||
|
OperateFunctionArg::with_name(
|
||||||
|
"int1",
|
||||||
|
DataType::Int(None)
|
||||||
|
)
|
||||||
|
]),
|
||||||
|
return_type: Some(DataType::Boolean),
|
||||||
|
language: Some("plpgsql".into()),
|
||||||
|
behavior: None,
|
||||||
|
called_on_null: None,
|
||||||
|
parallel: None,
|
||||||
|
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
|
||||||
|
(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()
|
||||||
|
))),
|
||||||
|
if_not_exists: false,
|
||||||
|
using: None,
|
||||||
|
determinism_specifier: None,
|
||||||
|
options: None,
|
||||||
|
remote_connection: None,
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
let sql3 = r#"CREATE OR REPLACE FUNCTION check_values_different(a INT, b INT) RETURNS BOOLEAN LANGUAGE plpgsql AS $$
|
||||||
|
BEGIN
|
||||||
|
IF a <> b THEN
|
||||||
|
RETURN TRUE;
|
||||||
|
ELSE
|
||||||
|
RETURN FALSE;
|
||||||
|
END IF;
|
||||||
|
END;
|
||||||
|
$$"#;
|
||||||
|
assert_eq!(
|
||||||
|
pg_and_generic().verified_stmt(sql3),
|
||||||
|
Statement::CreateFunction(CreateFunction {
|
||||||
|
or_alter: false,
|
||||||
|
or_replace: true,
|
||||||
|
temporary: false,
|
||||||
|
name: ObjectName::from(vec![Ident::new("check_values_different")]),
|
||||||
|
args: Some(vec![
|
||||||
|
OperateFunctionArg::with_name(
|
||||||
|
"a",
|
||||||
|
DataType::Int(None)
|
||||||
|
),
|
||||||
|
OperateFunctionArg::with_name(
|
||||||
|
"b",
|
||||||
|
DataType::Int(None)
|
||||||
|
),
|
||||||
|
]),
|
||||||
|
return_type: Some(DataType::Boolean),
|
||||||
|
language: Some("plpgsql".into()),
|
||||||
|
behavior: None,
|
||||||
|
called_on_null: None,
|
||||||
|
parallel: None,
|
||||||
|
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
|
||||||
|
(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()
|
||||||
|
))),
|
||||||
|
if_not_exists: false,
|
||||||
|
using: None,
|
||||||
|
determinism_specifier: None,
|
||||||
|
options: None,
|
||||||
|
remote_connection: None,
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
let sql4 = r#"CREATE OR REPLACE FUNCTION check_values_different(int1 INT, int2 INT) RETURNS BOOLEAN LANGUAGE plpgsql AS $$
|
||||||
|
BEGIN
|
||||||
|
IF int1 <> int2 THEN
|
||||||
|
RETURN TRUE;
|
||||||
|
ELSE
|
||||||
|
RETURN FALSE;
|
||||||
|
END IF;
|
||||||
|
END;
|
||||||
|
$$"#;
|
||||||
|
assert_eq!(
|
||||||
|
pg_and_generic().verified_stmt(sql4),
|
||||||
|
Statement::CreateFunction(CreateFunction {
|
||||||
|
or_alter: false,
|
||||||
|
or_replace: true,
|
||||||
|
temporary: false,
|
||||||
|
name: ObjectName::from(vec![Ident::new("check_values_different")]),
|
||||||
|
args: Some(vec![
|
||||||
|
OperateFunctionArg::with_name(
|
||||||
|
"int1",
|
||||||
|
DataType::Int(None)
|
||||||
|
),
|
||||||
|
OperateFunctionArg::with_name(
|
||||||
|
"int2",
|
||||||
|
DataType::Int(None)
|
||||||
|
),
|
||||||
|
]),
|
||||||
|
return_type: Some(DataType::Boolean),
|
||||||
|
language: Some("plpgsql".into()),
|
||||||
|
behavior: None,
|
||||||
|
called_on_null: None,
|
||||||
|
parallel: None,
|
||||||
|
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
|
||||||
|
(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()
|
||||||
|
))),
|
||||||
|
if_not_exists: false,
|
||||||
|
using: None,
|
||||||
|
determinism_specifier: None,
|
||||||
|
options: None,
|
||||||
|
remote_connection: None,
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
let sql5 = r#"CREATE OR REPLACE FUNCTION foo(a TIMESTAMP WITH TIME ZONE, b VARCHAR) RETURNS BOOLEAN LANGUAGE plpgsql AS $$
|
||||||
|
BEGIN
|
||||||
|
RETURN TRUE;
|
||||||
|
END;
|
||||||
|
$$"#;
|
||||||
|
assert_eq!(
|
||||||
|
pg_and_generic().verified_stmt(sql5),
|
||||||
|
Statement::CreateFunction(CreateFunction {
|
||||||
|
or_alter: false,
|
||||||
|
or_replace: true,
|
||||||
|
temporary: false,
|
||||||
|
name: ObjectName::from(vec![Ident::new("foo")]),
|
||||||
|
args: Some(vec![
|
||||||
|
OperateFunctionArg::with_name(
|
||||||
|
"a",
|
||||||
|
DataType::Timestamp(None, TimezoneInfo::WithTimeZone)
|
||||||
|
),
|
||||||
|
OperateFunctionArg::with_name("b", DataType::Varchar(None)),
|
||||||
|
]),
|
||||||
|
return_type: Some(DataType::Boolean),
|
||||||
|
language: Some("plpgsql".into()),
|
||||||
|
behavior: None,
|
||||||
|
called_on_null: None,
|
||||||
|
parallel: None,
|
||||||
|
function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
|
||||||
|
(Value::DollarQuotedString(DollarQuotedString {
|
||||||
|
value: "\n BEGIN\n RETURN TRUE;\n END;\n ".to_owned(),
|
||||||
|
tag: None
|
||||||
|
}))
|
||||||
|
.with_empty_span()
|
||||||
|
))),
|
||||||
|
if_not_exists: false,
|
||||||
|
using: None,
|
||||||
|
determinism_specifier: None,
|
||||||
|
options: None,
|
||||||
|
remote_connection: None,
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
let incorrect_sql = "CREATE FUNCTION add(function(struct<a,b> int64), b INTEGER) RETURNS INTEGER LANGUAGE SQL IMMUTABLE STRICT PARALLEL SAFE AS 'select $1 + $2;'";
|
||||||
|
assert!(pg().parse_sql_statements(incorrect_sql).is_err(),);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parse_create_function() {
|
fn parse_create_function() {
|
||||||
let sql = "CREATE FUNCTION add(INTEGER, INTEGER) RETURNS INTEGER LANGUAGE SQL IMMUTABLE STRICT PARALLEL SAFE AS 'select $1 + $2;'";
|
let sql = "CREATE FUNCTION add(INTEGER, INTEGER) RETURNS INTEGER LANGUAGE SQL IMMUTABLE STRICT PARALLEL SAFE AS 'select $1 + $2;'";
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue