feat(parser): parse PIVOT DEFAULT ON NULL, ClickHouse FINAL/SAMPLE, INSERT TOP, MySQL partition actions, ROWS FROM, array subscripts, nested GROUPING SETS

CHANGELOG.md

@@ -8,6 +8,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- Parse and faithfully render Snowflake PIVOT defaults, ClickHouse FINAL/SAMPLE
+  (also on joined tables, `TableModifiers.Sample`), SQL Server INSERT TOP, MySQL /
+  MariaDB ALTER TABLE partition actions (ADD/DROP [IF [NOT] EXISTS], TRUNCATE,
+  REORGANIZE, COALESCE, EXCHANGE, ANALYZE/CHECK/OPTIMIZE/REBUILD/REPAIR,
+  DISCARD/IMPORT, REMOVE/UPGRADE PARTITIONING, ALL), PostgreSQL ROWS FROM and
+  array subscripts, and ROLLUP/CUBE/nested GROUPING SETS inside GROUPING SETS;
+  cover identity constraints. Partition actions outside the structured grammar
+  still parse as as-written `AlterTableOther` text.
 - `ast.PivotClause.AggregateAlias`, `ast.PivotClause.ExtraAggregates`
   (`[]ast.PivotAggregate`), `ast.UnpivotClause.NullHandling`,
   `ast.UnpivotClause.ValueColumns` and `ast.UnpivotNullHandling`; rendered by
@@ -36,6 +44,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   identifier delimiters through rendering.
 
 ### Changed
