mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-07-07 17:04:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Parse Snowflake USE ROLE and USE SECONDARY ROLES (#1578)

Browse source
This commit is contained in:
Yoav Cohen 2024-12-05 19:17:52 +01:00 • committed by GitHub
parent dd7ba72a0b
commit 7b50ac31c3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 115 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

41
src/ast/dcl.rs
View file

@ -28,7 +28,7 @@ use serde::{Deserialize, Serialize};
#[cfg(feature = "visitor")]
use sqlparser_derive::{Visit, VisitMut};
use super::{Expr, Ident, Password};
use super::{display_comma_separated, Expr, Ident, Password};
use crate::ast::{display_separated, ObjectName};
/// An option in `ROLE` statement.
@ -204,12 +204,14 @@ impl fmt::Display for AlterRoleOperation {
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum Use {
Catalog(ObjectName), // e.g. `USE CATALOG foo.bar`
Schema(ObjectName), // e.g. `USE SCHEMA foo.bar`
Database(ObjectName), // e.g. `USE DATABASE foo.bar`
Warehouse(ObjectName), // e.g. `USE WAREHOUSE foo.bar`
Object(ObjectName), // e.g. `USE foo.bar`
Default, // e.g. `USE DEFAULT`
Catalog(ObjectName), // e.g. `USE CATALOG foo.bar`
Schema(ObjectName), // e.g. `USE SCHEMA foo.bar`
Database(ObjectName), // e.g. `USE DATABASE foo.bar`
Warehouse(ObjectName), // e.g. `USE WAREHOUSE foo.bar`
Role(ObjectName), // e.g. `USE ROLE PUBLIC`
SecondaryRoles(SecondaryRoles), // e.g. `USE SECONDARY ROLES ALL`
Object(ObjectName), // e.g. `USE foo.bar`
Default, // e.g. `USE DEFAULT`
}
impl fmt::Display for Use {
@ -220,8 +222,33 @@ impl fmt::Display for Use {
Use::Schema(name) => write!(f, "SCHEMA {}", name),
Use::Database(name) => write!(f, "DATABASE {}", name),
Use::Warehouse(name) => write!(f, "WAREHOUSE {}", name),
Use::Role(name) => write!(f, "ROLE {}", name),
Use::SecondaryRoles(secondary_roles) => {
write!(f, "SECONDARY ROLES {}", secondary_roles)
}
Use::Object(name) => write!(f, "{}", name),
Use::Default => write!(f, "DEFAULT"),
}
}
}
/// Snowflake `SECONDARY ROLES` USE variant
/// See: <https://docs.snowflake.com/en/sql-reference/sql/use-secondary-roles>
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum SecondaryRoles {
All,
None,
List(Vec<Ident>),
}
impl fmt::Display for SecondaryRoles {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
SecondaryRoles::All => write!(f, "ALL"),
SecondaryRoles::None => write!(f, "NONE"),
SecondaryRoles::List(roles) => write!(f, "{}", display_comma_separated(roles)),
}
}
}

4
src/ast/mod.rs
View file

