Skip to content

Commit

Permalink
trino: Avoid correlated subqueries in IN unnest
Browse files Browse the repository at this point in the history
We also add some tests for quoted identifiers, which do actually work
correctly in our lexer/parser, but for reasons I had totally forgotten.
We also needed to add a way to quote identifers in `sql_quote!`. I ended
up choosing Rust syntax.
  • Loading branch information
emk committed Nov 8, 2023
1 parent 0645dca commit bced1d5
Show file tree
Hide file tree
Showing 6 changed files with 47 additions and 13 deletions.
12 changes: 9 additions & 3 deletions joinery_macros/src/sql_quote.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,9 +64,15 @@ fn emit_sql_token_exprs(
}
TokenTree::Ident(ident) => {
let ident_str = ident.to_string();
sql_token_exprs.push(quote_spanned! { ident_str.span() =>
__tokens.push(Token::ident(#ident_str, Span::Unknown))
});
if let Some(ident_str) = ident_str.strip_prefix("r#") {
sql_token_exprs.push(quote_spanned! { ident_str.span() =>
__tokens.push(Token::quoted_ident(#ident_str, Span::Unknown))
});
} else {
sql_token_exprs.push(quote_spanned! { ident_str.span() =>
__tokens.push(Token::ident(#ident_str, Span::Unknown))
});
}
}
TokenTree::Punct(punct) => {
let punct_str = punct.to_string();
Expand Down
2 changes: 1 addition & 1 deletion src/drivers/trino/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ impl Driver for TrinoDriver {
Box::new(transforms::ArraySelectToSubquery),
Box::new(transforms::QualifyToSubquery),
Box::<transforms::ExpandExcept>::default(),
Box::new(transforms::InUnnestToInSelect),
Box::new(transforms::InUnnestToContains),
Box::new(transforms::CountifToCase),
Box::new(transforms::IndexFromOne),
Box::new(transforms::IsBoolToCase),
Expand Down
18 changes: 17 additions & 1 deletion src/tokenizer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,11 @@ impl Token {
Self::Ident(Ident::new(ident, span))
}

/// Construct a new quoted identifier token.
pub fn quoted_ident(ident: &str, span: Span) -> Self {
Self::Ident(Ident::new_quoted(ident, span))
}

/// Construct a new punctuation token.
pub fn punct(punct: &str, span: Span) -> Self {
Self::Punct(Punct::new(punct, span))
Expand Down Expand Up @@ -375,13 +380,24 @@ pub struct Ident {
}

impl Ident {
/// Create a new `Ident` with no source location.
/// Create a new `Ident`.
pub fn new(name: &str, span: Span) -> Self {
Self {
token: RawToken::new(name, span),
name: name.to_owned(),
}
}

/// Create a new quoted `Ident` that can't be mistaken for a keyword.
pub fn new_quoted(name: &str, span: Span) -> Self {
// Yes, the upper level parser tells these appart by looking at the
// uppercased `RawToken` and checking _that_ against the keyword list.
let quoted = format!("`{}`", name);
Self {
token: RawToken::new(&quoted, span),
name: name.to_owned(),
}
}
}

impl Spanned for Ident {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,22 +8,22 @@ use crate::{

use super::{Transform, TransformExtra};

/// Transform `val IN UNNEST(expr)` into `val IN (SELECT * FROM UNNEST(expr))`.
/// Transform `val IN UNNEST(expr)` into `CONTAINS(expr, val)`.
#[derive(VisitorMut)]
#[visitor(Expression(enter))]
pub struct InUnnestToInSelect;
pub struct InUnnestToContains;

impl InUnnestToInSelect {
impl InUnnestToContains {
fn enter_expression(&mut self, expr: &mut Expression) {
if let Expression::In(InExpression {
left,
not_token,
in_token,
value_set: InValueSet::Unnest { expression, .. },
..
}) = expr
{
let replacement = sql_quote! {
#left #not_token #in_token (SELECT * FROM UNNEST(#expression))
(#not_token r#CONTAINS(#expression, #left))
}
.try_into_expression()
.expect("generated SQL should always parse");
Expand All @@ -32,7 +32,7 @@ impl InUnnestToInSelect {
}
}

impl Transform for InUnnestToInSelect {
impl Transform for InUnnestToContains {
fn name(&self) -> &'static str {
"InUnnestToInSelect"
}
Expand Down
4 changes: 2 additions & 2 deletions src/transforms/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ pub use self::{
countif_to_case::CountifToCase,
expand_except::ExpandExcept,
if_to_case::IfToCase,
in_unnest_to_in_select::InUnnestToInSelect,
in_unnest_to_contains::InUnnestToContains,
index_from_one::IndexFromOne,
index_from_zero::IndexFromZero,
is_bool_to_case::IsBoolToCase,
Expand All @@ -31,7 +31,7 @@ mod clean_up_temp_manually;
mod countif_to_case;
mod expand_except;
mod if_to_case;
mod in_unnest_to_in_select;
mod in_unnest_to_contains;
mod index_from_one;
mod index_from_zero;
mod is_bool_to_case;
Expand Down
12 changes: 12 additions & 0 deletions tests/sql/lexical/quoted_identifiers.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
-- Quoted identifiers.

CREATE OR REPLACE TABLE __result1 AS
SELECT 'SELECT' AS `SELECT`, 'FROM' AS `FROM`, 'CONTAINS' AS `CONTAINS`;

CREATE OR REPLACE TABLE __expected1 (
`SELECT` STRING,
`FROM` STRING,
`CONTAINS` STRING,
);
INSERT INTO __expected1 VALUES
('SELECT', 'FROM', 'CONTAINS');

0 comments on commit bced1d5

Please sign in to comment.