+- Parsed ClickHouse SAMPLE sizes and offsets now use inspectable expression
+  fields, which take rendering precedence over the retained legacy string fields.
+- ClickHouse SAMPLE accepts only decimal numbers (hex such as `0x10` is
+  rejected) and renders before ARRAY JOIN; a SAMPLE clause with no valid size is
+  omitted instead of rendered as a bare `SAMPLE`.
+- A column with more than one of DEFAULT, an identity clause and a generation
+  expression is rejected; signed numbers in identity options render as `-1`.
 - **Dialect extension fields**: the new dialect-specific clauses live behind
   optional pointers so that common statements do not carry them:
   `SelectStatement.Extras` (`SelectExtras`: SELECT modifiers, INTO, LIMIT BY,

pkg/formatter/dialect_constructs_roundtrip_test.go

@@ -0,0 +1,149 @@
+package formatter_test
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/ajitpratap0/GoSQLX/pkg/formatter"
+	"github.com/ajitpratap0/GoSQLX/pkg/gosqlx"
+	"github.com/ajitpratap0/GoSQLX/pkg/models"
+	"github.com/ajitpratap0/GoSQLX/pkg/sql/ast"
+	"github.com/ajitpratap0/GoSQLX/pkg/sql/keywords"
+)
+
+func TestDialectConstructsRoundTripAndInspection(t *testing.T) {
+	cases := []struct {
+		dialect keywords.SQLDialect
+		sql     string
+	}{
+		{keywords.DialectSnowflake, `SELECT * FROM t PIVOT (SUM(amt) FOR month IN (ANY ORDER BY month) DEFAULT ON NULL (0)) AS p`},
+		{keywords.DialectSnowflake, `SELECT * FROM t PIVOT (SUM(amt) FOR month IN ('JAN', 'FEB') DEFAULT ON NULL (0)) AS p`},
+		{keywords.DialectClickHouse, `SELECT * FROM hits FINAL SAMPLE 0.1 WHERE x = 1`},
+		{keywords.DialectClickHouse, `SELECT * FROM hits FINAL SAMPLE 1/10 OFFSET 1/2 WHERE x = 1`},
+		{keywords.DialectClickHouse, `SELECT * FROM hits FINAL`},
+		{keywords.DialectClickHouse, `SELECT * FROM hits SAMPLE 0.1`},
+		{keywords.DialectClickHouse, `SELECT * FROM hits AS h FINAL SAMPLE 0.1`},
+		{keywords.DialectSQLServer, `INSERT TOP (10) INTO t (a) SELECT a FROM s`},
+		{keywords.DialectSQLServer, `INSERT TOP (10) PERCENT INTO t (a) SELECT a FROM s`},
+		{keywords.DialectSQLServer, `INSERT TOP (@n) t (a) SELECT a FROM s`},
+		{keywords.DialectSQLServer, `INSERT "TOP" (a) SELECT a FROM s`},
+		{keywords.DialectPostgreSQL, `CREATE TABLE t (id bigint GENERATED ALWAYS AS IDENTITY (START WITH 10 INCREMENT BY 5) PRIMARY KEY, b int GENERATED BY DEFAULT AS IDENTITY NOT NULL, c int NOT NULL GENERATED ALWAYS AS IDENTITY)`},
+		{keywords.DialectMySQL, `ALTER TABLE t PARTITION BY RANGE (YEAR(d)) (PARTITION p0 VALUES LESS THAN (2020), PARTITION p1 VALUES LESS THAN MAXVALUE)`},
+		{keywords.DialectMySQL, `ALTER TABLE t ADD PARTITION (PARTITION p2 VALUES LESS THAN (2030))`},
+		{keywords.DialectMySQL, `ALTER TABLE t DROP PARTITION p0`},
+		{keywords.DialectMySQL, `ALTER TABLE t TRUNCATE PARTITION p1`},
+		{keywords.DialectMySQL, `ALTER TABLE t REORGANIZE PARTITION p1 INTO (PARTITION p2 VALUES LESS THAN (2030), PARTITION p3 VALUES LESS THAN MAXVALUE)`},
+		{keywords.DialectMySQL, `ALTER TABLE t REMOVE PARTITIONING`},
+		{keywords.DialectPostgreSQL, `SELECT * FROM ROWS FROM (generate_series(1,3), unnest(ARRAY['a','b'])) WITH ORDINALITY AS r(n, s, i)`},
+		{keywords.DialectPostgreSQL, `SELECT * FROM ROWS FROM (f(x) AS (a int, b text))`},
+		{keywords.DialectPostgreSQL, `SELECT (ARRAY[1,2,3])[2]`},
+		{keywords.DialectPostgreSQL, `SELECT ARRAY[1,2,3][2]`},
+		{keywords.DialectPostgreSQL, `SELECT ARRAY[[1,2],[3,4]][1][2]`},
+		{keywords.DialectPostgreSQL, `SELECT ARRAY[1,2,3][1:2]`},
+	}
+	for _, dialect := range []keywords.SQLDialect{keywords.DialectPostgreSQL, keywords.DialectSnowflake, keywords.DialectSQLServer, keywords.DialectOracle} {
+		cases = append(cases, struct {
+			dialect keywords.SQLDialect
+			sql     string
+		}{dialect, `SELECT a, b, SUM(c) FROM t GROUP BY GROUPING SETS ((a), ROLLUP(b), CUBE(a, b), ())`})
+	}
+	for _, tc := range cases {
+		t.Run(string(tc.dialect)+"/"+tc.sql, func(t *testing.T) {
+			tree, err := gosqlx.ParseWithDialect(tc.sql, tc.dialect)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer ast.ReleaseAST(tree)
+			checkDialectConstructInspection(t, tree)
+			formatted, err := gosqlx.FormatWithDialect(tc.sql, gosqlx.DefaultFormatOptions(), tc.dialect)
+			if err != nil {
+				t.Fatal(err)
+			}
+			outputs := []string{tree.SQL(), formatted, formatter.FormatAST(tree, ast.CompactStyle()), formatter.FormatAST(tree, ast.ReadableStyle())}
+			clearDialectConstructLocations(reflect.ValueOf(tree.Statements))
+			for _, output := range outputs {
+				again, err := gosqlx.ParseWithDialect(output, tc.dialect)
+				if err != nil {
+					t.Fatalf("reparse %q: %v", output, err)
+				}
+				defer ast.ReleaseAST(again)
+				clearDialectConstructLocations(reflect.ValueOf(again.Statements))
+				if !reflect.DeepEqual(tree.Statements, again.Statements) {
+					t.Errorf("AST changed after rendering %q", output)
+				}
+				checkDialectConstructInspection(t, again)
+			}
+		})
+	}
+}
+
+// Source positions change with whitespace; every other field must be identical.
+func clearDialectConstructLocations(v reflect.Value) {
+	if !v.IsValid() {
+		return
+	}
+	if v.Type() == reflect.TypeFor[models.Location]() && v.CanSet() {
+		v.SetZero()
+		return
+	}
+	switch v.Kind() {
+	case reflect.Interface, reflect.Pointer:
+		if !v.IsNil() {
+			clearDialectConstructLocations(v.Elem())
+		}
+	case reflect.Struct:
+		for i := 0; i < v.NumField(); i++ {
+			if v.Type().Field(i).IsExported() {
+				clearDialectConstructLocations(v.Field(i))
+			}
+		}
+	case reflect.Slice:
+		for i := 0; i < v.Len(); i++ {
+			clearDialectConstructLocations(v.Index(i))
+		}
+	}
+}
+
+// Compare traversal against the fields themselves so a missing Children edge
+// cannot hide an expression from the acceptance tests.
+func checkDialectConstructInspection(t *testing.T, tree *ast.AST) {
+	t.Helper()
+	visited := map[ast.Expression]bool{}
+	ast.Inspect(tree, func(node ast.Node) bool {
+		if e, ok := node.(ast.Expression); ok {
+			visited[e] = true
+		}
+		return true
+	})
+	var check func(reflect.Value)
+	check = func(v reflect.Value) {
+		if !v.IsValid() {
+			return
+		}
+		if v.Kind() == reflect.Pointer && v.IsNil() {
+			return
+		}
+		if v.CanInterface() {
+			if e, ok := v.Interface().(ast.Expression); ok && !visited[e] {
+				t.Errorf("expression %T (%s) is not reachable through Inspect", e, e.TokenLiteral())
+			}
+		}
+		switch v.Kind() {
+		case reflect.Interface, reflect.Pointer:
+			if !v.IsNil() {
+				check(v.Elem())
+			}
+		case reflect.Struct:
+			for i := 0; i < v.NumField(); i++ {
+				if v.Type().Field(i).IsExported() {
+					check(v.Field(i))
+				}
+			}
+		case reflect.Slice:
+			for i := 0; i < v.Len(); i++ {
+				check(v.Index(i))
+			}
+		}
+	}
+	check(reflect.ValueOf(tree.Statements))
+}