@ -43,7 +43,9 @@ pub use self::data_type::{
ArrayElemTypeDef, CharLengthUnits, CharacterLength, DataType, EnumMember, ExactNumberInfo,
StructBracketKind, TimezoneInfo,
};
pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue, Use};
pub use self::dcl::{
AlterRoleOperation, ResetConfig, RoleOption, SecondaryRoles, SetConfigValue, Use,
};
pub use self::ddl::{
AlterColumnOperation, AlterIndexOperation, AlterPolicyOperation, AlterTableOperation,
ClusteredBy, ColumnDef, ColumnOption, ColumnOptionDef, ColumnPolicy, ColumnPolicyProperty,

17
src/ast/spans.rs
View file

@ -3,11 +3,11 @@ use core::iter;
use crate::tokenizer::Span;
use super::{
AlterColumnOperation, AlterIndexOperation, AlterTableOperation, Array, Assignment,
AssignmentTarget, CloseCursor, ClusteredIndex, ColumnDef, ColumnOption, ColumnOptionDef,
ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex, CreateTable,
CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr,
ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
dcl::SecondaryRoles, AlterColumnOperation, AlterIndexOperation, AlterTableOperation, Array,
Assignment, AssignmentTarget, CloseCursor, ClusteredIndex, ColumnDef, ColumnOption,
ColumnOptionDef, ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex,
CreateTable, CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem,
Expr, ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr, HavingBound,
IlikeSelectItem, Insert, Interpolate, InterpolateExpr, Join, JoinConstraint, JoinOperator,
JsonPath, JsonPathElem, LateralView, MatchRecognizePattern, Measure, NamedWindowDefinition,
@ -484,6 +484,13 @@ impl Spanned for Use {
Use::Schema(object_name) => object_name.span(),
Use::Database(object_name) => object_name.span(),
Use::Warehouse(object_name) => object_name.span(),
Use::Role(object_name) => object_name.span(),
Use::SecondaryRoles(secondary_roles) => {
if let SecondaryRoles::List(roles) = secondary_roles {
return union_spans(roles.iter().map(|i| i.span));
}
Span::empty()
}
Use::Object(object_name) => object_name.span(),
Use::Default => Span::empty(),
}

2
src/keywords.rs
View file

@ -664,6 +664,7 @@ define_keywords!(
RIGHT,
RLIKE,
ROLE,
ROLES,
ROLLBACK,
ROLLUP,
ROOT,
@ -682,6 +683,7 @@ define_keywords!(
SCROLL,
SEARCH,
SECOND,
SECONDARY,
SECRET,
SECURITY,
SELECT,

39
src/parser/mod.rs
View file

@ -10093,23 +10093,46 @@ impl<'a> Parser<'a> {
} else if dialect_of!(self is DatabricksDialect) {
self.parse_one_of_keywords(&[Keyword::CATALOG, Keyword::DATABASE, Keyword::SCHEMA])
} else if dialect_of!(self is SnowflakeDialect) {
self.parse_one_of_keywords(&[Keyword::DATABASE, Keyword::SCHEMA, Keyword::WAREHOUSE])
self.parse_one_of_keywords(&[
Keyword::DATABASE,
Keyword::SCHEMA,
Keyword::WAREHOUSE,
Keyword::ROLE,
Keyword::SECONDARY,
])
} else {
None // No specific keywords for other dialects, including GenericDialect
};
let obj_name = self.parse_object_name(false)?;
let result = match parsed_keyword {
Some(Keyword::CATALOG) => Use::Catalog(obj_name),
Some(Keyword::DATABASE) => Use::Database(obj_name),
Some(Keyword::SCHEMA) => Use::Schema(obj_name),
Some(Keyword::WAREHOUSE) => Use::Warehouse(obj_name),
_ => Use::Object(obj_name),
let result = if matches!(parsed_keyword, Some(Keyword::SECONDARY)) {
self.parse_secondary_roles()?
} else {
let obj_name = self.parse_object_name(false)?;
match parsed_keyword {
Some(Keyword::CATALOG) => Use::Catalog(obj_name),
Some(Keyword::DATABASE) => Use::Database(obj_name),
Some(Keyword::SCHEMA) => Use::Schema(obj_name),
Some(Keyword::WAREHOUSE) => Use::Warehouse(obj_name),
Some(Keyword::ROLE) => Use::Role(obj_name),
_ => Use::Object(obj_name),
}
};
Ok(Statement::Use(result))
}
fn parse_secondary_roles(&mut self) -> Result<Use, ParserError> {
self.expect_keyword(Keyword::ROLES)?;
if self.parse_keyword(Keyword::NONE) {
Ok(Use::SecondaryRoles(SecondaryRoles::None))
} else if self.parse_keyword(Keyword::ALL) {
Ok(Use::SecondaryRoles(SecondaryRoles::All))
} else {
let roles = self.parse_comma_separated(|parser| parser.parse_identifier(false))?;
Ok(Use::SecondaryRoles(SecondaryRoles::List(roles)))
}
}
pub fn parse_table_and_joins(&mut self) -> Result<TableWithJoins, ParserError> {
let relation = self.parse_table_factor()?;
// Note that for keywords to be properly handled here, they need to be

14
tests/sqlparser_common.rs
View file

@ -4066,8 +4066,8 @@ fn test_alter_table_with_on_cluster() {
Statement::AlterTable {
name, on_cluster, ..
} => {
std::assert_eq!(name.to_string(), "t");
std::assert_eq!(on_cluster, Some(Ident::with_quote('\'', "cluster")));
assert_eq!(name.to_string(), "t");
assert_eq!(on_cluster, Some(Ident::with_quote('\'', "cluster")));
}
_ => unreachable!(),
}
@ -4078,15 +4078,15 @@ fn test_alter_table_with_on_cluster() {
Statement::AlterTable {
name, on_cluster, ..
} => {
std::assert_eq!(name.to_string(), "t");
std::assert_eq!(on_cluster, Some(Ident::new("cluster_name")));
assert_eq!(name.to_string(), "t");
assert_eq!(on_cluster, Some(Ident::new("cluster_name")));
}
_ => unreachable!(),
}
let res = all_dialects()
.parse_sql_statements("ALTER TABLE t ON CLUSTER 123 ADD CONSTRAINT bar PRIMARY KEY (baz)");
std::assert_eq!(
assert_eq!(
res.unwrap_err(),
ParserError::ParserError("Expected: identifier, found: 123".to_string())
)
@ -11226,7 +11226,7 @@ fn test_group_by_nothing() {
let Select { group_by, .. } = all_dialects_where(|d| d.supports_group_by_expr())
.verified_only_select("SELECT count(1) FROM t GROUP BY ()");
{
std::assert_eq!(
assert_eq!(
GroupByExpr::Expressions(vec![Expr::Tuple(vec![])], vec![]),
group_by
);
@ -11235,7 +11235,7 @@ fn test_group_by_nothing() {
let Select { group_by, .. } = all_dialects_where(|d| d.supports_group_by_expr())
.verified_only_select("SELECT name, count(1) FROM t GROUP BY name, ()");
{
std::assert_eq!(
assert_eq!(
GroupByExpr::Expressions(
vec![
Identifier(Ident::new("name".to_string())),

34
tests/sqlparser_snowflake.rs
View file

@ -2649,7 +2649,7 @@ fn parse_use() {
let quote_styles = ['\'', '"', '`'];
for object_name in &valid_object_names {
// Test single identifier without quotes
std::assert_eq!(
assert_eq!(
snowflake().verified_stmt(&format!("USE {}", object_name)),
Statement::Use(Use::Object(ObjectName(vec![Ident::new(
object_name.to_string()
@ -2657,7 +2657,7 @@ fn parse_use() {
);
for &quote in &quote_styles {
// Test single identifier with different type of quotes
std::assert_eq!(
assert_eq!(
snowflake().verified_stmt(&format!("USE {}{}{}", quote, object_name, quote)),
Statement::Use(Use::Object(ObjectName(vec![Ident::with_quote(
quote,
@ -2669,7 +2669,7 @@ fn parse_use() {
for &quote in &quote_styles {
// Test double identifier with different type of quotes
std::assert_eq!(
assert_eq!(
snowflake().verified_stmt(&format!("USE {0}CATALOG{0}.{0}my_schema{0}", quote)),
Statement::Use(Use::Object(ObjectName(vec![
Ident::with_quote(quote, "CATALOG"),
@ -2678,7 +2678,7 @@ fn parse_use() {
);
}
// Test double identifier without quotes
std::assert_eq!(
assert_eq!(
snowflake().verified_stmt("USE mydb.my_schema"),
Statement::Use(Use::Object(ObjectName(vec![
Ident::new("mydb"),
@ -2688,37 +2688,55 @@ fn parse_use() {
for &quote in &quote_styles {
// Test single and double identifier with keyword and different type of quotes
std::assert_eq!(
assert_eq!(
snowflake().verified_stmt(&format!("USE DATABASE {0}my_database{0}", quote)),
Statement::Use(Use::Database(ObjectName(vec![Ident::with_quote(
quote,
"my_database".to_string(),
)])))
);
std::assert_eq!(
assert_eq!(
snowflake().verified_stmt(&format!("USE SCHEMA {0}my_schema{0}", quote)),
Statement::Use(Use::Schema(ObjectName(vec![Ident::with_quote(
quote,
"my_schema".to_string(),
)])))
);
std::assert_eq!(
assert_eq!(
snowflake().verified_stmt(&format!("USE SCHEMA {0}CATALOG{0}.{0}my_schema{0}", quote)),
Statement::Use(Use::Schema(ObjectName(vec![
Ident::with_quote(quote, "CATALOG"),
Ident::with_quote(quote, "my_schema")
])))
);
assert_eq!(
snowflake().verified_stmt(&format!("USE ROLE {0}my_role{0}", quote)),
Statement::Use(Use::Role(ObjectName(vec![Ident::with_quote(
quote,
"my_role".to_string(),
)])))
);
assert_eq!(
snowflake().verified_stmt(&format!("USE WAREHOUSE {0}my_wh{0}", quote)),
Statement::Use(Use::Warehouse(ObjectName(vec![Ident::with_quote(
quote,
"my_wh".to_string(),
)])))
);
}
// Test invalid syntax - missing identifier
let invalid_cases = ["USE SCHEMA", "USE DATABASE", "USE WAREHOUSE"];
for sql in &invalid_cases {
std::assert_eq!(
assert_eq!(
snowflake().parse_sql_statements(sql).unwrap_err(),
ParserError::ParserError("Expected: identifier, found: EOF".to_string()),
);
}
snowflake().verified_stmt("USE SECONDARY ROLES ALL");
snowflake().verified_stmt("USE SECONDARY ROLES NONE");
snowflake().verified_stmt("USE SECONDARY ROLES r1, r2, r3");
}
#[test]