mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-30 18:57:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Support DECLARE syntax for snowflake and bigquery (#1122)

Browse source
This commit is contained in:
Ifeanyi Ubah 2024-02-29 13:34:00 +01:00 • committed by GitHub
parent f75bb4be20
commit 57113a9344
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 741 additions and 68 deletions
Download patch file Download diff file Expand all files Collapse all files

64
tests/sqlparser_bigquery.rs
View file

@ -1278,6 +1278,70 @@ fn test_select_wildcard_with_replace() {
assert_eq!(expected, select.projection[0]);
}
#[test]
fn parse_big_query_declare() {
for (sql, expected_names, expected_data_type, expected_assigned_expr) in [
(
"DECLARE x INT64",
vec![Ident::new("x")],
Some(DataType::Int64),
None,
),
(
"DECLARE x INT64 DEFAULT 42",
vec![Ident::new("x")],
Some(DataType::Int64),
Some(DeclareAssignment::Default(Box::new(Expr::Value(number(
"42",
))))),
),
(
"DECLARE x, y, z INT64 DEFAULT 42",
vec![Ident::new("x"), Ident::new("y"), Ident::new("z")],
Some(DataType::Int64),
Some(DeclareAssignment::Default(Box::new(Expr::Value(number(
"42",
))))),
),
(
"DECLARE x DEFAULT 42",
vec![Ident::new("x")],
None,
Some(DeclareAssignment::Default(Box::new(Expr::Value(number(
"42",
))))),
),
] {
match bigquery().verified_stmt(sql) {
Statement::Declare { mut stmts } => {
assert_eq!(1, stmts.len());
let Declare {
names,
data_type,
assignment: assigned_expr,
..
} = stmts.swap_remove(0);
assert_eq!(expected_names, names);
assert_eq!(expected_data_type, data_type);
assert_eq!(expected_assigned_expr, assigned_expr);
}
_ => unreachable!(),
}
}
let error_sql = "DECLARE x";
assert_eq!(
ParserError::ParserError("Expected a data type name, found: EOF".to_owned()),
bigquery().parse_sql_statements(error_sql).unwrap_err()
);
let error_sql = "DECLARE x 42";
assert_eq!(
ParserError::ParserError("Expected a data type name, found: 42".to_owned()),
bigquery().parse_sql_statements(error_sql).unwrap_err()
);
}
fn bigquery() -> TestedDialects {
TestedDialects {
dialects: vec![Box::new(BigQueryDialect {})],

256
tests/sqlparser_snowflake.rs
View file

@ -579,6 +579,262 @@ fn test_drop_stage() {
.one_statement_parses_to("DROP STAGE IF EXISTS s1", "DROP STAGE IF EXISTS s1");
}
#[test]
fn parse_snowflake_declare_cursor() {
for (sql, expected_name, expected_assigned_expr, expected_query_projections) in [
(
"DECLARE c1 CURSOR FOR SELECT id, price FROM invoices",
"c1",
None,
Some(vec!["id", "price"]),
),
(
"DECLARE c1 CURSOR FOR res",
"c1",
Some(DeclareAssignment::For(
Expr::Identifier(Ident::new("res")).into(),
)),
None,
),
] {
match snowflake().verified_stmt(sql) {
Statement::Declare { mut stmts } => {
assert_eq!(1, stmts.len());
let Declare {
names,
data_type,
declare_type,
assignment: assigned_expr,
for_query,
..
} = stmts.swap_remove(0);
assert_eq!(vec![Ident::new(expected_name)], names);
assert!(data_type.is_none());
assert_eq!(Some(DeclareType::Cursor), declare_type);
assert_eq!(expected_assigned_expr, assigned_expr);
assert_eq!(
expected_query_projections,
for_query.as_ref().map(|q| {
match q.body.as_ref() {
SetExpr::Select(q) => q
.projection
.iter()
.map(|item| match item {
SelectItem::UnnamedExpr(Expr::Identifier(ident)) => {
ident.value.as_str()
}
_ => unreachable!(),
})
.collect::<Vec<_>>(),
_ => unreachable!(),
}
})
)
}
_ => unreachable!(),
}
}
let error_sql = "DECLARE c1 CURSOR SELECT id FROM invoices";
assert_eq!(
ParserError::ParserError("Expected FOR, found: SELECT".to_owned()),
snowflake().parse_sql_statements(error_sql).unwrap_err()
);
let error_sql = "DECLARE c1 CURSOR res";
assert_eq!(
ParserError::ParserError("Expected FOR, found: res".to_owned()),
snowflake().parse_sql_statements(error_sql).unwrap_err()
);
}
#[test]
fn parse_snowflake_declare_result_set() {
for (sql, expected_name, expected_assigned_expr) in [
(
"DECLARE res RESULTSET DEFAULT 42",
"res",
Some(DeclareAssignment::Default(Expr::Value(number("42")).into())),
),
(
"DECLARE res RESULTSET := 42",
"res",
Some(DeclareAssignment::DuckAssignment(
Expr::Value(number("42")).into(),
)),
),
("DECLARE res RESULTSET", "res", None),
] {
match snowflake().verified_stmt(sql) {
Statement::Declare { mut stmts } => {
assert_eq!(1, stmts.len());
let Declare {
names,
data_type,
declare_type,
assignment: assigned_expr,
for_query,
..
} = stmts.swap_remove(0);
assert_eq!(vec![Ident::new(expected_name)], names);
assert!(data_type.is_none());
assert!(for_query.is_none());
assert_eq!(Some(DeclareType::ResultSet), declare_type);
assert_eq!(expected_assigned_expr, assigned_expr);
}
_ => unreachable!(),
}
}
let sql = "DECLARE res RESULTSET DEFAULT (SELECT price FROM invoices)";
assert_eq!(snowflake().verified_stmt(sql).to_string(), sql);
let error_sql = "DECLARE res RESULTSET DEFAULT";
assert_eq!(
ParserError::ParserError("Expected an expression:, found: EOF".to_owned()),
snowflake().parse_sql_statements(error_sql).unwrap_err()
);
let error_sql = "DECLARE res RESULTSET :=";
assert_eq!(
ParserError::ParserError("Expected an expression:, found: EOF".to_owned()),
snowflake().parse_sql_statements(error_sql).unwrap_err()
);
}
#[test]
fn parse_snowflake_declare_exception() {
for (sql, expected_name, expected_assigned_expr) in [
(
"DECLARE ex EXCEPTION (42, 'ERROR')",
"ex",
Some(DeclareAssignment::Expr(
Expr::Tuple(vec![
Expr::Value(number("42")),
Expr::Value(Value::SingleQuotedString("ERROR".to_string())),
])
.into(),
)),
),
("DECLARE ex EXCEPTION", "ex", None),
] {
match snowflake().verified_stmt(sql) {
Statement::Declare { mut stmts } => {
assert_eq!(1, stmts.len());
let Declare {
names,
data_type,
declare_type,
assignment: assigned_expr,
for_query,
..
} = stmts.swap_remove(0);
assert_eq!(vec![Ident::new(expected_name)], names);
assert!(data_type.is_none());
assert!(for_query.is_none());
assert_eq!(Some(DeclareType::Exception), declare_type);
assert_eq!(expected_assigned_expr, assigned_expr);
}
_ => unreachable!(),
}
}
}
#[test]
fn parse_snowflake_declare_variable() {
for (sql, expected_name, expected_data_type, expected_assigned_expr) in [
(
"DECLARE profit TEXT DEFAULT 42",
"profit",
Some(DataType::Text),
Some(DeclareAssignment::Default(Expr::Value(number("42")).into())),
),
(
"DECLARE profit DEFAULT 42",
"profit",
None,
Some(DeclareAssignment::Default(Expr::Value(number("42")).into())),
),
("DECLARE profit TEXT", "profit", Some(DataType::Text), None),
("DECLARE profit", "profit", None, None),
] {
match snowflake().verified_stmt(sql) {
Statement::Declare { mut stmts } => {
assert_eq!(1, stmts.len());
let Declare {
names,
data_type,
declare_type,
assignment: assigned_expr,
for_query,
..
} = stmts.swap_remove(0);
assert_eq!(vec![Ident::new(expected_name)], names);
assert!(for_query.is_none());
assert_eq!(expected_data_type, data_type);
assert_eq!(None, declare_type);
assert_eq!(expected_assigned_expr, assigned_expr);
}
_ => unreachable!(),
}
}
snowflake().one_statement_parses_to("DECLARE profit;", "DECLARE profit");
let error_sql = "DECLARE profit INT 2";
assert_eq!(
ParserError::ParserError("Expected end of statement, found: 2".to_owned()),
snowflake().parse_sql_statements(error_sql).unwrap_err()
);
let error_sql = "DECLARE profit INT DEFAULT";
assert_eq!(
ParserError::ParserError("Expected an expression:, found: EOF".to_owned()),
snowflake().parse_sql_statements(error_sql).unwrap_err()
);
let error_sql = "DECLARE profit DEFAULT";
assert_eq!(
ParserError::ParserError("Expected an expression:, found: EOF".to_owned()),
snowflake().parse_sql_statements(error_sql).unwrap_err()
);
}
#[test]
fn parse_snowflake_declare_multi_statements() {
let sql = concat!(
"DECLARE profit DEFAULT 42; ",
"res RESULTSET DEFAULT (SELECT price FROM invoices); ",
"c1 CURSOR FOR res; ",
"ex EXCEPTION (-20003, 'ERROR: Could not create table.')"
);
match snowflake().verified_stmt(sql) {
Statement::Declare { stmts } => {
let actual = stmts
.iter()
.map(|stmt| (stmt.names[0].value.as_str(), stmt.declare_type.clone()))
.collect::<Vec<_>>();
assert_eq!(
vec![
("profit", None),
("res", Some(DeclareType::ResultSet)),
("c1", Some(DeclareType::Cursor)),
("ex", Some(DeclareType::Exception)),
],
actual
);
}
_ => unreachable!(),
}
let error_sql = "DECLARE profit DEFAULT 42 c1 CURSOR FOR res;";
assert_eq!(
ParserError::ParserError("Expected end of statement, found: c1".to_owned()),
snowflake().parse_sql_statements(error_sql).unwrap_err()
);
}
#[test]
fn test_create_stage() {
let sql = "CREATE STAGE s1.s2";