Add alias support in CAST, SUBSTRING, and TRIM functions (#107)

Author: kyleconroy

ast/ast.go

@@ -73,6 +73,7 @@ type SelectQuery struct {
 	OrderBy     []*OrderByElement     `json:"order_by,omitempty"`
 	Limit            Expression            `json:"limit,omitempty"`
 	LimitBy          []Expression          `json:"limit_by,omitempty"`
+	LimitByLimit     Expression            `json:"limit_by_limit,omitempty"`     // LIMIT value before BY (e.g., LIMIT 1 BY x LIMIT 3)
 	LimitByHasLimit  bool                  `json:"limit_by_has_limit,omitempty"` // true if LIMIT BY was followed by another LIMIT
 	Offset           Expression            `json:"offset,omitempty"`
 	Settings           []*SettingExpr        `json:"settings,omitempty"`
@@ -126,9 +127,10 @@ func (t *TablesInSelectQuery) End() token.Position { return t.Position }
 
 // TablesInSelectQueryElement represents a single table element in a SELECT.
 type TablesInSelectQueryElement struct {
-	Position token.Position   `json:"-"`
-	Table    *TableExpression `json:"table"`
-	Join     *TableJoin       `json:"join,omitempty"`
+	Position  token.Position    `json:"-"`
+	Table     *TableExpression  `json:"table,omitempty"`
+	Join      *TableJoin        `json:"join,omitempty"`
+	ArrayJoin *ArrayJoinClause  `json:"array_join,omitempty"` // For ARRAY JOIN as table element
 }
 
 func (t *TablesInSelectQueryElement) Pos() token.Position { return t.Position }
@@ -225,6 +227,7 @@ type InsertQuery struct {
 	Table       string           `json:"table,omitempty"`
 	Function    *FunctionCall    `json:"function,omitempty"` // For INSERT INTO FUNCTION syntax
 	Columns     []*Identifier    `json:"columns,omitempty"`
+	AllColumns  bool             `json:"all_columns,omitempty"` // For (*) syntax meaning all columns
 	PartitionBy Expression       `json:"partition_by,omitempty"` // For PARTITION BY clause
 	Infile      string           `json:"infile,omitempty"`       // For FROM INFILE clause
 	Compression string           `json:"compression,omitempty"`  // For COMPRESSION clause
@@ -271,6 +274,7 @@ type CreateQuery struct {
 	CreateUser       bool                 `json:"create_user,omitempty"`
 	AlterUser        bool                 `json:"alter_user,omitempty"`
 	HasAuthenticationData bool            `json:"has_authentication_data,omitempty"`
+	AuthenticationValues []string         `json:"authentication_values,omitempty"` // Password/hash values from IDENTIFIED BY
 	CreateDictionary   bool                              `json:"create_dictionary,omitempty"`
 	DictionaryAttrs    []*DictionaryAttributeDeclaration `json:"dictionary_attrs,omitempty"`
 	DictionaryDef      *DictionaryDefinition             `json:"dictionary_def,omitempty"`
@@ -531,6 +535,7 @@ type AlterCommand struct {
 	ProjectionName string               `json:"projection_name,omitempty"` // For DROP/MATERIALIZE/CLEAR PROJECTION
 	StatisticsColumns []string          `json:"statistics_columns,omitempty"` // For ADD/DROP/CLEAR/MATERIALIZE STATISTICS
 	StatisticsTypes   []*FunctionCall   `json:"statistics_types,omitempty"`   // For ADD/MODIFY STATISTICS TYPE
+	Comment           string            `json:"comment,omitempty"`            // For COMMENT COLUMN
 }
 
 // Projection represents a projection definition.
@@ -698,11 +703,12 @@ const (
 
 // ExplainQuery represents an EXPLAIN statement.
 type ExplainQuery struct {
-	Position     token.Position `json:"-"`
-	ExplainType  ExplainType    `json:"explain_type"`
-	Statement    Statement      `json:"statement"`
-	HasSettings  bool           `json:"has_settings,omitempty"`
-	ExplicitType bool           `json:"explicit_type,omitempty"` // true if type was explicitly specified
+	Position       token.Position `json:"-"`
+	ExplainType    ExplainType    `json:"explain_type"`
+	Statement      Statement      `json:"statement"`
+	HasSettings    bool           `json:"has_settings,omitempty"`
+	ExplicitType   bool           `json:"explicit_type,omitempty"` // true if type was explicitly specified
+	OptionsString  string         `json:"options_string,omitempty"` // Formatted options like "actions = 1"
 }
 
 func (e *ExplainQuery) Pos() token.Position { return e.Position }
@@ -739,6 +745,7 @@ type OptimizeQuery struct {
 	Table     string         `json:"table"`
 	Partition Expression     `json:"partition,omitempty"`
 	Final     bool           `json:"final,omitempty"`
+	Cleanup   bool           `json:"cleanup,omitempty"`
 	Dedupe    bool           `json:"dedupe,omitempty"`
 	OnCluster string         `json:"on_cluster,omitempty"`
 }
@@ -772,6 +779,17 @@ func (s *SystemQuery) Pos() token.Position { return s.Position }
 func (s *SystemQuery) End() token.Position { return s.Position }
 func (s *SystemQuery) statementNode()      {}
 
+// TransactionControlQuery represents a transaction control statement (BEGIN, COMMIT, ROLLBACK, SET TRANSACTION SNAPSHOT).
+type TransactionControlQuery struct {
+	Position token.Position `json:"-"`
+	Action   string         `json:"action"` // "BEGIN", "COMMIT", "ROLLBACK", "SET_SNAPSHOT"
+	Snapshot int64          `json:"snapshot,omitempty"`
+}
+
+func (t *TransactionControlQuery) Pos() token.Position { return t.Position }
+func (t *TransactionControlQuery) End() token.Position { return t.Position }
+func (t *TransactionControlQuery) statementNode()      {}
+
 // RenamePair represents a single rename pair in RENAME TABLE.
 type RenamePair struct {
 	FromDatabase string `json:"from_database,omitempty"`
@@ -1064,6 +1082,7 @@ type Literal struct {
 	Position token.Position `json:"-"`
 	Type     LiteralType    `json:"type"`
 	Value    interface{}    `json:"value"`
+	Source   string         `json:"source,omitempty"`   // Original source text (for preserving 0.0 vs 0)
 	Negative bool           `json:"negative,omitempty"` // True if literal was explicitly negative (for -0)
 }
 
@@ -1126,6 +1145,7 @@ type Asterisk struct {
 	Table    string          `json:"table,omitempty"`   // for table.*
 	Except   []string        `json:"except,omitempty"`  // for * EXCEPT (col1, col2)
 	Replace  []*ReplaceExpr  `json:"replace,omitempty"` // for * REPLACE (expr AS col)
+	Apply    []string        `json:"apply,omitempty"`   // for * APPLY (func1) APPLY(func2)
 }
 
 func (a *Asterisk) Pos() token.Position { return a.Position }

internal/explain/explain.go

@@ -121,6 +121,8 @@ func Node(sb *strings.Builder, node interface{}, depth int) {
 		explainSetQuery(sb, indent)
 	case *ast.SystemQuery:
 		explainSystemQuery(sb, n, indent)
+	case *ast.TransactionControlQuery:
+		fmt.Fprintf(sb, "%sASTTransactionControl\n", indent)
 	case *ast.ExplainQuery:
 		explainExplainQuery(sb, n, indent, depth)
 	case *ast.ShowQuery:

internal/explain/expressions.go

@@ -529,8 +529,8 @@ func explainAliasedExpr(sb *strings.Builder, n *ast.AliasedExpr, depth int) {
 }
 
 func explainAsterisk(sb *strings.Builder, n *ast.Asterisk, indent string, depth int) {
-	// Check if there are any column transformers (EXCEPT, REPLACE)
-	hasTransformers := len(n.Except) > 0 || len(n.Replace) > 0
+	// Check if there are any column transformers (EXCEPT, REPLACE, APPLY)
+	hasTransformers := len(n.Except) > 0 || len(n.Replace) > 0 || len(n.Apply) > 0
 
 	if n.Table != "" {
 		if hasTransformers {
@@ -559,6 +559,8 @@ func explainColumnsTransformers(sb *strings.Builder, n *ast.Asterisk, indent str
 	if len(n.Replace) > 0 {
 		transformerCount++
 	}
+	// Each APPLY adds one transformer
+	transformerCount += len(n.Apply)
 
 	fmt.Fprintf(sb, "%sColumnsTransformerList (children %d)\n", indent, transformerCount)
 
@@ -583,6 +585,11 @@ func explainColumnsTransformers(sb *strings.Builder, n *ast.Asterisk, indent str
 			}
 		}
 	}
+
+	// Each APPLY function gets its own ColumnsApplyTransformer
+	for range n.Apply {
+		fmt.Fprintf(sb, "%s ColumnsApplyTransformer\n", indent)
+	}
 }
 
 func explainColumnsMatcher(sb *strings.Builder, n *ast.ColumnsMatcher, indent string, depth int) {

internal/explain/format.go

@@ -277,6 +277,7 @@ func formatBinaryExprForType(expr *ast.BinaryExpr) string {
 func NormalizeFunctionName(name string) string {
 	// ClickHouse normalizes certain function names in EXPLAIN AST
 	normalized := map[string]string{
+		"trim":       "trimBoth",
 		"ltrim":      "trimLeft",
 		"rtrim":      "trimRight",
 		"lcase":      "lower",
@@ -374,6 +375,10 @@ func formatExprAsString(expr ast.Expression) string {
 			}
 			return fmt.Sprintf("%d", e.Value)
 		case ast.LiteralFloat:
+			// Use Source field if available to preserve original representation (e.g., "0.0")
+			if e.Source != "" {
+				return e.Source
+			}
 			if e.Negative {
 				switch v := e.Value.(type) {
 				case float64:

internal/explain/select.go

@@ -9,8 +9,26 @@ import (
 
 func explainSelectIntersectExceptQuery(sb *strings.Builder, n *ast.SelectIntersectExceptQuery, indent string, depth int) {
 	fmt.Fprintf(sb, "%sSelectIntersectExceptQuery (children %d)\n", indent, len(n.Selects))
-	for _, sel := range n.Selects {
-		Node(sb, sel, depth+1)
+
+	// ClickHouse wraps first operand in SelectWithUnionQuery when EXCEPT is present
+	hasExcept := false
+	for _, op := range n.Operators {
+		if op == "EXCEPT" {
+			hasExcept = true
+			break
+		}
+	}
+
+	childIndent := strings.Repeat(" ", depth+1)
+	for i, sel := range n.Selects {
+		if hasExcept && i == 0 {
+			// Wrap first operand in SelectWithUnionQuery -> ExpressionList format
+			fmt.Fprintf(sb, "%sSelectWithUnionQuery (children 1)\n", childIndent)
+			fmt.Fprintf(sb, "%s ExpressionList (children 1)\n", childIndent)
+			Node(sb, sel, depth+3)
+		} else {
+			Node(sb, sel, depth+1)
+		}
 	}
 }
 
@@ -111,16 +129,31 @@ func explainSelectQuery(sb *strings.Builder, n *ast.SelectQuery, indent string,
 	if n.Offset != nil {
 		Node(sb, n.Offset, depth+1)
 	}
-	// LIMIT
-	if n.Limit != nil {
-		Node(sb, n.Limit, depth+1)
-	}
-	// LIMIT BY - only output when there's no ORDER BY and no second LIMIT (matches ClickHouse behavior)
-	if len(n.LimitBy) > 0 && len(n.OrderBy) == 0 && !n.LimitByHasLimit {
+	// LIMIT BY handling
+	if n.LimitByLimit != nil {
+		// Case: LIMIT n BY x LIMIT m -> output LimitByLimit, LimitBy, Limit
+		Node(sb, n.LimitByLimit, depth+1)
+		if len(n.LimitBy) > 0 {
+			fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(n.LimitBy))
+			for _, expr := range n.LimitBy {
+				Node(sb, expr, depth+2)
+			}
+		}
+		if n.Limit != nil {
+			Node(sb, n.Limit, depth+1)
+		}
+	} else if len(n.LimitBy) > 0 {
+		// Case: LIMIT n BY x (no second LIMIT) -> output Limit, then LimitBy
+		if n.Limit != nil {
+			Node(sb, n.Limit, depth+1)
+		}
 		fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(n.LimitBy))
 		for _, expr := range n.LimitBy {
 			Node(sb, expr, depth+2)
 		}
+	} else if n.Limit != nil {
+		// Case: plain LIMIT n (no BY)
+		Node(sb, n.Limit, depth+1)
 	}
 	// SETTINGS is output at SelectQuery level only when NOT after FORMAT
 	// When SettingsAfterFormat is true, it's output at SelectWithUnionQuery level instead
@@ -285,10 +318,13 @@ func countSelectQueryChildren(n *ast.SelectQuery) int {
 	if len(n.OrderBy) > 0 {
 		count++
 	}
+	if n.LimitByLimit != nil {
+		count++ // LIMIT n in "LIMIT n BY x LIMIT m"
+	}
 	if n.Limit != nil {
 		count++
 	}
-	if len(n.LimitBy) > 0 && len(n.OrderBy) == 0 && !n.LimitByHasLimit {
+	if len(n.LimitBy) > 0 {
 		count++
 	}
 	if n.Offset != nil {