Fix INTERSECT/EXCEPT operator precedence and WITH clause inheritance (#111)

Author: kyleconroy

ast/ast.go

@@ -272,11 +272,12 @@ type CreateQuery struct {
 	Columns          []*ColumnDeclaration `json:"columns,omitempty"`
 	Indexes          []*IndexDefinition   `json:"indexes,omitempty"`
 	Projections      []*Projection        `json:"projections,omitempty"`
-	Constraints      []*Constraint        `json:"constraints,omitempty"`
-	Engine           *EngineClause        `json:"engine,omitempty"`
-	OrderBy          []Expression         `json:"order_by,omitempty"`
-	PartitionBy      Expression           `json:"partition_by,omitempty"`
-	PrimaryKey       []Expression         `json:"primary_key,omitempty"`
+	Constraints         []*Constraint        `json:"constraints,omitempty"`
+	ColumnsPrimaryKey   []Expression         `json:"columns_primary_key,omitempty"` // PRIMARY KEY in column list
+	Engine              *EngineClause        `json:"engine,omitempty"`
+	OrderBy             []Expression         `json:"order_by,omitempty"`
+	PartitionBy         Expression           `json:"partition_by,omitempty"`
+	PrimaryKey          []Expression         `json:"primary_key,omitempty"`
 	SampleBy         Expression           `json:"sample_by,omitempty"`
 	TTL              *TTLClause           `json:"ttl,omitempty"`
 	Settings         []*SettingExpr       `json:"settings,omitempty"`
@@ -502,6 +503,7 @@ type DropQuery struct {
 	Policy          string             `json:"policy,omitempty"`   // For DROP POLICY
 	RowPolicy       string             `json:"row_policy,omitempty"` // For DROP ROW POLICY
 	SettingsProfile string             `json:"settings_profile,omitempty"` // For DROP SETTINGS PROFILE
+	Index           string             `json:"index,omitempty"`            // For DROP INDEX
 	Temporary       bool               `json:"temporary,omitempty"`
 	OnCluster       string             `json:"on_cluster,omitempty"`
 	DropDatabase    bool               `json:"drop_database,omitempty"`
@@ -642,6 +644,7 @@ const (
 	AlterMovePartition      AlterCommandType = "MOVE_PARTITION"
 	AlterFreezePartition    AlterCommandType = "FREEZE_PARTITION"
 	AlterFreeze             AlterCommandType = "FREEZE"
+	AlterApplyPatches       AlterCommandType = "APPLY_PATCHES"
 	AlterDeleteWhere        AlterCommandType = "DELETE_WHERE"
 	AlterUpdate             AlterCommandType = "UPDATE"
 	AlterAddProjection      AlterCommandType = "ADD_PROJECTION"
@@ -669,6 +672,18 @@ func (t *TruncateQuery) Pos() token.Position { return t.Position }
 func (t *TruncateQuery) End() token.Position { return t.Position }
 func (t *TruncateQuery) statementNode()      {}
 
+// DeleteQuery represents a lightweight DELETE statement.
+type DeleteQuery struct {
+	Position token.Position `json:"-"`
+	Database string         `json:"database,omitempty"`
+	Table    string         `json:"table"`
+	Where    Expression     `json:"where,omitempty"`
+}
+
+func (d *DeleteQuery) Pos() token.Position { return d.Position }
+func (d *DeleteQuery) End() token.Position { return d.Position }
+func (d *DeleteQuery) statementNode()      {}
+
 // UseQuery represents a USE statement.
 type UseQuery struct {
 	Position token.Position `json:"-"`
@@ -1106,10 +1121,12 @@ func (d *DropWorkloadQuery) statementNode()      {}
 
 // CreateIndexQuery represents a CREATE INDEX statement.
 type CreateIndexQuery struct {
-	Position  token.Position `json:"-"`
-	IndexName string         `json:"index_name"`
-	Table     string         `json:"table"`
-	Columns   []Expression   `json:"columns,omitempty"`
+	Position    token.Position `json:"-"`
+	IndexName   string         `json:"index_name"`
+	Table       string         `json:"table"`
+	Columns     []Expression   `json:"columns,omitempty"`
+	Type        string         `json:"type,omitempty"`        // Index type (minmax, bloom_filter, etc.)
+	Granularity int            `json:"granularity,omitempty"` // GRANULARITY value
 }
 
 func (c *CreateIndexQuery) Pos() token.Position { return c.Position }

internal/explain/explain.go

@@ -238,6 +238,8 @@ func Node(sb *strings.Builder, node interface{}, depth int) {
 		explainOptimizeQuery(sb, n, indent, depth)
 	case *ast.TruncateQuery:
 		explainTruncateQuery(sb, n, indent)
+	case *ast.DeleteQuery:
+		explainDeleteQuery(sb, n, indent, depth)
 	case *ast.CheckQuery:
 		explainCheckQuery(sb, n, indent)
 	case *ast.CreateIndexQuery:

internal/explain/expressions.go

@@ -864,6 +864,35 @@ func explainWithElement(sb *strings.Builder, n *ast.WithElement, indent string,
 				return
 			}
 		}
+		// Arrays containing non-literal expressions should be rendered as Function array
+		if e.Type == ast.LiteralArray {
+			if exprs, ok := e.Value.([]ast.Expression); ok {
+				needsFunctionFormat := false
+				for _, elem := range exprs {
+					if !isSimpleLiteralOrNegation(elem) {
+						needsFunctionFormat = true
+						break
+					}
+				}
+				if needsFunctionFormat {
+					// Render as Function array with alias
+					if n.Name != "" {
+						fmt.Fprintf(sb, "%sFunction array (alias %s) (children %d)\n", indent, n.Name, 1)
+					} else {
+						fmt.Fprintf(sb, "%sFunction array (children %d)\n", indent, 1)
+					}
+					if len(exprs) > 0 {
+						fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(exprs))
+					} else {
+						fmt.Fprintf(sb, "%s ExpressionList\n", indent)
+					}
+					for _, elem := range exprs {
+						Node(sb, elem, depth+2)
+					}
+					return
+				}
+			}
+		}
 		if n.Name != "" {
 			fmt.Fprintf(sb, "%sLiteral %s (alias %s)\n", indent, FormatLiteral(e), n.Name)
 		} else {

internal/explain/format.go

@@ -357,6 +357,8 @@ func NormalizeFunctionName(name string) string {
 		"least":      "least",
 		"concat_ws":  "concat",
 		"position":   "position",
+		"date_diff":  "dateDiff",
+		"datediff":   "dateDiff",
 		// SQL standard ANY/ALL subquery operators - simple cases
 		"anyequals":    "in",
 		"allnotequals": "notIn",

internal/explain/functions.go

@@ -372,9 +372,9 @@ func explainDateAddSubWithInterval(sb *strings.Builder, opFunc string, arg1, arg
 }
 
 // handleDateDiff handles DATE_DIFF/DATEDIFF
-// DATE_DIFF(unit, date1, date2) -> dateDiff('unit', date1, date2)
+// DATE_DIFF(unit, date1, date2[, timezone]) -> dateDiff('unit', date1, date2[, timezone])
 func handleDateDiff(sb *strings.Builder, n *ast.FunctionCall, alias string, indent string, depth int) bool {
-	if len(n.Arguments) != 3 {
+	if len(n.Arguments) < 3 || len(n.Arguments) > 4 {
 		return false
 	}
 
@@ -392,12 +392,17 @@ func handleDateDiff(sb *strings.Builder, n *ast.FunctionCall, alias string, inde
 		return false
 	}
 
+	argCount := 3
+	if len(n.Arguments) == 4 {
+		argCount = 4
+	}
+
 	if alias != "" {
 		fmt.Fprintf(sb, "%sFunction dateDiff (alias %s) (children %d)\n", indent, alias, 1)
 	} else {
 		fmt.Fprintf(sb, "%sFunction dateDiff (children %d)\n", indent, 1)
 	}
-	fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, 3)
+	fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, argCount)
 
 	// First arg: unit as lowercase string literal
 	fmt.Fprintf(sb, "%s  Literal \\'%s\\'\n", indent, strings.ToLower(unitName))
@@ -406,6 +411,11 @@ func handleDateDiff(sb *strings.Builder, n *ast.FunctionCall, alias string, inde
 	Node(sb, date1Arg, depth+2)
 	Node(sb, date2Arg, depth+2)
 
+	// Fourth arg: optional timezone
+	if len(n.Arguments) == 4 {
+		Node(sb, n.Arguments[3], depth+2)
+	}
+
 	return true
 }
 
@@ -1334,13 +1344,10 @@ func explainExtractExprWithAlias(sb *strings.Builder, n *ast.ExtractExpr, alias
 	// EXTRACT is represented as Function toYear, toMonth, etc.
 	// ClickHouse uses specific function names for date/time extraction
 	fnName := extractFieldToFunction(n.Field)
-	// Use alias from parameter, or fall back to expression's alias
-	effectiveAlias := alias
-	if effectiveAlias == "" {
-		effectiveAlias = n.Alias
-	}
-	if effectiveAlias != "" {
-		fmt.Fprintf(sb, "%sFunction %s (alias %s) (children %d)\n", indent, fnName, effectiveAlias, 1)
+	// Only use the external alias parameter (from explicit AS on EXTRACT itself)
+	// NOT the alias from the From expression - that stays on the inner expression
+	if alias != "" {
+		fmt.Fprintf(sb, "%sFunction %s (alias %s) (children %d)\n", indent, fnName, alias, 1)
 	} else {
 		fmt.Fprintf(sb, "%sFunction %s (children %d)\n", indent, fnName, 1)
 	}

internal/explain/select.go

@@ -19,19 +19,165 @@ func explainSelectIntersectExceptQuery(sb *strings.Builder, n *ast.SelectInterse
 		}
 	}
 
+	// Check if first operand has a WITH clause to be inherited by subsequent operands
+	var inheritedWith []ast.Expression
+	if len(n.Selects) > 0 {
+		inheritedWith = extractWithClause(n.Selects[0])
+	}
+
 	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)
+			// But if it's already a SelectWithUnionQuery, don't double-wrap
+			if _, isUnion := sel.(*ast.SelectWithUnionQuery); isUnion {
+				Node(sb, sel, depth+1)
+			} else {
+				fmt.Fprintf(sb, "%sSelectWithUnionQuery (children 1)\n", childIndent)
+				fmt.Fprintf(sb, "%s ExpressionList (children 1)\n", childIndent)
+				Node(sb, sel, depth+3)
+			}
+		} else if i > 0 && len(inheritedWith) > 0 {
+			// Subsequent operands inherit the WITH clause from the first operand
+			explainSelectQueryWithInheritedWith(sb, sel, inheritedWith, depth+1)
 		} else {
 			Node(sb, sel, depth+1)
 		}
 	}
 }
 
+// extractWithClause extracts the WITH clause from a statement (if it's a SelectQuery)
+func extractWithClause(stmt ast.Statement) []ast.Expression {
+	switch s := stmt.(type) {
+	case *ast.SelectQuery:
+		return s.With
+	case *ast.SelectWithUnionQuery:
+		// Check the first select in the union
+		if len(s.Selects) > 0 {
+			return extractWithClause(s.Selects[0])
+		}
+	}
+	return nil
+}
+
+// explainSelectQueryWithInheritedWith outputs a SELECT with an inherited WITH clause
+// The inherited WITH clause is output AFTER the columns (not before, like a regular WITH)
+func explainSelectQueryWithInheritedWith(sb *strings.Builder, stmt ast.Statement, inheritedWith []ast.Expression, depth int) {
+	sq, ok := stmt.(*ast.SelectQuery)
+	if !ok {
+		// Not a SelectQuery, output normally
+		Node(sb, stmt, depth)
+		return
+	}
+
+	// If the SelectQuery already has a WITH clause, output normally
+	if len(sq.With) > 0 {
+		Node(sb, stmt, depth)
+		return
+	}
+
+	// Output SelectQuery with inherited WITH clause after columns
+	indent := strings.Repeat(" ", depth)
+	children := countSelectQueryChildren(sq) + 1 // +1 for inherited WITH clause
+	fmt.Fprintf(sb, "%sSelectQuery (children %d)\n", indent, children)
+
+	// Columns (ExpressionList) - output BEFORE inherited WITH
+	fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(sq.Columns))
+	for _, col := range sq.Columns {
+		Node(sb, col, depth+2)
+	}
+
+	// Inherited WITH clause (ExpressionList) - output AFTER columns
+	fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(inheritedWith))
+	for _, w := range inheritedWith {
+		Node(sb, w, depth+2)
+	}
+
+	// FROM (including ARRAY JOIN as part of TablesInSelectQuery)
+	if sq.From != nil || sq.ArrayJoin != nil {
+		TablesWithArrayJoin(sb, sq.From, sq.ArrayJoin, depth+1)
+	}
+	// PREWHERE
+	if sq.PreWhere != nil {
+		Node(sb, sq.PreWhere, depth+1)
+	}
+	// WHERE
+	if sq.Where != nil {
+		Node(sb, sq.Where, depth+1)
+	}
+	// GROUP BY (skip for GROUP BY ALL which doesn't output an expression list)
+	if len(sq.GroupBy) > 0 && !sq.GroupByAll {
+		fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(sq.GroupBy))
+		for _, g := range sq.GroupBy {
+			Node(sb, g, depth+2)
+		}
+	}
+	// HAVING
+	if sq.Having != nil {
+		Node(sb, sq.Having, depth+1)
+	}
+	// QUALIFY
+	if sq.Qualify != nil {
+		Node(sb, sq.Qualify, depth+1)
+	}
+	// WINDOW clause
+	if len(sq.Window) > 0 {
+		fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(sq.Window))
+		for range sq.Window {
+			fmt.Fprintf(sb, "%s  WindowListElement\n", indent)
+		}
+	}
+	// ORDER BY
+	if len(sq.OrderBy) > 0 {
+		fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(sq.OrderBy))
+		for _, o := range sq.OrderBy {
+			Node(sb, o, depth+2)
+		}
+	}
+	// INTERPOLATE
+	if len(sq.Interpolate) > 0 {
+		fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(sq.Interpolate))
+		for _, i := range sq.Interpolate {
+			Node(sb, i, depth+2)
+		}
+	}
+	// OFFSET
+	if sq.Offset != nil {
+		Node(sb, sq.Offset, depth+1)
+	}
+	// LIMIT BY handling
+	if sq.LimitByLimit != nil {
+		Node(sb, sq.LimitByLimit, depth+1)
+		if len(sq.LimitBy) > 0 {
+			fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(sq.LimitBy))
+			for _, expr := range sq.LimitBy {
+				Node(sb, expr, depth+2)
+			}
+		}
+		if sq.Limit != nil {
+			Node(sb, sq.Limit, depth+1)
+		}
+	} else if len(sq.LimitBy) > 0 {
+		if sq.Limit != nil {
+			Node(sb, sq.Limit, depth+1)
+		}
+		fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(sq.LimitBy))
+		for _, expr := range sq.LimitBy {
+			Node(sb, expr, depth+2)
+		}
+	} else if sq.Limit != nil {
+		Node(sb, sq.Limit, depth+1)
+	}
+	// SETTINGS
+	if len(sq.Settings) > 0 && !sq.SettingsAfterFormat {
+		fmt.Fprintf(sb, "%s Set\n", indent)
+	}
+	// TOP clause
+	if sq.Top != nil {
+		Node(sb, sq.Top, depth+1)
+	}
+}
+
 func explainSelectWithUnionQuery(sb *strings.Builder, n *ast.SelectWithUnionQuery, indent string, depth int) {
 	if n == nil {
 		return