pkg/formatter/render.go

@@ -301,14 +301,15 @@ func renderSelect(s *ast.SelectStatement, opts ast.FormatOptions) string {
 		sb.WriteString(joinSQL(&j, f))
 	}
 
-	if s.ArrayJoin != nil {
+	// ClickHouse writes SAMPLE with the table expression, before ARRAY JOIN.
+	if sample := ast.SampleExpressionSQL(s.Sample, f.kw); sample != "" {
 		sb.WriteString(f.clauseSep())
-		sb.WriteString(arrayJoinSQL(s.ArrayJoin, f))
+		sb.WriteString(sample)
 	}
 
-	if s.Sample != nil {
+	if s.ArrayJoin != nil {
 		sb.WriteString(f.clauseSep())
-		sb.WriteString(sampleSQL(s.Sample, f))
+		sb.WriteString(arrayJoinSQL(s.ArrayJoin, f))
 	}
 
 	if s.PrewhereClause != nil {
@@ -1179,6 +1180,8 @@ func tableRefSQL(t *ast.TableReference, f *nodeFormatter) string {
 		sb.WriteString("(")
 		sb.WriteString(f.tableSubquerySQL(t.Subquery))
 		sb.WriteString(")")
+	} else if rows := ast.RowsFromSQL(t, f.kw); rows != "" {
+		sb.WriteString(rows)
 	} else if t.TableFunc != nil {
 		sb.WriteString(f.exprSQL(t.TableFunc))
 		if t.WithOrdinality {
@@ -1241,33 +1244,6 @@ func tableRefSQL(t *ast.TableReference, f *nodeFormatter) string {
 	return sb.String()
 }
 
-// sampleSQL renders a ClickHouse SAMPLE clause.
-func sampleSQL(s *ast.SampleClause, f *nodeFormatter) string {
-	var sb strings.Builder
-	sb.WriteString(f.kw("SAMPLE"))
-	sb.WriteString(" ")
-	// The ratio and offset parts are numeric literals; anything else can only
-	// come from a hand-built AST and is dropped rather than written raw.
-	if ast.IsPlainNumberText(s.Value) {
-		sb.WriteString(s.Value)
-		if ast.IsPlainNumberText(s.Denominator) {
-			sb.WriteString("/")
-			sb.WriteString(s.Denominator)
-		}
-	}
-	if ast.IsPlainNumberText(s.Offset) {
-		sb.WriteString(" ")
-		sb.WriteString(f.kw("OFFSET"))
-		sb.WriteString(" ")
-		sb.WriteString(s.Offset)
-		if ast.IsPlainNumberText(s.OffsetDenominator) {
-			sb.WriteString("/")
-			sb.WriteString(s.OffsetDenominator)
-		}
-	}
-	return sb.String()
-}
-
 // joinSQL renders a JOIN clause.
 func joinSQL(j *ast.JoinClause, f *nodeFormatter) string {
 	var sb strings.Builder

pkg/formatter/render_dml.go

@@ -139,6 +139,7 @@ func writeInsertHead(sb *strings.Builder, f *nodeFormatter, i *ast.InsertStateme
 	sb.WriteString(f.kw("INSERT"))
 	sb.WriteString(" ")
 	writeDMLModifiers(sb, f, i.DialectExtras().Modifiers)
+	writeDMLTop(sb, f, i.DialectExtras().Top)
 	sb.WriteString(f.kw("INTO"))
 	sb.WriteString(" ")
 	if i.DialectExtras().TableFunction != nil {

pkg/formatter/render_expression_case.go

@@ -250,12 +250,33 @@ func (f *nodeFormatter) exprSQL(e ast.Expression) string {
 		copy.Expressions = c.exprs(v.Expressions)
 		return c.result(&copy)
 	case *ast.GroupingSetsExpression:
-		copy := *v
-		copy.Sets = make([][]ast.Expression, len(v.Sets))
-		for i := range v.Sets {
-			copy.Sets[i] = c.exprs(v.Sets[i])
+		// Rendered here rather than through slots: SQL() leaves a set that is
+		// a lone ROLLUP, CUBE or GROUPING SETS unparenthesised, and a slot
+		// would hide that element's type from it.
+		var b strings.Builder
+		b.WriteString("grouping sets(")
+		for i, set := range v.Sets {
+			if i > 0 {
+				b.WriteString(", ")
+			}
+			if len(set) == 1 {
+				switch set[0].(type) {
+				case *ast.RollupExpression, *ast.CubeExpression, *ast.GroupingSetsExpression:
+					b.WriteString(f.exprSQL(set[0]))
+					continue
+				}
+			}
+			b.WriteByte('(')
+			for j, e := range set {
+				if j > 0 {
+					b.WriteString(", ")
+				}
+				b.WriteString(f.exprSQL(e))
+			}
+			b.WriteByte(')')
 		}
-		return c.result(&copy)
+		b.WriteByte(')')
+		return b.String()
 	case *ast.OuterJoinMarker:
 		copy := *v
 		copy.Expr = c.expr(v.Expr)

pkg/formatter/sample_expression_safety_test.go

@@ -0,0 +1,64 @@
+package formatter_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/ajitpratap0/GoSQLX/pkg/formatter"
+	"github.com/ajitpratap0/GoSQLX/pkg/gosqlx"
+	"github.com/ajitpratap0/GoSQLX/pkg/sql/ast"
+	"github.com/ajitpratap0/GoSQLX/pkg/sql/keywords"
+)
+
+func TestDialectConstructsSampleFormatterSafety(t *testing.T) {
+	evil := "1; DROP TABLE audit; --"
+	number := func(v string) ast.Expression { return &ast.LiteralValue{Value: v, Type: "int"} }
+	var nilRatio *ast.BinaryExpression
+	clauses := []*ast.SampleClause{
+		{Size: number(evil)},
+		{Size: number("1"), OffsetSize: number(evil)},
+		{Size: &ast.BinaryExpression{Left: number(evil), Operator: "/", Right: number("10")}},
+		{Size: &ast.BinaryExpression{Left: number("1"), Operator: "/", Right: number(evil)}},
+		{Size: number("1"), OffsetSize: &ast.BinaryExpression{Left: number(evil), Operator: "/", Right: number("2")}},
+		{Size: number("1"), OffsetSize: &ast.BinaryExpression{Left: number("1"), Operator: "/", Right: number(evil)}},
+		{Size: nilRatio},
+		{Size: number("1"), OffsetSize: nilRatio},
+	}
+	for i, clause := range clauses {
+		tree := &ast.AST{Statements: []ast.Statement{&ast.SelectStatement{Columns: []ast.Expression{&ast.Identifier{Name: "x"}}, From: []ast.TableReference{{Name: "hits"}}, Sample: clause}}}
+		for _, output := range []string{tree.SQL(), formatter.FormatAST(tree, ast.CompactStyle())} {
+			if strings.Contains(output, "DROP TABLE") || strings.Contains(output, ";") || strings.Contains(output, "--") {
+				t.Errorf("case %d emitted unsafe text: %q", i, output)
+			}
+		}
+	}
+}
+
+func TestDialectConstructsSampleCanonicalMutation(t *testing.T) {
+	tree, err := gosqlx.ParseWithDialect("SELECT * FROM hits SAMPLE 1/10 OFFSET 1/2", keywords.DialectClickHouse)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ast.ReleaseAST(tree)
+	stmt, ok := tree.Statements[0].(*ast.SelectStatement)
+	if !ok {
+		t.Fatalf("got %T", tree.Statements[0])
+	}
+	for _, output := range []string{tree.SQL(), formatter.FormatAST(tree, ast.CompactStyle())} {
+		if !strings.Contains(output, "SAMPLE 1/10 OFFSET 1/2") {
+			t.Fatalf("normal ratio changed: %q", output)
+		}
+	}
+	stmt.Sample.Size = &ast.LiteralValue{Value: "0.25", Type: "float"}
+	stmt.Sample.OffsetSize = &ast.LiteralValue{Value: "0.5", Type: "FLOAT"}
+	for _, output := range []string{tree.SQL(), formatter.FormatAST(tree, ast.CompactStyle())} {
+		if !strings.Contains(output, "SAMPLE 0.25 OFFSET 0.5") {
+			t.Fatalf("canonical mutation ignored: %q", output)
+		}
+		again, err := gosqlx.ParseWithDialect(output, keywords.DialectClickHouse)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer ast.ReleaseAST(again)
+	}
+}

pkg/formatter/table_factor_render_test.go

@@ -319,12 +319,18 @@ func TestRender_TableReferenceRendersEveryField(t *testing.T) {
 func TestRender_TableModifiersRendersEveryField(t *testing.T) {
 	num := func(v string) *ast.LiteralValue { return &ast.LiteralValue{Value: v, Type: "int"} }
 	samples := map[string]func(*ast.TableModifiers){
+		"RowsFrom": func(m *ast.TableModifiers) {
+			m.RowsFrom = []ast.RowsFromItem{{Function: &ast.FunctionCall{Name: "generate_series", Arguments: []ast.Expression{num("1"), num("3")}}}}
+		},
 		"Values": func(m *ast.TableModifiers) {
 			m.Values = &ast.Values{Rows: [][]ast.Expression{{num("1")}}}
 		},
 		"TableSample": func(m *ast.TableModifiers) {
 			m.TableSample = &ast.TableSampleClause{Keyword: "SAMPLE", Args: []ast.Expression{num("5")}}
 		},
+		"Sample": func(m *ast.TableModifiers) {
+			m.Sample = &ast.SampleClause{Size: &ast.LiteralValue{Value: "0.5", Type: "float"}}
+		},
 		"Flashback":  func(m *ast.TableModifiers) { m.Flashback = &ast.FlashbackClause{Kind: "SCN", Point: num("9")} },
 		"DBLink":     func(m *ast.TableModifiers) { m.DBLink = "lnk" },
 		"Partitions": func(m *ast.TableModifiers) { m.Partitions = []string{"p0"} },