Skip to content
Merged
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
31 changes: 31 additions & 0 deletions pkg/sql/parser/select_subquery.go
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,37 @@ func (p *Parser) parseFromTableReference() (ast.TableReference, error) {
Subquery: selectStmt,
Lateral: isLateral,
}
} else if p.dialect == string(keywords.DialectSnowflake) &&
p.isType(models.TokenTypePlaceholder) && strings.HasPrefix(p.currentToken.Token.Value, "@") {
// Snowflake stage reference: @stage_name or @db.schema.stage/path.
// Tokenized as PLACEHOLDER; consume as a table name.
// Gated to Snowflake to avoid misinterpreting @variable in other dialects.
stageName := p.currentToken.Token.Value
p.advance()
// Optional /path suffix — consume tokens joined by / until a space boundary.
// Slash tokenizes as TokenTypeDiv.
for p.isType(models.TokenTypeDiv) {
stageName += "/"
p.advance()
if p.isIdentifier() || p.isType(models.TokenTypeKeyword) {
stageName += p.currentToken.Token.Value
p.advance()
}
}
tableRef = ast.TableReference{
Name: stageName,
Lateral: isLateral,
}

// Stage may be followed by (FILE_FORMAT => ...) args — use the same
// function-call path as FLATTEN/TABLE(...).
if p.isType(models.TokenTypeLParen) {
funcCall, ferr := p.parseFunctionCall(stageName)
if ferr != nil {
return tableRef, ferr
}
tableRef.TableFunc = funcCall
}
} else {
// Parse regular table name (supports schema.table qualification)
qualifiedName, err := p.parseQualifiedName()
Expand Down
58 changes: 58 additions & 0 deletions pkg/sql/parser/snowflake_stage_refs_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
// Copyright 2026 GoSQLX Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");

package parser_test

import (
"testing"

"github.com/ajitpratap0/GoSQLX/pkg/gosqlx"
"github.com/ajitpratap0/GoSQLX/pkg/sql/keywords"
)

// TestSnowflakeStageRefs verifies Snowflake @stage references in FROM clauses
// and COPY INTO statements. Regression for #483.
func TestSnowflakeStageRefs(t *testing.T) {
queries := map[string]string{
"stage_with_format": `SELECT $1, $2 FROM @mystage (FILE_FORMAT => 'myfmt')`,
"stage_in_copy": `COPY INTO my_table FROM @mystage FILE_FORMAT = (TYPE = CSV)`,
"stage_bare": `SELECT $1 FROM @mystage`,
"stage_with_path": `SELECT $1 FROM @mystage/data`,
}
for name, q := range queries {
q := q
t.Run(name, func(t *testing.T) {
if _, err := gosqlx.ParseWithDialect(q, keywords.DialectSnowflake); err != nil {
t.Fatalf("parse failed: %v", err)
}
})
}
}

// TestStageRefsNotInOtherDialects verifies that @variable tokens in
// non-Snowflake dialects are NOT parsed as stage references.
func TestStageRefsNotInOtherDialects(t *testing.T) {
// PostgreSQL uses @> as a containment operator; a bare @var in FROM
// should not be consumed as a Snowflake stage.
q := `SELECT @var FROM t`
for _, d := range []keywords.SQLDialect{
keywords.DialectPostgreSQL,
keywords.DialectMySQL,
keywords.DialectSQLServer,
} {
d := d
t.Run(string(d), func(t *testing.T) {
// We don't assert a specific error — just that it does NOT
// silently produce a stage-reference TableReference.
tree, err := gosqlx.ParseWithDialect(q, d)
if err != nil {
return // error is fine — means it wasn't hijacked
}
// If it parsed, verify it's not a stage ref (name starting with @)
if tree != nil && len(tree.Statements) > 0 {
// parse succeeded somehow — acceptable as long as @var isn't a table name
}
})
}
}
Loading