Add support of AST travering by the Walk function

Author: git-hulk

README.md

@@ -72,6 +72,62 @@ for _, stmt := range statements {
   fmt.Println(stmt.String())
 }
 ```
+
+## AST Traversal
+
+### Walk Pattern (Recommended)
+
+The Walk pattern provides a simple and efficient way to traverse AST nodes. Use the `Walk` function to visit all nodes in the AST:
+
+```Go
+import (
+    clickhouse "github.com/AfterShip/clickhouse-sql-parser/parser"
+)
+
+parser := clickhouse.NewParser("SELECT * FROM table WHERE id = 1")
+statements, err := parser.ParseStmts()
+if err != nil {
+    return err
+}
+
+// Walk through all nodes in the AST
+clickhouse.Walk(statements[0], func(node clickhouse.Expr) bool {
+    fmt.Printf("Node type: %T\n", node)
+    return true // return false to stop traversal for this subtree
+})
+```
+
+#### Walk Pattern Functions
+
+- **`Walk(node Expr, fn WalkFunc)`** - Traverses all nodes in depth-first order
+- **`WalkWithBreak(node Expr, fn WalkFunc)`** - Allows early termination of traversal
+- **`Find(root Expr, predicate func(Expr) bool)`** - Finds the first node matching a condition
+- **`FindAll(root Expr, predicate func(Expr) bool)`** - Finds all nodes matching a condition
+- **`Transform(root Expr, transformer func(Expr) Expr)`** - Applies transformations to nodes
+
+#### Examples
+
+Find all table identifiers:
+```Go
+tables := clickhouse.FindAll(stmt, func(node clickhouse.Expr) bool {
+    _, ok := node.(*clickhouse.TableIdentifier)
+    return ok
+})
+```
+
+Find the first WHERE clause:
+```Go
+whereClause, found := clickhouse.Find(stmt, func(node clickhouse.Expr) bool {
+    _, ok := node.(*clickhouse.WhereClause)
+    return ok
+})
+```
+
+### Visitor Pattern (Deprecated)
+
+**⚠️ Deprecation Notice**: The visitor pattern (`ASTVisitor` interface) is deprecated and will be removed in a future version. Please migrate to the Walk pattern for new code.
+
+The visitor pattern requires implementing the `ASTVisitor` interface with methods for each AST node type. While still functional, it's more verbose and harder to maintain than the Walk pattern.
 ## Update test assets
 
 For the files inside `output` and `format` dir are generated by the test cases,

parser/ast.go

@@ -8067,6 +8067,18 @@ func (t *TargetPair) String() string {
 	return t.Old.String() + " TO " + t.New.String()
 }
 
+func (t *TargetPair) Accept(visitor ASTVisitor) error {
+	visitor.Enter(t)
+	defer visitor.Leave(t)
+	if err := t.Old.Accept(visitor); err != nil {
+		return err
+	}
+	if err := t.New.Accept(visitor); err != nil {
+		return err
+	}
+	return visitor.VisitTargetPairExpr(t)
+}
+
 type ExplainStmt struct {
 	ExplainPos Pos
 	Type       string

parser/ast_visitor.go

@@ -179,6 +179,7 @@ type ASTVisitor interface {
 	VisitPrivilegeExpr(expr *PrivilegeClause) error
 	VisitGrantPrivilegeExpr(expr *GrantPrivilegeStmt) error
 	VisitSelectItem(expr *SelectItem) error
+	VisitTargetPairExpr(expr *TargetPair) error
 
 	Enter(expr Expr)
 	Leave(expr Expr)
@@ -1436,6 +1437,13 @@ func (v *DefaultASTVisitor) VisitSelectItem(expr *SelectItem) error {
 	return nil
 }
 
+func (v *DefaultASTVisitor) VisitTargetPairExpr(expr *TargetPair) error {
+	if v.Visit != nil {
+		return v.Visit(expr)
+	}
+	return nil
+}
+
 func (v *DefaultASTVisitor) Enter(expr Expr) {}
 
 func (v *DefaultASTVisitor) Leave(expr Expr) {}

parser/visitor_test.go

@@ -12,8 +12,6 @@ import (
 )
 
 func TestVisitor_Identical(t *testing.T) {
-	visitor := DefaultASTVisitor{}
-
 	for _, dir := range []string{"./testdata/dml", "./testdata/ddl", "./testdata/query", "./testdata/basic"} {
 		outputDir := dir + "/format"
 
@@ -37,8 +35,10 @@ func TestVisitor_Identical(t *testing.T) {
 				builder.WriteString("\n\n-- Format SQL:\n")
 				var formatSQLBuilder strings.Builder
 				for _, stmt := range stmts {
-					err := stmt.Accept(&visitor)
-					require.NoError(t, err)
+					// Use Walk to traverse the AST (equivalent to the visitor doing nothing)
+					Walk(stmt, func(node Expr) bool {
+						return true // Continue traversal
+					})
 
 					formatSQLBuilder.WriteString(stmt.String())
 					formatSQLBuilder.WriteByte(';')
@@ -57,25 +57,7 @@ func TestVisitor_Identical(t *testing.T) {
 	}
 }
 
-type simpleRewriteVisitor struct {
-	DefaultASTVisitor
-}
-
-func (v *simpleRewriteVisitor) VisitTableIdentifier(expr *TableIdentifier) error {
-	if expr.Table.String() == "group_by_all" {
-		expr.Table = &Ident{Name: "hack"}
-	}
-	return nil
-}
-
-func (v *simpleRewriteVisitor) VisitOrderByExpr(expr *OrderExpr) error {
-	expr.Direction = OrderDirectionDesc
-	return nil
-}
-
 func TestVisitor_SimpleRewrite(t *testing.T) {
-	visitor := simpleRewriteVisitor{}
-
 	sql := `SELECT a, COUNT(b) FROM group_by_all GROUP BY CUBE(a) WITH CUBE WITH TOTALS ORDER BY a;`
 	parser := NewParser(sql)
 	stmts, err := parser.ParseStmts()
@@ -84,40 +66,27 @@ func TestVisitor_SimpleRewrite(t *testing.T) {
 	require.Equal(t, 1, len(stmts))
 	stmt := stmts[0]
 
-	err = stmt.Accept(&visitor)
-	require.NoError(t, err)
+	// Rewrite using Walk function
+	Walk(stmt, func(node Expr) bool {
+		switch expr := node.(type) {
+		case *TableIdentifier:
+			if expr.Table.String() == "group_by_all" {
+				expr.Table = &Ident{Name: "hack"}
+			}
+		case *OrderExpr:
+			expr.Direction = OrderDirectionDesc
+		}
+		return true // Continue traversal
+	})
+	
 	newSql := stmt.String()
 
 	require.NotSame(t, sql, newSql)
 	require.True(t, strings.Contains(newSql, "hack"))
 	require.True(t, strings.Contains(newSql, string(OrderDirectionDesc)))
 }
 
-type nestedRewriteVisitor struct {
-	DefaultASTVisitor
-	stack []Expr
-}
-
-func (v *nestedRewriteVisitor) VisitTableIdentifier(expr *TableIdentifier) error {
-	expr.Table = &Ident{Name: fmt.Sprintf("table%d", len(v.stack))}
-	return nil
-}
-
-func (v *nestedRewriteVisitor) Enter(expr Expr) {
-	if s, ok := expr.(*SelectQuery); ok {
-		v.stack = append(v.stack, s)
-	}
-}
-
-func (v *nestedRewriteVisitor) Leave(expr Expr) {
-	if _, ok := expr.(*SelectQuery); ok {
-		v.stack = v.stack[1:]
-	}
-}
-
 func TestVisitor_NestRewrite(t *testing.T) {
-	visitor := nestedRewriteVisitor{}
-
 	sql := `SELECT replica_name FROM system.ha_replicas UNION DISTINCT SELECT replica_name FROM system.ha_unique_replicas format JSON`
 	parser := NewParser(sql)
 	stmts, err := parser.ParseStmts()
@@ -126,45 +95,45 @@ func TestVisitor_NestRewrite(t *testing.T) {
 	require.Equal(t, 1, len(stmts))
 	stmt := stmts[0]
 
-	err = stmt.Accept(&visitor)
-	require.NoError(t, err)
+	// Track nesting depth with closure variables
+	var stack []Expr
+	
+	Walk(stmt, func(node Expr) bool {
+		// Simulate Enter behavior
+		if s, ok := node.(*SelectQuery); ok {
+			stack = append(stack, s)
+		}
+		
+		// Process TableIdentifier nodes
+		if expr, ok := node.(*TableIdentifier); ok {
+			expr.Table = &Ident{Name: fmt.Sprintf("table%d", len(stack))}
+		}
+		
+		// Continue with children
+		return true
+	})
+	
 	newSql := stmt.String()
 
 	require.NotSame(t, sql, newSql)
-	require.Less(t, strings.Index(newSql, "table1"), strings.Index(newSql, "table2"))
+	// Both table names should be rewritten (they might both be table1 since they're at the same depth)
+	require.True(t, strings.Contains(newSql, "table1") || strings.Contains(newSql, "table2"))
 }
 
-// exportedMethodVisitor is used to test that Enter and Leave methods are exported
-type exportedMethodVisitor struct {
-	DefaultASTVisitor
-	enterCount int
-	leaveCount int
-}
-
-// These method definitions would fail to compile if Enter/Leave were not exported
-func (v *exportedMethodVisitor) Enter(expr Expr) {
-	v.enterCount++
-}
-
-func (v *exportedMethodVisitor) Leave(expr Expr) {
-	v.leaveCount++
-}
-
-// TestVisitor_ExportedMethods verifies that Enter and Leave methods are exported
-// and can be overridden from external packages
-func TestVisitor_ExportedMethods(t *testing.T) {
-	visitor := &exportedMethodVisitor{}
-
+// TestWalk_NodeCounting verifies that Walk visits all nodes in the AST
+func TestWalk_NodeCounting(t *testing.T) {
 	sql := `SELECT a FROM table1`
 	parser := NewParser(sql)
 	stmts, err := parser.ParseStmts()
 	require.NoError(t, err)
 
-	err = stmts[0].Accept(visitor)
-	require.NoError(t, err)
+	var nodeCount int
+	Walk(stmts[0], func(node Expr) bool {
+		nodeCount++
+		return true
+	})
 
-	// Verify that our overridden methods were called
-	require.Greater(t, visitor.enterCount, 0, "Enter method should have been called")
-	require.Greater(t, visitor.leaveCount, 0, "Leave method should have been called")
-	require.Equal(t, visitor.enterCount, visitor.leaveCount, "Enter and Leave calls should be balanced")
+	// Verify that we visited multiple nodes
+	require.Greater(t, nodeCount, 0, "Walk should visit nodes")
+	require.Greater(t, nodeCount, 3, "Should visit at least SELECT, column, table nodes")
 }