feat(parser): Snowflake SAMPLE/TABLESAMPLE clause (#483)

Author: Ajit Pratap Singh

pkg/sql/parser/pivot.go

@@ -374,3 +374,22 @@ func (p *Parser) isSnowflakeTimeTravelStart() bool {
 	}
 	return false
 }
+
+// isSampleKeyword returns true if the current token is SAMPLE or TABLESAMPLE
+// followed by '(' or a sampling-method keyword, indicating a sampling clause
+// rather than a table alias. Used to prevent the FROM-alias parser from
+// consuming these tokens.
+func (p *Parser) isSampleKeyword() bool {
+	upper := strings.ToUpper(p.currentToken.Token.Value)
+	if upper != "SAMPLE" && upper != "TABLESAMPLE" {
+		return false
+	}
+	// Require '(' or a method keyword as lookahead to disambiguate from
+	// a table actually named "sample".
+	next := p.peekToken().Token
+	if next.Type == models.TokenTypeLParen {
+		return true
+	}
+	nextUpper := strings.ToUpper(next.Value)
+	return nextUpper == "BERNOULLI" || nextUpper == "SYSTEM" || nextUpper == "BLOCK" || nextUpper == "ROW"
+}

pkg/sql/parser/select_subquery.go

@@ -95,6 +95,41 @@ func (p *Parser) parseFromTableReference() (ast.TableReference, error) {
 			tableRef.TableFunc = funcCall
 		}
 
+		// Snowflake / ANSI SAMPLE or TABLESAMPLE clause on a table reference:
+		//   SAMPLE [BERNOULLI | SYSTEM | BLOCK | ROW] (N [ROWS])
+		//   TABLESAMPLE [method] (N [ROWS])
+		// Consume permissively — the method and paren block are consumed
+		// but not yet modeled on the AST.
+		if strings.EqualFold(p.currentToken.Token.Value, "SAMPLE") ||
+			strings.EqualFold(p.currentToken.Token.Value, "TABLESAMPLE") {
+			p.advance() // SAMPLE / TABLESAMPLE
+			// Optional method name
+			upper := strings.ToUpper(p.currentToken.Token.Value)
+			if upper == "BERNOULLI" || upper == "SYSTEM" || upper == "BLOCK" || upper == "ROW" {
+				p.advance()
+			}
+			// (N [ROWS]) block
+			if p.isType(models.TokenTypeLParen) {
+				depth := 0
+				for {
+					t := p.currentToken.Token.Type
+					if t == models.TokenTypeEOF {
+						break
+					}
+					if t == models.TokenTypeLParen {
+						depth++
+					} else if t == models.TokenTypeRParen {
+						depth--
+						if depth == 0 {
+							p.advance()
+							break
+						}
+					}
+					p.advance()
+				}
+			}
+		}
+
 		// Snowflake time-travel / change-tracking clauses:
 		//   AT (TIMESTAMP => ...)
 		//   BEFORE (STATEMENT => ...)
@@ -113,7 +148,7 @@ func (p *Parser) parseFromTableReference() (ast.TableReference, error) {
 	// Similarly, START followed by WITH is a hierarchical query seed, not an alias.
 	// Don't consume PIVOT/UNPIVOT as a table alias — they are contextual
 	// keywords in SQL Server/Oracle and must reach the pivot-clause parser below.
-	if (p.isIdentifier() || p.isType(models.TokenTypeAs)) && !p.isMariaDBClauseStart() && !p.isPivotKeyword() && !p.isUnpivotKeyword() && !p.isQualifyKeyword() && !p.isMinusSetOp() && !p.isSnowflakeTimeTravelStart() {
+	if (p.isIdentifier() || p.isType(models.TokenTypeAs)) && !p.isMariaDBClauseStart() && !p.isPivotKeyword() && !p.isUnpivotKeyword() && !p.isQualifyKeyword() && !p.isMinusSetOp() && !p.isSnowflakeTimeTravelStart() && !p.isSampleKeyword() {
 		if p.isType(models.TokenTypeAs) {
 			p.advance() // Consume AS
 			if !p.isIdentifier() {
@@ -227,7 +262,7 @@ func (p *Parser) parseJoinedTableRef(joinType string) (ast.TableReference, error
 	// Similarly, START followed by WITH is a hierarchical query seed, not an alias.
 	// Don't consume PIVOT/UNPIVOT as a table alias — they are contextual
 	// keywords in SQL Server/Oracle and must reach the pivot-clause parser below.
-	if (p.isIdentifier() || p.isType(models.TokenTypeAs)) && !p.isMariaDBClauseStart() && !p.isPivotKeyword() && !p.isUnpivotKeyword() && !p.isQualifyKeyword() && !p.isMinusSetOp() && !p.isSnowflakeTimeTravelStart() {
+	if (p.isIdentifier() || p.isType(models.TokenTypeAs)) && !p.isMariaDBClauseStart() && !p.isPivotKeyword() && !p.isUnpivotKeyword() && !p.isQualifyKeyword() && !p.isMinusSetOp() && !p.isSnowflakeTimeTravelStart() && !p.isSampleKeyword() {
 		if p.isType(models.TokenTypeAs) {
 			p.advance()
 			if !p.isIdentifier() {

pkg/sql/parser/snowflake_sample_test.go

@@ -0,0 +1,33 @@
+// 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"
+)
+
+// TestSnowflakeSample verifies Snowflake SAMPLE / TABLESAMPLE clause on
+// table references. Regression for #483.
+func TestSnowflakeSample(t *testing.T) {
+	queries := map[string]string{
+		"sample_pct":         `SELECT * FROM users SAMPLE (10)`,
+		"tablesample_rows":   `SELECT * FROM users TABLESAMPLE (100 ROWS)`,
+		"sample_bernoulli":   `SELECT * FROM users SAMPLE BERNOULLI (5)`,
+		"tablesample_system": `SELECT * FROM users TABLESAMPLE SYSTEM (1)`,
+		"sample_block":       `SELECT * FROM users SAMPLE BLOCK (10)`,
+		"sample_with_where":  `SELECT * FROM users SAMPLE (50) WHERE id > 100`,
+	}
+	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)
+			}
+		})
+	}
+}