Snowflake: Improve accuracy of lookahead in implicit LIMIT alias (#1941)

This commit is contained in:
Yoav Cohen 2025-07-17 11:08:47 +03:00 committed by GitHub
parent 4d9338638f
commit 92db20673b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 45 additions and 11 deletions

View file

@ -23,8 +23,8 @@ use crate::ast::helpers::stmt_data_loading::{
FileStagingCommand, StageLoadSelectItem, StageLoadSelectItemKind, StageParamsObject,
};
use crate::ast::{
ColumnOption, ColumnPolicy, ColumnPolicyProperty, CopyIntoSnowflakeKind, Ident,
IdentityParameters, IdentityProperty, IdentityPropertyFormatKind, IdentityPropertyKind,
ColumnOption, ColumnPolicy, ColumnPolicyProperty, CopyIntoSnowflakeKind, DollarQuotedString,
Ident, IdentityParameters, IdentityProperty, IdentityPropertyFormatKind, IdentityPropertyKind,
IdentityPropertyOrder, ObjectName, ObjectNamePart, RowAccessPolicy, ShowObjects, SqlOption,
Statement, TagsColumnOption, WrappedCollection,
};
@ -307,22 +307,22 @@ impl Dialect for SnowflakeDialect {
// they are not followed by other tokens that may change their meaning
// e.g. `SELECT * EXCEPT (col1) FROM tbl`
Keyword::EXCEPT
// e.g. `SELECT 1 LIMIT 5`
| Keyword::LIMIT
// e.g. `SELECT 1 OFFSET 5 ROWS`
| Keyword::OFFSET
// e.g. `INSERT INTO t SELECT 1 RETURNING *`
| Keyword::RETURNING if !matches!(parser.peek_token_ref().token, Token::Comma | Token::EOF) =>
{
false
}
// e.g. `SELECT 1 LIMIT 5` - not an alias
// e.g. `SELECT 1 OFFSET 5 ROWS` - not an alias
Keyword::LIMIT | Keyword::OFFSET if peek_for_limit_options(parser) => false,
// `FETCH` can be considered an alias as long as it's not followed by `FIRST`` or `NEXT`
// which would give it a different meanings, for example:
// `SELECT 1 FETCH FIRST 10 ROWS` - not an alias
// `SELECT 1 FETCH 10` - not an alias
Keyword::FETCH if parser.peek_one_of_keywords(&[Keyword::FIRST, Keyword::NEXT]).is_some()
|| matches!(parser.peek_token().token, Token::Number(_, _)) =>
|| peek_for_limit_options(parser) =>
{
false
}
@ -351,20 +351,23 @@ impl Dialect for SnowflakeDialect {
match kw {
// The following keywords can be considered an alias as long as
// they are not followed by other tokens that may change their meaning
Keyword::LIMIT
| Keyword::RETURNING
Keyword::RETURNING
| Keyword::INNER
| Keyword::USING
| Keyword::PIVOT
| Keyword::UNPIVOT
| Keyword::EXCEPT
| Keyword::MATCH_RECOGNIZE
| Keyword::OFFSET
if !matches!(parser.peek_token_ref().token, Token::SemiColon | Token::EOF) =>
{
false
}
// `LIMIT` can be considered an alias as long as it's not followed by a value. For example:
// `SELECT * FROM tbl LIMIT WHERE 1=1` - alias
// `SELECT * FROM tbl LIMIT 3` - not an alias
Keyword::LIMIT | Keyword::OFFSET if peek_for_limit_options(parser) => false,
// `FETCH` can be considered an alias as long as it's not followed by `FIRST`` or `NEXT`
// which would give it a different meanings, for example:
// `SELECT * FROM tbl FETCH FIRST 10 ROWS` - not an alias
@ -373,7 +376,7 @@ impl Dialect for SnowflakeDialect {
if parser
.peek_one_of_keywords(&[Keyword::FIRST, Keyword::NEXT])
.is_some()
|| matches!(parser.peek_token().token, Token::Number(_, _)) =>
|| peek_for_limit_options(parser) =>
{
false
}
@ -387,6 +390,7 @@ impl Dialect for SnowflakeDialect {
{
false
}
Keyword::GLOBAL if parser.peek_keyword(Keyword::FULL) => false,
// Reserved keywords by the Snowflake dialect, which seem to be less strictive
@ -472,6 +476,18 @@ impl Dialect for SnowflakeDialect {
}
}
// Peeks ahead to identify tokens that are expected after
// a LIMIT/FETCH keyword.
fn peek_for_limit_options(parser: &Parser) -> bool {
match &parser.peek_token_ref().token {
Token::Number(_, _) | Token::Placeholder(_) => true,
Token::SingleQuotedString(val) if val.is_empty() => true,
Token::DollarQuotedString(DollarQuotedString { value, .. }) if value.is_empty() => true,
Token::Word(w) if w.keyword == Keyword::NULL => true,
_ => false,
}
}
fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result<Statement, ParserError> {
let stage = parse_snowflake_stage_name(parser)?;
let pattern = if parser.parse_keyword(Keyword::PATTERN) {

View file

@ -3535,6 +3535,15 @@ fn test_sql_keywords_as_select_item_aliases() {
.parse_sql_statements(&format!("SELECT 1 {kw}"))
.is_err());
}
// LIMIT is alias
snowflake().one_statement_parses_to("SELECT 1 LIMIT", "SELECT 1 AS LIMIT");
// LIMIT is not an alias
snowflake().verified_stmt("SELECT 1 LIMIT 1");
snowflake().verified_stmt("SELECT 1 LIMIT $1");
snowflake().verified_stmt("SELECT 1 LIMIT ''");
snowflake().verified_stmt("SELECT 1 LIMIT NULL");
snowflake().verified_stmt("SELECT 1 LIMIT $$$$");
}
#[test]
@ -3586,6 +3595,15 @@ fn test_sql_keywords_as_table_aliases() {
.parse_sql_statements(&format!("SELECT * FROM tbl {kw}"))
.is_err());
}
// LIMIT is alias
snowflake().one_statement_parses_to("SELECT * FROM tbl LIMIT", "SELECT * FROM tbl AS LIMIT");
// LIMIT is not an alias
snowflake().verified_stmt("SELECT * FROM tbl LIMIT 1");
snowflake().verified_stmt("SELECT * FROM tbl LIMIT $1");
snowflake().verified_stmt("SELECT * FROM tbl LIMIT ''");
snowflake().verified_stmt("SELECT * FROM tbl LIMIT NULL");
snowflake().verified_stmt("SELECT * FROM tbl LIMIT $$$$");
}
#[test]