Skip to content

Commit 6dd521a

Browse files
tomershaniiitomershaniii
committed
Support for projection prefix operator (CONNECT_BY_ROOT)
1 parent 62495f2 commit 6dd521a

18 files changed

Lines changed: 1147 additions & 700 deletions

‎src/ast/query.rs‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -622,9 +622,13 @@ pub enum SelectItemQualifiedWildcardKind {
622622
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
623623
pub enum SelectItem {
624624
/// Any expression, not followed by `[ AS ] alias`
625-
UnnamedExpr(Expr),
625+
UnnamedExpr { expr: Expr, prefix: Option<Expr> },
626626
/// An expression, followed by `[ AS ] alias`
627-
ExprWithAlias { expr: Expr, alias: Ident },
627+
ExprWithAlias {
628+
expr: Expr,
629+
alias: Ident,
630+
prefix: Option<Expr>,
631+
},
628632
/// An expression, followed by a wildcard expansion.
629633
/// e.g. `alias.*`, `STRUCT<STRING>('foo').*`
630634
QualifiedWildcard(SelectItemQualifiedWildcardKind, WildcardAdditionalOptions),
@@ -907,8 +911,22 @@ impl fmt::Display for ReplaceSelectElement {
907911
impl fmt::Display for SelectItem {
908912
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
909913
match &self {
910-
SelectItem::UnnamedExpr(expr) => write!(f, "{expr}"),
911-
SelectItem::ExprWithAlias { expr, alias } => write!(f, "{expr} AS {alias}"),
914+
SelectItem::UnnamedExpr { expr, prefix } => {
915+
if let Some(expr) = prefix {
916+
write!(f, "{expr} ")?
917+
}
918+
write!(f, "{expr}")
919+
}
920+
SelectItem::ExprWithAlias {
921+
expr,
922+
alias,
923+
prefix,
924+
} => {
925+
if let Some(expr) = prefix {
926+
write!(f, "{expr} ")?
927+
}
928+
write!(f, "{expr} AS {alias}")
929+
}
912930
SelectItem::QualifiedWildcard(kind, additional_options) => {
913931
write!(f, "{kind}")?;
914932
write!(f, "{additional_options}")?;

‎src/ast/spans.rs‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1722,8 +1722,20 @@ impl Spanned for SelectItemQualifiedWildcardKind {
17221722
impl Spanned for SelectItem {
17231723
fn span(&self) -> Span {
17241724
match self {
1725-
SelectItem::UnnamedExpr(expr) => expr.span(),
1726-
SelectItem::ExprWithAlias { expr, alias } => expr.span().union(&alias.span),
1725+
SelectItem::UnnamedExpr { expr, prefix } => expr
1726+
.span()
1727+
.union_opt(&prefix.as_ref().map(|expr| expr.span())),
1728+
1729+
SelectItem::ExprWithAlias {
1730+
expr,
1731+
alias,
1732+
prefix,
1733+
} => {
1734+
// let x = &prefix.map(|i| i.span());
1735+
expr.span()
1736+
.union(&alias.span)
1737+
.union_opt(&prefix.as_ref().map(|i| i.span()))
1738+
}
17271739
SelectItem::QualifiedWildcard(kind, wildcard_additional_options) => union_spans(
17281740
[kind.span()]
17291741
.into_iter()

‎src/dialect/mod.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -888,6 +888,12 @@ pub trait Dialect: Debug + Any {
888888
keywords::RESERVED_FOR_TABLE_FACTOR
889889
}
890890

891+
/// Returns reserved keywords for projection item prefix operator
892+
/// e.g. SELECT CONNECT_BY_ROOT name FROM Tbl2 (Snowflake)
893+
fn get_reserved_keywords_for_select_item_operator(&self) -> &[Keyword] {
894+
&[]
895+
}
896+
891897
/// Returns true if this dialect supports the `TABLESAMPLE` option
892898
/// before the table alias option. For example:
893899
///

‎src/dialect/snowflake.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -346,6 +346,13 @@ impl Dialect for SnowflakeDialect {
346346
fn supports_group_by_expr(&self) -> bool {
347347
true
348348
}
349+
350+
/// See: <https://docs.snowflake.com/en/sql-reference/constructs/connect-by>
351+
fn get_reserved_keywords_for_select_item_operator(&self) -> &[Keyword] {
352+
const RESERVED_KEYWORDS_FOR_SELECT_ITEM_OPERATOR: [Keyword; 1] = [Keyword::CONNECT_BY_ROOT];
353+
354+
&RESERVED_KEYWORDS_FOR_SELECT_ITEM_OPERATOR
355+
}
349356
}
350357

351358
fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result<Statement, ParserError> {

‎src/keywords.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,7 @@ define_keywords!(
207207
CONNECT,
208208
CONNECTION,
209209
CONNECTOR,
210+
CONNECT_BY_ROOT,
210211
CONSTRAINT,
211212
CONTAINS,
212213
CONTINUE,

‎src/parser/mod.rs‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1237,6 +1237,23 @@ impl<'a> Parser<'a> {
12371237
}
12381238
}
12391239

1240+
//Select item operators
1241+
fn parse_select_item_prefix_by_reserved_word(&mut self) -> Result<Option<Expr>, ParserError> {
1242+
if let Some(kw) = self.peek_one_of_keywords(
1243+
self.dialect
1244+
.get_reserved_keywords_for_select_item_operator(),
1245+
) {
1246+
if let TokenWithSpan {
1247+
span,
1248+
token: Token::Word(word),
1249+
} = self.expect_keyword(kw)?
1250+
{
1251+
return Ok(Some(Expr::Identifier(word.clone().into_ident(span))));
1252+
}
1253+
}
1254+
Ok(None)
1255+
}
1256+
12401257
/// Tries to parse an expression by matching the specified word to known keywords that have a special meaning in the dialect.
12411258
/// Returns `None if no match is found.
12421259
fn parse_expr_prefix_by_reserved_word(
@@ -13732,6 +13749,10 @@ impl<'a> Parser<'a> {
1373213749

1373313750
/// Parse a comma-delimited list of projections after SELECT
1373413751
pub fn parse_select_item(&mut self) -> Result<SelectItem, ParserError> {
13752+
let prefix = self
13753+
.maybe_parse(|parser| parser.parse_select_item_prefix_by_reserved_word())?
13754+
.flatten();
13755+
1373513756
match self.parse_wildcard_expr()? {
1373613757
Expr::QualifiedWildcard(prefix, token) => Ok(SelectItem::QualifiedWildcard(
1373713758
SelectItemQualifiedWildcardKind::ObjectName(prefix),
@@ -13762,6 +13783,7 @@ impl<'a> Parser<'a> {
1376213783
Ok(SelectItem::ExprWithAlias {
1376313784
expr: *right,
1376413785
alias,
13786+
prefix,
1376513787
})
1376613788
}
1376713789
expr if self.dialect.supports_select_expr_star()
@@ -13776,8 +13798,12 @@ impl<'a> Parser<'a> {
1377613798
expr => self
1377713799
.maybe_parse_select_item_alias()
1377813800
.map(|alias| match alias {
13779-
Some(alias) => SelectItem::ExprWithAlias { expr, alias },
13780-
None => SelectItem::UnnamedExpr(expr),
13801+
Some(alias) => SelectItem::ExprWithAlias {
13802+
expr,
13803+
alias,
13804+
prefix,
13805+
},
13806+
None => SelectItem::UnnamedExpr { expr, prefix },
1378113807
}),
1378213808
}
1378313809
}

‎src/test_utils.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,7 @@ pub fn only<T>(v: impl IntoIterator<Item = T>) -> T {
311311

312312
pub fn expr_from_projection(item: &SelectItem) -> &Expr {
313313
match item {
314-
SelectItem::UnnamedExpr(expr) => expr,
314+
SelectItem::UnnamedExpr { expr, .. } => expr,
315315
_ => panic!("Expected UnnamedExpr"),
316316
}
317317
}

‎tests/sqlparser_bigquery.rs‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1626,14 +1626,17 @@ fn parse_hyphenated_table_identifiers() {
16261626
"SELECT foo - bar.x FROM t"
16271627
)
16281628
.projection[0],
1629-
SelectItem::UnnamedExpr(Expr::BinaryOp {
1630-
left: Box::new(Expr::Identifier(Ident::new("foo"))),
1631-
op: BinaryOperator::Minus,
1632-
right: Box::new(Expr::CompoundIdentifier(vec![
1633-
Ident::new("bar"),
1634-
Ident::new("x")
1635-
]))
1636-
})
1629+
SelectItem::UnnamedExpr {
1630+
expr: Expr::BinaryOp {
1631+
left: Box::new(Expr::Identifier(Ident::new("foo"))),
1632+
op: BinaryOperator::Minus,
1633+
right: Box::new(Expr::CompoundIdentifier(vec![
1634+
Ident::new("bar"),
1635+
Ident::new("x")
1636+
]))
1637+
},
1638+
prefix: None
1639+
}
16371640
);
16381641
}
16391642

‎tests/sqlparser_clickhouse.rs‎

Lines changed: 42 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -44,22 +44,25 @@ fn parse_map_access_expr() {
4444
select_token: AttachedToken::empty(),
4545
top: None,
4646
top_before_distinct: false,
47-
projection: vec![UnnamedExpr(Expr::CompoundFieldAccess {
48-
root: Box::new(Identifier(Ident {
49-
value: "string_values".to_string(),
50-
quote_style: None,
51-
span: Span::empty(),
52-
})),
53-
access_chain: vec![AccessExpr::Subscript(Subscript::Index {
54-
index: call(
55-
"indexOf",
56-
[
57-
Expr::Identifier(Ident::new("string_names")),
58-
Expr::value(Value::SingleQuotedString("endpoint".to_string()))
59-
]
60-
),
61-
})],
62-
})],
47+
projection: vec![UnnamedExpr {
48+
expr: Expr::CompoundFieldAccess {
49+
root: Box::new(Identifier(Ident {
50+
value: "string_values".to_string(),
51+
quote_style: None,
52+
span: Span::empty(),
53+
})),
54+
access_chain: vec![AccessExpr::Subscript(Subscript::Index {
55+
index: call(
56+
"indexOf",
57+
[
58+
Expr::Identifier(Ident::new("string_names")),
59+
Expr::value(Value::SingleQuotedString("endpoint".to_string()))
60+
]
61+
),
62+
})],
63+
},
64+
prefix: None
65+
}],
6366
into: None,
6467
from: vec![TableWithJoins {
6568
relation: table_from_name(ObjectName::from(vec![Ident::new("foos")])),
@@ -205,7 +208,11 @@ fn parse_delimited_identifiers() {
205208
expr_from_projection(&select.projection[1]),
206209
);
207210
match &select.projection[2] {
208-
SelectItem::ExprWithAlias { expr, alias } => {
211+
SelectItem::ExprWithAlias {
212+
expr,
213+
alias,
214+
prefix: _,
215+
} => {
209216
assert_eq!(&Expr::Identifier(Ident::with_quote('"', "simple id")), expr);
210217
assert_eq!(&Ident::with_quote('"', "column alias"), alias);
211218
}
@@ -315,8 +322,14 @@ fn parse_alter_table_add_projection() {
315322
name: "my_name".into(),
316323
select: ProjectionSelect {
317324
projection: vec![
318-
UnnamedExpr(Identifier(Ident::new("a"))),
319-
UnnamedExpr(Identifier(Ident::new("b"))),
325+
UnnamedExpr {
326+
expr: Identifier(Ident::new("a")),
327+
prefix: None
328+
},
329+
UnnamedExpr {
330+
expr: Identifier(Ident::new("b")),
331+
prefix: None
332+
},
320333
],
321334
group_by: Some(GroupByExpr::Expressions(
322335
vec![Identifier(Ident::new("a"))],
@@ -1000,7 +1013,10 @@ fn parse_select_parametric_function() {
10001013
let projection: &Vec<SelectItem> = query.body.as_select().unwrap().projection.as_ref();
10011014
assert_eq!(projection.len(), 1);
10021015
match &projection[0] {
1003-
UnnamedExpr(Expr::Function(f)) => {
1016+
UnnamedExpr {
1017+
expr: Expr::Function(f),
1018+
..
1019+
} => {
10041020
let args = match &f.args {
10051021
FunctionArguments::List(ref args) => args,
10061022
_ => unreachable!(),
@@ -1418,9 +1434,12 @@ fn parse_create_table_on_commit_and_as_query() {
14181434
assert_eq!(on_commit, Some(OnCommit::PreserveRows));
14191435
assert_eq!(
14201436
query.unwrap().body.as_select().unwrap().projection,
1421-
vec![UnnamedExpr(Expr::Value(
1422-
(Value::Number("1".parse().unwrap(), false)).with_empty_span()
1423-
))]
1437+
vec![UnnamedExpr {
1438+
expr: Expr::Value(
1439+
(Value::Number("1".parse().unwrap(), false)).with_empty_span()
1440+
),
1441+
prefix: None
1442+
}]
14241443
);
14251444
}
14261445
_ => unreachable!(),

0 commit comments

Comments
 (0)