Fixed select dollar column from stage for snowflake (#2165)

Author: romanoff

src/dialect/mod.rs

@@ -49,6 +49,7 @@ pub use self::mysql::MySqlDialect;
 pub use self::oracle::OracleDialect;
 pub use self::postgresql::PostgreSqlDialect;
 pub use self::redshift::RedshiftSqlDialect;
+pub use self::snowflake::parse_snowflake_stage_name;
 pub use self::snowflake::SnowflakeDialect;
 pub use self::sqlite::SQLiteDialect;
 

src/dialect/snowflake.rs

@@ -1230,7 +1230,7 @@ pub fn parse_stage_name_identifier(parser: &mut Parser) -> Result<Ident, ParserE
                 parser.prev_token();
                 break;
             }
-            Token::RParen => {
+            Token::LParen | Token::RParen => {
                 parser.prev_token();
                 break;
             }
@@ -1248,6 +1248,8 @@ pub fn parse_stage_name_identifier(parser: &mut Parser) -> Result<Ident, ParserE
     Ok(Ident::new(ident))
 }
 
+/// Parses a Snowflake stage name, which may start with `@` for internal stages.
+/// Examples: `@mystage`, `@namespace.stage`, `schema.table`
 pub fn parse_snowflake_stage_name(parser: &mut Parser) -> Result<ObjectName, ParserError> {
     match parser.next_token().token {
         Token::AtSign => {

src/parser/mod.rs

@@ -1284,6 +1284,11 @@ impl<'a> Parser<'a> {
                                 // SQLite has single-quoted identifiers
                                 id_parts.push(Ident::with_quote('\'', s))
                             }
+                            Token::Placeholder(s) => {
+                                // Snowflake uses $1, $2, etc. for positional column references
+                                // in staged data queries like: SELECT t.$1 FROM @stage t
+                                id_parts.push(Ident::new(s))
+                            }
                             Token::Mul => {
                                 return Ok(Expr::QualifiedWildcard(
                                     ObjectName::from(id_parts),
@@ -1946,6 +1951,13 @@ impl<'a> Parser<'a> {
                         chain.push(AccessExpr::Dot(expr));
                         self.advance_token(); // The consumed string
                     }
+                    Token::Placeholder(s) => {
+                        // Snowflake uses $1, $2, etc. for positional column references
+                        // in staged data queries like: SELECT t.$1 FROM @stage t
+                        let expr = Expr::Identifier(Ident::with_span(next_token.span, s));
+                        chain.push(AccessExpr::Dot(expr));
+                        self.advance_token(); // The consumed placeholder
+                    }
                     // Fallback to parsing an arbitrary expression, but restrict to expression
                     // types that are valid after the dot operator. This ensures that e.g.
                     // `T.interval` is parsed as a compound identifier, not as an interval
@@ -15435,6 +15447,9 @@ impl<'a> Parser<'a> {
             && self.peek_keyword_with_tokens(Keyword::SEMANTIC_VIEW, &[Token::LParen])
         {
             self.parse_semantic_view_table_factor()
+        } else if self.peek_token_ref().token == Token::AtSign {
+            // Stage reference: @mystage or @namespace.stage (e.g. Snowflake)
+            self.parse_snowflake_stage_table_factor()
         } else {
             let name = self.parse_object_name(true)?;
 
@@ -15531,6 +15546,37 @@ impl<'a> Parser<'a> {
         }
     }
 
+    /// Parse a Snowflake stage reference as a table factor.
+    /// Handles syntax like: `@mystage1 (file_format => 'myformat', pattern => '...')`
+    ///
+    /// See: <https://docs.snowflake.com/en/user-guide/querying-stage>
+    fn parse_snowflake_stage_table_factor(&mut self) -> Result<TableFactor, ParserError> {
+        // Parse the stage name starting with @
+        let name = crate::dialect::parse_snowflake_stage_name(self)?;
+
+        // Parse optional stage options like (file_format => 'myformat', pattern => '...')
+        let args = if self.consume_token(&Token::LParen) {
+            Some(self.parse_table_function_args()?)
+        } else {
+            None
+        };
+
+        let alias = self.maybe_parse_table_alias()?;
+
+        Ok(TableFactor::Table {
+            name,
+            alias,
+            args,
+            with_hints: vec![],
+            version: None,
+            partitions: vec![],
+            with_ordinality: false,
+            json_path: None,
+            sample: None,
+            index_hints: vec![],
+        })
+    }
+
     fn maybe_parse_table_sample(&mut self) -> Result<Option<Box<TableSample>>, ParserError> {
         let modifier = if self.parse_keyword(Keyword::TABLESAMPLE) {
             TableSampleModifier::TableSample

tests/sqlparser_snowflake.rs

@@ -4878,3 +4878,15 @@ fn test_truncate_table_if_exists() {
     snowflake().verified_stmt("TRUNCATE TABLE my_table");
     snowflake().verified_stmt("TRUNCATE IF EXISTS my_table");
 }
+
+#[test]
+fn test_select_dollar_column_from_stage() {
+    // With table function args and alias
+    snowflake().verified_stmt("SELECT t.$1, t.$2 FROM @mystage1(file_format => 'myformat', pattern => '.*data.*[.]csv.gz') t");
+    // Without table function args, with alias
+    snowflake().verified_stmt("SELECT t.$1, t.$2 FROM @mystage1 t");
+    // Without table function args, without alias
+    snowflake().verified_stmt("SELECT $1, $2 FROM @mystage1");
+    // With table function args, without alias
+    snowflake().verified_stmt("SELECT $1, $2 FROM @mystage1(file_format => 'myformat')");
+}