Skip to content

Support for projection item prefix operator (CONNECT_BY_ROOT) #1780

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 22 additions & 4 deletions src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -622,9 +622,13 @@ pub enum SelectItemQualifiedWildcardKind {
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum SelectItem {
/// Any expression, not followed by `[ AS ] alias`
UnnamedExpr(Expr),
UnnamedExpr { expr: Expr, prefix: Option<Expr> },
/// An expression, followed by `[ AS ] alias`
ExprWithAlias { expr: Expr, alias: Ident },
ExprWithAlias {
expr: Expr,
alias: Ident,
prefix: Option<Expr>,
},
/// An expression, followed by a wildcard expansion.
/// e.g. `alias.*`, `STRUCT<STRING>('foo').*`
QualifiedWildcard(SelectItemQualifiedWildcardKind, WildcardAdditionalOptions),
Expand Down Expand Up @@ -907,8 +911,22 @@ impl fmt::Display for ReplaceSelectElement {
impl fmt::Display for SelectItem {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match &self {
SelectItem::UnnamedExpr(expr) => write!(f, "{expr}"),
SelectItem::ExprWithAlias { expr, alias } => write!(f, "{expr} AS {alias}"),
SelectItem::UnnamedExpr { expr, prefix } => {
if let Some(expr) = prefix {
write!(f, "{expr} ")?
}
write!(f, "{expr}")
}
SelectItem::ExprWithAlias {
expr,
alias,
prefix,
} => {
if let Some(expr) = prefix {
write!(f, "{expr} ")?
}
write!(f, "{expr} AS {alias}")
}
SelectItem::QualifiedWildcard(kind, additional_options) => {
write!(f, "{kind}")?;
write!(f, "{additional_options}")?;
Expand Down
16 changes: 14 additions & 2 deletions src/ast/spans.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1722,8 +1722,20 @@ impl Spanned for SelectItemQualifiedWildcardKind {
impl Spanned for SelectItem {
fn span(&self) -> Span {
match self {
SelectItem::UnnamedExpr(expr) => expr.span(),
SelectItem::ExprWithAlias { expr, alias } => expr.span().union(&alias.span),
SelectItem::UnnamedExpr { expr, prefix } => expr
.span()
.union_opt(&prefix.as_ref().map(|expr| expr.span())),

SelectItem::ExprWithAlias {
expr,
alias,
prefix,
} => {
// let x = &prefix.map(|i| i.span());
expr.span()
.union(&alias.span)
.union_opt(&prefix.as_ref().map(|i| i.span()))
}
SelectItem::QualifiedWildcard(kind, wildcard_additional_options) => union_spans(
[kind.span()]
.into_iter()
Expand Down
6 changes: 6 additions & 0 deletions src/dialect/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -888,6 +888,12 @@ pub trait Dialect: Debug + Any {
keywords::RESERVED_FOR_TABLE_FACTOR
}

/// Returns reserved keywords that may prefix a select item expression
/// e.g. `SELECT CONNECT_BY_ROOT name FROM Tbl2` (Snowflake)
fn get_reserved_keywords_for_select_item(&self) -> &[Keyword] {
&[]
}

/// Returns true if this dialect supports the `TABLESAMPLE` option
/// before the table alias option. For example:
///
Expand Down
7 changes: 7 additions & 0 deletions src/dialect/snowflake.rs
Original file line number Diff line number Diff line change
Expand Up @@ -346,6 +346,13 @@ impl Dialect for SnowflakeDialect {
fn supports_group_by_expr(&self) -> bool {
true
}

/// See: <https://docs.snowflake.com/en/sql-reference/constructs/connect-by>
fn get_reserved_keywords_for_select_item_operator(&self) -> &[Keyword] {
const RESERVED_KEYWORDS_FOR_SELECT_ITEM_OPERATOR: [Keyword; 1] = [Keyword::CONNECT_BY_ROOT];
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could we move this to the top of the file maybe? thinking it would be more visible there and if the list happens to get long it doesnt add to the function's length


&RESERVED_KEYWORDS_FOR_SELECT_ITEM_OPERATOR
}
}

fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result<Statement, ParserError> {
Expand Down
1 change: 1 addition & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,7 @@ define_keywords!(
CONNECT,
CONNECTION,
CONNECTOR,
CONNECT_BY_ROOT,
CONSTRAINT,
CONTAINS,
CONTINUE,
Expand Down
30 changes: 28 additions & 2 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1237,6 +1237,23 @@ impl<'a> Parser<'a> {
}
}

//Select item operators
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
//Select item operators
/// Select item operators

for the doc comment can we describe what the function does?

fn parse_select_item_prefix_by_reserved_word(&mut self) -> Result<Option<Expr>, ParserError> {
if let Some(kw) = self.peek_one_of_keywords(
self.dialect
.get_reserved_keywords_for_select_item_operator(),
) {
if let TokenWithSpan {
span,
token: Token::Word(word),
} = self.expect_keyword(kw)?
{
return Ok(Some(Expr::Identifier(word.clone().into_ident(span))));
}
}
Ok(None)
}

/// Tries to parse an expression by matching the specified word to known keywords that have a special meaning in the dialect.
/// Returns `None if no match is found.
fn parse_expr_prefix_by_reserved_word(
Expand Down Expand Up @@ -13732,6 +13749,10 @@ impl<'a> Parser<'a> {

/// Parse a comma-delimited list of projections after SELECT
pub fn parse_select_item(&mut self) -> Result<SelectItem, ParserError> {
let prefix = self
.maybe_parse(|parser| parser.parse_select_item_prefix_by_reserved_word())?
.flatten();

match self.parse_wildcard_expr()? {
Comment on lines +13752 to 13756
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thinking introducing the prefix field into the select item is a bit more invasive and would be nice to avoid if it makes sense. I noticed this use case is similar in representation to IntroducedString, maybe we can repurpose that one to be generic.

I'm thinking something like this?

changing that enum variant into

Expr::Prefixed { prefix: Ident, expr: Expr }

Then impl wise here we could wrap the self.parse_wildcard_expr() call with the logic to optionally parse the prefix

fn parse_select_item_expr() {
  let prefix = self.parse_one_of_keywords(self.dialect.get_reserved_keywords_for_select_item());

  let expr = self.parse_wildcard_expr()?

  if let Some(prefix) = prefix {
    Expr::Prefixed {
      prefix: Ident::new(prefix),
      expr,
    }
  } else {
    expr
  }
}

Expr::QualifiedWildcard(prefix, token) => Ok(SelectItem::QualifiedWildcard(
SelectItemQualifiedWildcardKind::ObjectName(prefix),
Expand Down Expand Up @@ -13762,6 +13783,7 @@ impl<'a> Parser<'a> {
Ok(SelectItem::ExprWithAlias {
expr: *right,
alias,
prefix,
})
}
expr if self.dialect.supports_select_expr_star()
Expand All @@ -13776,8 +13798,12 @@ impl<'a> Parser<'a> {
expr => self
.maybe_parse_select_item_alias()
.map(|alias| match alias {
Some(alias) => SelectItem::ExprWithAlias { expr, alias },
None => SelectItem::UnnamedExpr(expr),
Some(alias) => SelectItem::ExprWithAlias {
expr,
alias,
prefix,
},
None => SelectItem::UnnamedExpr { expr, prefix },
}),
}
}
Expand Down
2 changes: 1 addition & 1 deletion src/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -311,7 +311,7 @@ pub fn only<T>(v: impl IntoIterator<Item = T>) -> T {

pub fn expr_from_projection(item: &SelectItem) -> &Expr {
match item {
SelectItem::UnnamedExpr(expr) => expr,
SelectItem::UnnamedExpr { expr, .. } => expr,
_ => panic!("Expected UnnamedExpr"),
}
}
Expand Down
19 changes: 11 additions & 8 deletions tests/sqlparser_bigquery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1626,14 +1626,17 @@ fn parse_hyphenated_table_identifiers() {
"SELECT foo - bar.x FROM t"
)
.projection[0],
SelectItem::UnnamedExpr(Expr::BinaryOp {
left: Box::new(Expr::Identifier(Ident::new("foo"))),
op: BinaryOperator::Minus,
right: Box::new(Expr::CompoundIdentifier(vec![
Ident::new("bar"),
Ident::new("x")
]))
})
SelectItem::UnnamedExpr {
expr: Expr::BinaryOp {
left: Box::new(Expr::Identifier(Ident::new("foo"))),
op: BinaryOperator::Minus,
right: Box::new(Expr::CompoundIdentifier(vec![
Ident::new("bar"),
Ident::new("x")
]))
},
prefix: None
}
);
}

Expand Down
65 changes: 42 additions & 23 deletions tests/sqlparser_clickhouse.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,22 +44,25 @@ fn parse_map_access_expr() {
select_token: AttachedToken::empty(),
top: None,
top_before_distinct: false,
projection: vec![UnnamedExpr(Expr::CompoundFieldAccess {
root: Box::new(Identifier(Ident {
value: "string_values".to_string(),
quote_style: None,
span: Span::empty(),
})),
access_chain: vec![AccessExpr::Subscript(Subscript::Index {
index: call(
"indexOf",
[
Expr::Identifier(Ident::new("string_names")),
Expr::value(Value::SingleQuotedString("endpoint".to_string()))
]
),
})],
})],
projection: vec![UnnamedExpr {
expr: Expr::CompoundFieldAccess {
root: Box::new(Identifier(Ident {
value: "string_values".to_string(),
quote_style: None,
span: Span::empty(),
})),
access_chain: vec![AccessExpr::Subscript(Subscript::Index {
index: call(
"indexOf",
[
Expr::Identifier(Ident::new("string_names")),
Expr::value(Value::SingleQuotedString("endpoint".to_string()))
]
),
})],
},
prefix: None
}],
into: None,
from: vec![TableWithJoins {
relation: table_from_name(ObjectName::from(vec![Ident::new("foos")])),
Expand Down Expand Up @@ -205,7 +208,11 @@ fn parse_delimited_identifiers() {
expr_from_projection(&select.projection[1]),
);
match &select.projection[2] {
SelectItem::ExprWithAlias { expr, alias } => {
SelectItem::ExprWithAlias {
expr,
alias,
prefix: _,
} => {
assert_eq!(&Expr::Identifier(Ident::with_quote('"', "simple id")), expr);
assert_eq!(&Ident::with_quote('"', "column alias"), alias);
}
Expand Down Expand Up @@ -315,8 +322,14 @@ fn parse_alter_table_add_projection() {
name: "my_name".into(),
select: ProjectionSelect {
projection: vec![
UnnamedExpr(Identifier(Ident::new("a"))),
UnnamedExpr(Identifier(Ident::new("b"))),
UnnamedExpr {
expr: Identifier(Ident::new("a")),
prefix: None
},
UnnamedExpr {
expr: Identifier(Ident::new("b")),
prefix: None
},
],
group_by: Some(GroupByExpr::Expressions(
vec![Identifier(Ident::new("a"))],
Expand Down Expand Up @@ -1000,7 +1013,10 @@ fn parse_select_parametric_function() {
let projection: &Vec<SelectItem> = query.body.as_select().unwrap().projection.as_ref();
assert_eq!(projection.len(), 1);
match &projection[0] {
UnnamedExpr(Expr::Function(f)) => {
UnnamedExpr {
expr: Expr::Function(f),
..
} => {
let args = match &f.args {
FunctionArguments::List(ref args) => args,
_ => unreachable!(),
Expand Down Expand Up @@ -1418,9 +1434,12 @@ fn parse_create_table_on_commit_and_as_query() {
assert_eq!(on_commit, Some(OnCommit::PreserveRows));
assert_eq!(
query.unwrap().body.as_select().unwrap().projection,
vec![UnnamedExpr(Expr::Value(
(Value::Number("1".parse().unwrap(), false)).with_empty_span()
))]
vec![UnnamedExpr {
expr: Expr::Value(
(Value::Number("1".parse().unwrap(), false)).with_empty_span()
),
prefix: None
}]
);
}
_ => unreachable!(),
Expand Down
Loading