Skip to content

Commit b267953

Browse files
committed
Fixed select dollar column from stage for snowflake
1 parent 46f2234 commit b267953

4 files changed

Lines changed: 55 additions & 1 deletion

File tree

src/dialect/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ pub use self::mysql::MySqlDialect;
4949
pub use self::oracle::OracleDialect;
5050
pub use self::postgresql::PostgreSqlDialect;
5151
pub use self::redshift::RedshiftSqlDialect;
52+
pub use self::snowflake::parse_snowflake_stage_name;
5253
pub use self::snowflake::SnowflakeDialect;
5354
pub use self::sqlite::SQLiteDialect;
5455
use crate::ast::{ColumnOption, Expr, GranteesType, Ident, ObjectNamePart, Statement};

src/dialect/snowflake.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1149,7 +1149,7 @@ pub fn parse_stage_name_identifier(parser: &mut Parser) -> Result<Ident, ParserE
11491149
parser.prev_token();
11501150
break;
11511151
}
1152-
Token::RParen => {
1152+
Token::LParen | Token::RParen => {
11531153
parser.prev_token();
11541154
break;
11551155
}
@@ -1167,6 +1167,8 @@ pub fn parse_stage_name_identifier(parser: &mut Parser) -> Result<Ident, ParserE
11671167
Ok(Ident::new(ident))
11681168
}
11691169

1170+
/// Parses a Snowflake stage name, which may start with `@` for internal stages.
1171+
/// Examples: `@mystage`, `@namespace.stage`, `schema.table`
11701172
pub fn parse_snowflake_stage_name(parser: &mut Parser) -> Result<ObjectName, ParserError> {
11711173
match parser.next_token().token {
11721174
Token::AtSign => {

src/parser/mod.rs

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1273,6 +1273,11 @@ impl<'a> Parser<'a> {
12731273
// SQLite has single-quoted identifiers
12741274
id_parts.push(Ident::with_quote('\'', s))
12751275
}
1276+
Token::Placeholder(s) => {
1277+
// Snowflake uses $1, $2, etc. for positional column references
1278+
// in staged data queries like: SELECT t.$1 FROM @stage t
1279+
id_parts.push(Ident::new(s))
1280+
}
12761281
Token::Mul => {
12771282
return Ok(Expr::QualifiedWildcard(
12781283
ObjectName::from(id_parts),
@@ -1898,6 +1903,13 @@ impl<'a> Parser<'a> {
18981903
chain.push(AccessExpr::Dot(expr));
18991904
self.advance_token(); // The consumed string
19001905
}
1906+
Token::Placeholder(s) => {
1907+
// Snowflake uses $1, $2, etc. for positional column references
1908+
// in staged data queries like: SELECT t.$1 FROM @stage t
1909+
let expr = Expr::Identifier(Ident::with_span(next_token.span, s));
1910+
chain.push(AccessExpr::Dot(expr));
1911+
self.advance_token(); // The consumed placeholder
1912+
}
19011913
// Fallback to parsing an arbitrary expression.
19021914
_ => match self.parse_subexpr(self.dialect.prec_value(Precedence::Period))? {
19031915
// If we get back a compound field access or identifier,
@@ -15103,6 +15115,11 @@ impl<'a> Parser<'a> {
1510315115
&& self.peek_keyword_with_tokens(Keyword::SEMANTIC_VIEW, &[Token::LParen])
1510415116
{
1510515117
self.parse_semantic_view_table_factor()
15118+
} else if dialect_of!(self is SnowflakeDialect)
15119+
&& self.peek_token_ref().token == Token::AtSign
15120+
{
15121+
// Snowflake stage reference: @mystage or @namespace.stage
15122+
self.parse_snowflake_stage_table_factor()
1510615123
} else {
1510715124
let name = self.parse_object_name(true)?;
1510815125

@@ -15199,6 +15216,35 @@ impl<'a> Parser<'a> {
1519915216
}
1520015217
}
1520115218

15219+
/// Parse a Snowflake stage reference as a table factor.
15220+
/// Handles syntax like: `@mystage1 (file_format => 'myformat', pattern => '...')`
15221+
fn parse_snowflake_stage_table_factor(&mut self) -> Result<TableFactor, ParserError> {
15222+
// Parse the stage name starting with @
15223+
let name = crate::dialect::parse_snowflake_stage_name(self)?;
15224+
15225+
// Parse optional stage options like (file_format => 'myformat', pattern => '...')
15226+
let args = if self.consume_token(&Token::LParen) {
15227+
Some(self.parse_table_function_args()?)
15228+
} else {
15229+
None
15230+
};
15231+
15232+
let alias = self.maybe_parse_table_alias()?;
15233+
15234+
Ok(TableFactor::Table {
15235+
name,
15236+
alias,
15237+
args,
15238+
with_hints: vec![],
15239+
version: None,
15240+
partitions: vec![],
15241+
with_ordinality: false,
15242+
json_path: None,
15243+
sample: None,
15244+
index_hints: vec![],
15245+
})
15246+
}
15247+
1520215248
fn maybe_parse_table_sample(&mut self) -> Result<Option<Box<TableSample>>, ParserError> {
1520315249
let modifier = if self.parse_keyword(Keyword::TABLESAMPLE) {
1520415250
TableSampleModifier::TableSample

tests/sqlparser_snowflake.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4533,3 +4533,8 @@ fn test_alter_external_table() {
45334533
snowflake()
45344534
.verified_stmt("ALTER EXTERNAL TABLE IF EXISTS some_table REFRESH 'year=2025/month=12/'");
45354535
}
4536+
4537+
#[test]
4538+
fn test_select_dollar_column_from_stage() {
4539+
snowflake().verified_stmt("SELECT t.$1, t.$2 FROM @mystage1(file_format => 'myformat', pattern => '.*data.*[.]csv.gz') t");
4540+
}

0 commit comments

Comments
 (0)