Add support for ANY/ALL comparison operators with subqueries (#110)

Author: kyleconroy

ast/ast.go

@@ -72,7 +72,8 @@ type SelectQuery struct {
 	Having      Expression            `json:"having,omitempty"`
 	Qualify     Expression            `json:"qualify,omitempty"`
 	Window      []*WindowDefinition   `json:"window,omitempty"`
-	OrderBy     []*OrderByElement     `json:"order_by,omitempty"`
+	OrderBy      []*OrderByElement      `json:"order_by,omitempty"`
+	Interpolate  []*InterpolateElement  `json:"interpolate,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)
@@ -212,6 +213,17 @@ type OrderByElement struct {
 func (o *OrderByElement) Pos() token.Position { return o.Position }
 func (o *OrderByElement) End() token.Position { return o.Position }
 
+// InterpolateElement represents a single column interpolation in INTERPOLATE clause.
+// Example: INTERPOLATE (value AS value + 1)
+type InterpolateElement struct {
+	Position token.Position `json:"-"`
+	Column   string         `json:"column"`
+	Value    Expression     `json:"value,omitempty"` // nil if just column name
+}
+
+func (i *InterpolateElement) Pos() token.Position { return i.Position }
+func (i *InterpolateElement) End() token.Position { return i.Position }
+
 // SettingExpr represents a setting expression.
 type SettingExpr struct {
 	Position token.Position `json:"-"`
@@ -284,6 +296,7 @@ type CreateQuery struct {
 	FunctionName       string                            `json:"function_name,omitempty"`
 	FunctionBody     Expression           `json:"function_body,omitempty"`
 	UserName         string               `json:"user_name,omitempty"`
+	Format           string               `json:"format,omitempty"` // For FORMAT clause
 }
 
 func (c *CreateQuery) Pos() token.Position { return c.Position }
@@ -493,6 +506,7 @@ type DropQuery struct {
 	OnCluster       string             `json:"on_cluster,omitempty"`
 	DropDatabase    bool               `json:"drop_database,omitempty"`
 	Sync            bool               `json:"sync,omitempty"`
+	Format          string             `json:"format,omitempty"` // For FORMAT clause
 }
 
 func (d *DropQuery) Pos() token.Position { return d.Position }
@@ -512,6 +526,20 @@ func (u *UndropQuery) Pos() token.Position { return u.Position }
 func (u *UndropQuery) End() token.Position { return u.Position }
 func (u *UndropQuery) statementNode()      {}
 
+// UpdateQuery represents a standalone UPDATE statement.
+// In ClickHouse, UPDATE is syntactic sugar for ALTER TABLE ... UPDATE
+type UpdateQuery struct {
+	Position    token.Position `json:"-"`
+	Database    string         `json:"database,omitempty"`
+	Table       string         `json:"table"`
+	Assignments []*Assignment  `json:"assignments"`
+	Where       Expression     `json:"where,omitempty"`
+}
+
+func (u *UpdateQuery) Pos() token.Position { return u.Position }
+func (u *UpdateQuery) End() token.Position { return u.Position }
+func (u *UpdateQuery) statementNode()      {}
+
 // AlterQuery represents an ALTER statement.
 type AlterQuery struct {
 	Position  token.Position  `json:"-"`
@@ -520,6 +548,7 @@ type AlterQuery struct {
 	Commands  []*AlterCommand `json:"commands"`
 	OnCluster string          `json:"on_cluster,omitempty"`
 	Settings  []*SettingExpr  `json:"settings,omitempty"`
+	Format    string          `json:"format,omitempty"` // For FORMAT clause
 }
 
 func (a *AlterQuery) Pos() token.Position { return a.Position }
@@ -624,6 +653,7 @@ const (
 	AlterDropStatistics     AlterCommandType = "DROP_STATISTICS"
 	AlterClearStatistics    AlterCommandType = "CLEAR_STATISTICS"
 	AlterMaterializeStatistics AlterCommandType = "MATERIALIZE_STATISTICS"
+	AlterModifyComment         AlterCommandType = "MODIFY_COMMENT"
 )
 
 // TruncateQuery represents a TRUNCATE statement.

internal/explain/explain.go

@@ -12,6 +12,10 @@ import (
 // This affects how negated literals with aliases are formatted
 var inSubqueryContext bool
 
+// inCreateQueryContext is a package-level flag to track when we're inside a CreateQuery
+// This affects whether FORMAT is output at SelectWithUnionQuery level (it shouldn't be, as CreateQuery outputs it)
+var inCreateQueryContext bool
+
 // Explain returns the EXPLAIN AST output for a statement, matching ClickHouse's format.
 func Explain(stmt ast.Statement) string {
 	var sb strings.Builder
@@ -57,6 +61,8 @@ func Node(sb *strings.Builder, node interface{}, depth int) {
 	// Expressions
 	case *ast.OrderByElement:
 		explainOrderByElement(sb, n, indent, depth)
+	case *ast.InterpolateElement:
+		explainInterpolateElement(sb, n, indent, depth)
 	case *ast.Identifier:
 		explainIdentifier(sb, n, indent)
 	case *ast.Literal:
@@ -236,6 +242,8 @@ func Node(sb *strings.Builder, node interface{}, depth int) {
 		explainCheckQuery(sb, n, indent)
 	case *ast.CreateIndexQuery:
 		explainCreateIndexQuery(sb, n, indent, depth)
+	case *ast.UpdateQuery:
+		explainUpdateQuery(sb, n, indent, depth)
 
 	// Types
 	case *ast.DataType:
@@ -262,6 +270,8 @@ func Node(sb *strings.Builder, node interface{}, depth int) {
 		explainDictionaryLayout(sb, n, indent, depth)
 	case *ast.DictionaryRange:
 		explainDictionaryRange(sb, n, indent, depth)
+	case *ast.Assignment:
+		explainAssignment(sb, n, indent, depth)
 
 	default:
 		// For unhandled types, just print the type name

internal/explain/expressions.go

@@ -235,6 +235,31 @@ func containsNonLiteralExpressions(exprs []ast.Expression) bool {
 	return false
 }
 
+// containsNonLiteralInNested checks if an array or tuple literal contains
+// non-literal elements at any nesting level (identifiers, function calls, etc.)
+func containsNonLiteralInNested(lit *ast.Literal) bool {
+	if lit.Type != ast.LiteralArray && lit.Type != ast.LiteralTuple {
+		return false
+	}
+	exprs, ok := lit.Value.([]ast.Expression)
+	if !ok {
+		return false
+	}
+	for _, e := range exprs {
+		// Check if this element is a non-literal (identifier, function call, etc.)
+		if _, isLit := e.(*ast.Literal); !isLit {
+			return true
+		}
+		// Recursively check nested arrays/tuples
+		if innerLit, ok := e.(*ast.Literal); ok {
+			if containsNonLiteralInNested(innerLit) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // containsTuples checks if a slice of expressions contains any tuple literals
 func containsTuples(exprs []ast.Expression) bool {
 	for _, e := range exprs {
@@ -377,10 +402,23 @@ func explainUnaryExpr(sb *strings.Builder, n *ast.UnaryExpr, indent string, dept
 				// Convert positive integer to negative
 				switch val := lit.Value.(type) {
 				case int64:
-					fmt.Fprintf(sb, "%sLiteral Int64_%d\n", indent, -val)
+					negVal := -val
+					// ClickHouse normalizes -0 to UInt64_0
+					if negVal == 0 {
+						fmt.Fprintf(sb, "%sLiteral UInt64_0\n", indent)
+					} else if negVal > 0 {
+						fmt.Fprintf(sb, "%sLiteral UInt64_%d\n", indent, negVal)
+					} else {
+						fmt.Fprintf(sb, "%sLiteral Int64_%d\n", indent, negVal)
+					}
 					return
 				case uint64:
-					fmt.Fprintf(sb, "%sLiteral Int64_-%d\n", indent, val)
+					// ClickHouse normalizes -0 to UInt64_0
+					if val == 0 {
+						fmt.Fprintf(sb, "%sLiteral UInt64_0\n", indent)
+					} else {
+						fmt.Fprintf(sb, "%sLiteral Int64_-%d\n", indent, val)
+					}
 					return
 				}
 			case ast.LiteralFloat:
@@ -432,11 +470,23 @@ func explainAliasedExpr(sb *strings.Builder, n *ast.AliasedExpr, depth int) {
 						needsFunctionFormat = true
 						break
 					}
+					// Also check if nested arrays/tuples contain non-literal elements
+					if lit, ok := expr.(*ast.Literal); ok {
+						if containsNonLiteralInNested(lit) {
+							needsFunctionFormat = true
+							break
+						}
+					}
 				}
 				if needsFunctionFormat {
 					// Render as Function tuple with alias
 					fmt.Fprintf(sb, "%sFunction tuple (alias %s) (children %d)\n", indent, escapeAlias(n.Alias), 1)
-					fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(exprs))
+					// For empty ExpressionList, don't include children count
+					if len(exprs) > 0 {
+						fmt.Fprintf(sb, "%s ExpressionList (children %d)\n", indent, len(exprs))
+					} else {
+						fmt.Fprintf(sb, "%s ExpressionList\n", indent)
+					}
 					for _, expr := range exprs {
 						Node(sb, expr, depth+2)
 					}
@@ -463,6 +513,11 @@ func explainAliasedExpr(sb *strings.Builder, n *ast.AliasedExpr, depth int) {
 						needsFunctionFormat = true
 						break
 					}
+					// Check for function calls - use Function array
+					if _, ok := expr.(*ast.FunctionCall); ok {
+						needsFunctionFormat = true
+						break
+					}
 				}
 				if needsFunctionFormat {
 					// Render as Function array with alias
@@ -577,6 +632,9 @@ func explainAliasedExpr(sb *strings.Builder, n *ast.AliasedExpr, depth int) {
 	case *ast.CaseExpr:
 		// CASE expressions with alias
 		explainCaseExprWithAlias(sb, e, n.Alias, indent, depth)
+	case *ast.ExistsExpr:
+		// EXISTS expressions with alias
+		explainExistsExprWithAlias(sb, e, n.Alias, indent, depth)
 	default:
 		// For other types, recursively explain and add alias info
 		Node(sb, n.Expr, depth)

internal/explain/format.go

@@ -74,6 +74,36 @@ func escapeStringLiteral(s string) string {
 	return sb.String()
 }
 
+// escapeStringForTypeParam escapes special characters for use in type parameters
+// Uses extra escaping because type strings are embedded inside another string literal
+func escapeStringForTypeParam(s string) string {
+	var sb strings.Builder
+	for i := 0; i < len(s); i++ {
+		b := s[i]
+		switch b {
+		case '\\':
+			sb.WriteString("\\\\\\\\\\\\\\\\") // backslash becomes 8 backslashes
+		case '\'':
+			sb.WriteString("\\\\\\\\\\'") // single quote becomes 5 backslashes + quote
+		case '\n':
+			sb.WriteString("\\\\\\\\n") // newline becomes \\\\n
+		case '\t':
+			sb.WriteString("\\\\\\\\t") // tab becomes \\\\t
+		case '\r':
+			sb.WriteString("\\\\\\\\r") // carriage return becomes \\\\r
+		case '\x00':
+			sb.WriteString("\\\\\\\\0") // null becomes \\\\0
+		case '\b':
+			sb.WriteString("\\\\\\\\b") // backspace becomes \\\\b
+		case '\f':
+			sb.WriteString("\\\\\\\\f") // form feed becomes \\\\f
+		default:
+			sb.WriteByte(b)
+		}
+	}
+	return sb.String()
+}
+
 // FormatLiteral formats a literal value for EXPLAIN AST output
 func FormatLiteral(lit *ast.Literal) string {
 	switch lit.Type {
@@ -270,7 +300,9 @@ func formatBinaryExprForType(expr *ast.BinaryExpr) string {
 	// Format left side
 	if lit, ok := expr.Left.(*ast.Literal); ok {
 		if lit.Type == ast.LiteralString {
-			left = fmt.Sprintf("\\\\\\'%s\\\\\\'", lit.Value)
+			// Use extra escaping for type parameters since they're embedded in another string literal
+			escaped := escapeStringForTypeParam(fmt.Sprintf("%v", lit.Value))
+			left = fmt.Sprintf("\\\\\\'%s\\\\\\'", escaped)
 		} else {
 			left = fmt.Sprintf("%v", lit.Value)
 		}
@@ -285,13 +317,24 @@ func formatBinaryExprForType(expr *ast.BinaryExpr) string {
 		right = fmt.Sprintf("%v", lit.Value)
 	} else if ident, ok := expr.Right.(*ast.Identifier); ok {
 		right = ident.Name()
+	} else if unary, ok := expr.Right.(*ast.UnaryExpr); ok {
+		// Handle unary expressions like -100
+		right = formatUnaryExprForType(unary)
 	} else {
 		right = fmt.Sprintf("%v", expr.Right)
 	}
 
 	return left + " " + expr.Op + " " + right
 }
 
+// formatUnaryExprForType formats a unary expression for use in type parameters (e.g., -100)
+func formatUnaryExprForType(expr *ast.UnaryExpr) string {
+	if lit, ok := expr.Operand.(*ast.Literal); ok {
+		return expr.Op + fmt.Sprintf("%v", lit.Value)
+	}
+	return expr.Op + fmt.Sprintf("%v", expr.Operand)
+}
+
 // NormalizeFunctionName normalizes function names to match ClickHouse's EXPLAIN AST output
 func NormalizeFunctionName(name string) string {
 	// ClickHouse normalizes certain function names in EXPLAIN AST
@@ -314,9 +357,9 @@ func NormalizeFunctionName(name string) string {
 		"least":      "least",
 		"concat_ws":  "concat",
 		"position":   "position",
-		// SQL standard ANY/ALL subquery operators
-		"anymatch":   "in",
-		"allmatch":   "notIn",
+		// SQL standard ANY/ALL subquery operators - simple cases
+		"anyequals":    "in",
+		"allnotequals": "notIn",
 	}
 	if n, ok := normalized[strings.ToLower(name)]; ok {
 		return n
@@ -351,6 +394,8 @@ func OperatorToFunction(op string) string {
 		return "lessOrEquals"
 	case ">=":
 		return "greaterOrEquals"
+	case "<=>":
+		return "isNotDistinctFrom"
 	case "AND":
 		return "and"
 	case "OR":