Implement Set and SetAll (and Get)

processor/jsonpath/jsonpath.go

@@ -23,6 +23,15 @@ const (
 	IndexStep
 )
 
+// TraversalMode determines how the traversal behaves
+type TraversalMode int
+
+const (
+	CollectMode     TraversalMode = iota // Just collect matched nodes
+	ModifyFirstMode                      // Modify first matching node
+	ModifyAllMode                        // Modify all matching nodes
+)
+
 func ParseJSONPath(path string) ([]JSONStep, error) {
 	if len(path) == 0 || path[0] != '$' {
 		return nil, fmt.Errorf("path must start with $")
@@ -96,120 +105,246 @@ func readIndex(path string, start int) (string, int) {
 	return path[start:i], i
 }
 
+// Get retrieves values from data at the specified JSONPath
 func Get(data interface{}, path string) []interface{} {
 	steps, err := ParseJSONPath(path)
 	if err != nil {
 		log.Println("Error parsing JSONPath:", err)
 		return nil
 	}
-	return EvaluateJSONPath(data, steps)
-}
-func Set(data interface{}, path string, value interface{}) {
-	steps, err := ParseJSONPath(path)
-	if err != nil {
-		log.Println("Error parsing JSONPath:", err)
-		return
-	}
-	node := EvaluateJSONPath(data, steps)
-	if len(node) == 0 {
-		log.Println("No node found for path:", path)
-		return
-	}
-	if len(node) > 1 {
-		log.Println("Multiple nodes found for path:", path)
-		return
-	}
-	node[0] = value
-}
-
-func EvaluateJSONPath(data interface{}, steps []JSONStep) []interface{} {
-	current := []interface{}{data}
-
-	for _, step := range steps {
-		var next []interface{}
-		for _, node := range current {
-			next = append(next, evalStep(node, step)...)
-		}
-		current = next
-	}
-
-	return current
-}
-
-func evalStep(node interface{}, step JSONStep) []interface{} {
-	switch step.Type {
-	case ChildStep:
-		return evalChild(node, step.Key)
-	case RecursiveDescentStep:
-		return evalRecursiveDescent(node, step.Key)
-	case WildcardStep:
-		return evalWildcard(node)
-	case IndexStep:
-		return evalIndex(node, step.Index)
-	case RootStep:
-		return []interface{}{node}
-	default:
-		log.Println("Unknown step type:", step.Type)
-		return nil
-	}
-}
-
-func evalChild(node interface{}, key string) []interface{} {
-	if m, ok := node.(map[string]interface{}); ok {
-		if val, exists := m[key]; exists {
-			return []interface{}{val}
-		}
-	}
-	return nil
-}
-
-func evalRecursiveDescent(node interface{}, targetKey string) []interface{} {
 	results := []interface{}{}
-	queue := []interface{}{node}
-
-	for len(queue) > 0 {
-		current := queue[0]
-		queue = queue[1:]
-
-		if targetKey == "*" {
-			results = append(results, current)
-		} else if m, ok := current.(map[string]interface{}); ok {
-			if val, exists := m[targetKey]; exists {
-				results = append(results, val)
-			}
-		}
-
-		if m, ok := current.(map[string]interface{}); ok {
-			for _, v := range m {
-				queue = append(queue, v)
-			}
-		} else if arr, ok := current.([]interface{}); ok {
-			queue = append(queue, arr...)
-		}
-	}
-
+	traverse(data, steps, CollectMode, nil, &results)
 	return results
 }
 
-func evalWildcard(node interface{}) []interface{} {
-	if m, ok := node.(map[string]interface{}); ok {
-		results := make([]interface{}, 0, len(m))
-		for _, v := range m {
-			results = append(results, v)
-		}
-		return results
+// Set updates the value at the specified JSONPath in the original data structure.
+func Set(data interface{}, path string, value interface{}) bool {
+	steps, err := ParseJSONPath(path)
+	if err != nil {
+		log.Println("Error parsing JSONPath:", err)
+		return false
 	}
-	return nil
+
+	if len(steps) <= 1 {
+		log.Println("Cannot set root node")
+		return false
+	}
+
+	success := false
+	traverse(data, steps, ModifyFirstMode, value, &success)
+	return success
 }
 
-func evalIndex(node interface{}, index int) []interface{} {
-	if arr, ok := node.([]interface{}); ok {
-		if index == -1 { // Wildcard [*]
-			return arr
+// SetAll updates all matching values at the specified JSONPath.
+func SetAll(data interface{}, path string, value interface{}) bool {
+	steps, err := ParseJSONPath(path)
+	if err != nil {
+		log.Println("Error parsing JSONPath:", err)
+		return false
+	}
+
+	if len(steps) <= 1 {
+		log.Println("Cannot set root node")
+		return false
+	}
+
+	success := false
+	traverse(data, steps, ModifyAllMode, value, &success)
+	return success
+}
+
+// traverse is the main entry point for JSONPath traversal
+func traverse(data interface{}, steps []JSONStep, mode TraversalMode, value interface{}, result interface{}) {
+	if len(steps) == 0 || data == nil {
+		return
+	}
+
+	// Skip root step
+	actualSteps := steps
+	if steps[0].Type == RootStep {
+		if len(steps) == 1 {
+			if mode == CollectMode {
+				results := result.(*[]interface{})
+				*results = append(*results, data)
+			}
+			return
 		}
-		if index >= 0 && index < len(arr) {
-			return []interface{}{arr[index]}
+		actualSteps = steps[1:]
+	}
+
+	// Start recursion
+	traverseSteps(data, actualSteps, mode, value, result)
+}
+
+// traverseSteps handles all path traversal with a unified algorithm
+func traverseSteps(node interface{}, steps []JSONStep, mode TraversalMode, value interface{}, result interface{}) {
+	if len(steps) == 0 {
+		// We've reached a terminal node
+		if mode == CollectMode {
+			results := result.(*[]interface{})
+			*results = append(*results, node)
+		}
+		return
+	}
+
+	// Extract current step info
+	step := steps[0]
+	isLastStep := len(steps) == 1
+	nextSteps := steps[1:]
+
+	// For modify modes, get the success pointer
+	var successPtr *bool
+	if mode != CollectMode {
+		successPtr = result.(*bool)
+	}
+
+	// Helper function to check if we should stop recursion
+	shouldStop := func() bool {
+		return mode == ModifyFirstMode && *successPtr
+	}
+
+	// Handle each step type
+	switch step.Type {
+	case ChildStep:
+		// Only process maps for child steps
+		m, ok := node.(map[string]interface{})
+		if !ok {
+			return
+		}
+
+		child, exists := m[step.Key]
+
+		// Handle modification on last step
+		if isLastStep && (mode == ModifyFirstMode || mode == ModifyAllMode) {
+			m[step.Key] = value
+			*successPtr = true
+			return
+		}
+
+		// Continue traversal with existing child
+		if exists {
+			traverseSteps(child, nextSteps, mode, value, result)
+			return
+		}
+
+		// Create missing intermediate nodes
+		if !isLastStep && (mode == ModifyFirstMode || mode == ModifyAllMode) {
+			var newNode interface{}
+			if len(nextSteps) > 0 && nextSteps[0].Type == IndexStep {
+				newNode = []interface{}{}
+			} else {
+				newNode = map[string]interface{}{}
+			}
+			m[step.Key] = newNode
+			traverseSteps(newNode, nextSteps, mode, value, result)
+		}
+
+	case IndexStep:
+		// Only process arrays for index steps
+		arr, ok := node.([]interface{})
+		if !ok {
+			return
+		}
+
+		// Handle wildcard index
+		if step.Index == -1 {
+			for i, item := range arr {
+				if isLastStep && (mode == ModifyFirstMode || mode == ModifyAllMode) {
+					arr[i] = value
+					*successPtr = true
+					if shouldStop() {
+						return
+					}
+				} else {
+					traverseSteps(item, nextSteps, mode, value, result)
+					if shouldStop() {
+						return
+					}
+				}
+			}
+			return
+		}
+
+		// Handle specific index
+		if step.Index >= 0 && step.Index < len(arr) {
+			if isLastStep && (mode == ModifyFirstMode || mode == ModifyAllMode) {
+				arr[step.Index] = value
+				*successPtr = true
+			} else {
+				traverseSteps(arr[step.Index], nextSteps, mode, value, result)
+			}
+		}
+
+	case RecursiveDescentStep:
+		// For recursive descent, first collect/modify match at this level if available
+		if m, ok := node.(map[string]interface{}); ok && step.Key != "*" {
+			if val, exists := m[step.Key]; exists {
+				if isLastStep {
+					if mode == CollectMode {
+						results := result.(*[]interface{})
+						*results = append(*results, val)
+					} else { // Modify modes
+						m[step.Key] = value
+						*successPtr = true
+						if shouldStop() {
+							return
+						}
+					}
+				} else if !isLastStep && mode != CollectMode {
+					// Continue with next steps for non-terminal direct matches
+					traverseSteps(val, nextSteps, mode, value, result)
+					if shouldStop() {
+						return
+					}
+				}
+			}
+		}
+
+		// For wildcard, collect this node
+		if step.Key == "*" && mode == CollectMode {
+			results := result.(*[]interface{})
+			*results = append(*results, node)
+		}
+
+		// Then continue recursion to all children
+		switch n := node.(type) {
+		case map[string]interface{}:
+			for _, v := range n {
+				traverseSteps(v, steps, mode, value, result) // Use same steps
+				if shouldStop() {
+					return
+				}
+			}
+		case []interface{}:
+			for _, v := range n {
+				traverseSteps(v, steps, mode, value, result) // Use same steps
+				if shouldStop() {
+					return
+				}
+			}
+		}
+
+	case WildcardStep:
+		// Only process maps for wildcard steps
+		m, ok := node.(map[string]interface{})
+		if !ok {
+			return
+		}
+
+		// Process all keys
+		for k, v := range m {
+			if isLastStep && (mode == ModifyFirstMode || mode == ModifyAllMode) {
+				m[k] = value
+				*successPtr = true
+				if shouldStop() {
+					return
+				}
+			} else {
+				traverseSteps(v, nextSteps, mode, value, result)
+				if shouldStop() {
+					return
+				}
+			}
 		}
 	}
-	return nil
 }

processor/jsonpath/jsonpath_get_set_test.go

@@ -0,0 +1,432 @@
+package jsonpath
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestGet(t *testing.T) {
+	tests := []struct {
+		name     string
+		data     map[string]interface{}
+		path     string
+		expected []interface{}
+	}{
+		{
+			name: "simple property",
+			data: map[string]interface{}{
+				"name": "John",
+				"age":  30,
+			},
+			path:     "$.name",
+			expected: []interface{}{"John"},
+		},
+		{
+			name: "nested property",
+			data: map[string]interface{}{
+				"user": map[string]interface{}{
+					"name": "John",
+					"age":  30,
+				},
+			},
+			path:     "$.user.name",
+			expected: []interface{}{"John"},
+		},
+		{
+			name: "array access",
+			data: map[string]interface{}{
+				"users": []interface{}{
+					map[string]interface{}{"name": "John", "age": 30},
+					map[string]interface{}{"name": "Jane", "age": 25},
+				},
+			},
+			path:     "$.users[1].name",
+			expected: []interface{}{"Jane"},
+		},
+		{
+			name: "wildcard",
+			data: map[string]interface{}{
+				"users": []interface{}{
+					map[string]interface{}{"name": "John", "age": 30},
+					map[string]interface{}{"name": "Jane", "age": 25},
+				},
+			},
+			path:     "$.users[*].name",
+			expected: []interface{}{"John", "Jane"},
+		},
+		{
+			name: "recursive descent",
+			data: map[string]interface{}{
+				"user": map[string]interface{}{
+					"name": "John",
+					"profile": map[string]interface{}{
+						"email": "john@example.com",
+					},
+				},
+				"admin": map[string]interface{}{
+					"email": "admin@example.com",
+				},
+			},
+			path:     "$..email",
+			expected: []interface{}{"john@example.com", "admin@example.com"},
+		},
+		{
+			name: "nonexistent path",
+			data: map[string]interface{}{
+				"user": map[string]interface{}{
+					"name": "John",
+				},
+			},
+			path:     "$.user.email",
+			expected: nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := Get(tt.data, tt.path)
+
+			// For nonexistent path, we expect either nil or empty slice
+			if tt.name == "nonexistent path" {
+				if len(result) > 0 {
+					t.Errorf("Get() returned %v, expected empty result", result)
+				}
+				return
+			}
+
+			// Check if lengths match
+			if len(result) != len(tt.expected) {
+				t.Errorf("Get() returned %d items, expected %d", len(result), len(tt.expected))
+				return
+			}
+
+			// For wildcard results, we need to check containment rather than exact order
+			if tt.name == "wildcard" || tt.name == "recursive descent" {
+				expectedMap := make(map[interface{}]bool)
+				for _, v := range tt.expected {
+					expectedMap[v] = true
+				}
+
+				for _, v := range result {
+					if !expectedMap[v] {
+						t.Errorf("Get() result contains unexpected value: %v", v)
+					}
+				}
+			} else {
+				// Otherwise check exact equality
+				if !reflect.DeepEqual(result, tt.expected) {
+					t.Errorf("Get() = %v, expected %v", result, tt.expected)
+				}
+			}
+		})
+	}
+}
+
+func TestSet(t *testing.T) {
+	t.Run("simple property", func(t *testing.T) {
+		data := map[string]interface{}{
+			"name": "John",
+			"age":  30,
+		}
+		success := Set(data, "$.name", "Jane")
+
+		if !success {
+			t.Errorf("Set() returned false, expected true")
+		}
+		if data["name"] != "Jane" {
+			t.Errorf("Set() failed: expected name to be 'Jane', got %v", data["name"])
+		}
+	})
+
+	t.Run("nested property", func(t *testing.T) {
+		data := map[string]interface{}{
+			"user": map[string]interface{}{
+				"name": "John",
+				"age":  30,
+			},
+		}
+		success := Set(data, "$.user.name", "Jane")
+
+		if !success {
+			t.Errorf("Set() returned false, expected true")
+		}
+		user, ok := data["user"].(map[string]interface{})
+		if !ok {
+			t.Fatalf("User is not a map")
+		}
+		if user["name"] != "Jane" {
+			t.Errorf("Set() failed: expected user.name to be 'Jane', got %v", user["name"])
+		}
+	})
+
+	t.Run("array element", func(t *testing.T) {
+		data := map[string]interface{}{
+			"users": []interface{}{
+				map[string]interface{}{"name": "John", "age": 30},
+				map[string]interface{}{"name": "Jane", "age": 25},
+			},
+		}
+		success := Set(data, "$.users[0].name", "Bob")
+
+		if !success {
+			t.Errorf("Set() returned false, expected true")
+		}
+		users, ok := data["users"].([]interface{})
+		if !ok {
+			t.Fatalf("Users is not a slice")
+		}
+		user0, ok := users[0].(map[string]interface{})
+		if !ok {
+			t.Fatalf("User is not a map")
+		}
+		if user0["name"] != "Bob" {
+			t.Errorf("Set() failed: expected users[0].name to be 'Bob', got %v", user0["name"])
+		}
+	})
+
+	t.Run("complex value", func(t *testing.T) {
+		data := map[string]interface{}{
+			"user": map[string]interface{}{
+				"name": "John",
+				"profile": map[string]interface{}{
+					"email": "john@example.com",
+				},
+			},
+		}
+
+		newProfile := map[string]interface{}{
+			"email": "john.doe@example.com",
+			"phone": "123-456-7890",
+		}
+
+		success := Set(data, "$.user.profile", newProfile)
+
+		if !success {
+			t.Errorf("Set() returned false, expected true")
+		}
+		userMap, ok := data["user"].(map[string]interface{})
+		if !ok {
+			t.Fatalf("User is not a map")
+		}
+
+		profile, ok := userMap["profile"].(map[string]interface{})
+		if !ok {
+			t.Fatalf("Profile is not a map")
+		}
+
+		if profile["email"] != "john.doe@example.com" || profile["phone"] != "123-456-7890" {
+			t.Errorf("Set() failed: expected profile to be updated with new values")
+		}
+	})
+
+	t.Run("create new property", func(t *testing.T) {
+		data := map[string]interface{}{
+			"user": map[string]interface{}{
+				"name": "John",
+			},
+		}
+
+		success := Set(data, "$.user.email", "john@example.com")
+
+		if !success {
+			t.Errorf("Set() returned false, expected true")
+		}
+		userMap, ok := data["user"].(map[string]interface{})
+		if !ok {
+			t.Fatalf("User is not a map")
+		}
+
+		if email, exists := userMap["email"]; !exists || email != "john@example.com" {
+			t.Errorf("Set() failed: expected user.email to be 'john@example.com', got %v", userMap["email"])
+		}
+	})
+
+	t.Run("create nested properties", func(t *testing.T) {
+		data := map[string]interface{}{
+			"user": map[string]interface{}{
+				"name": "John",
+			},
+		}
+
+		success := Set(data, "$.user.contact.email", "john@example.com")
+
+		if !success {
+			t.Errorf("Set() returned false, expected true")
+		}
+		userMap, ok := data["user"].(map[string]interface{})
+		if !ok {
+			t.Fatalf("User is not a map")
+		}
+
+		contact, ok := userMap["contact"].(map[string]interface{})
+		if !ok {
+			t.Fatalf("Contact is not a map")
+		}
+
+		if email, exists := contact["email"]; !exists || email != "john@example.com" {
+			t.Errorf("Set() failed: expected user.contact.email to be 'john@example.com', got %v", contact["email"])
+		}
+	})
+
+	t.Run("create array and element", func(t *testing.T) {
+		data := map[string]interface{}{
+			"user": map[string]interface{}{
+				"name": "John",
+			},
+		}
+
+		// This should create an empty addresses array, but won't be able to set index 0
+		// since the array is empty
+		success := Set(data, "$.user.addresses[0].street", "123 Main St")
+
+		// This shouldn't succeed because we can't create array elements that don't exist
+		if success {
+			t.Errorf("Set() returned true, expected false for out-of-bounds array index")
+		}
+	})
+
+	t.Run("multiple targets (should only update first)", func(t *testing.T) {
+		data := map[string]interface{}{
+			"users": []interface{}{
+				map[string]interface{}{"active": true},
+				map[string]interface{}{"active": true},
+			},
+		}
+
+		success := Set(data, "$.users[*].active", false)
+
+		if !success {
+			t.Errorf("Set() returned false, expected true")
+		}
+
+		users, ok := data["users"].([]interface{})
+		if !ok {
+			t.Fatalf("Users is not a slice")
+		}
+
+		user0, ok := users[0].(map[string]interface{})
+		if !ok {
+			t.Fatalf("User0 is not a map")
+		}
+
+		user1, ok := users[1].(map[string]interface{})
+		if !ok {
+			t.Fatalf("User1 is not a map")
+		}
+
+		// Only the first one should be changed
+		if active, exists := user0["active"]; !exists || active != false {
+			t.Errorf("Set() failed: expected users[0].active to be false, got %v", user0["active"])
+		}
+
+		// The second one should remain unchanged
+		if active, exists := user1["active"]; !exists || active != true {
+			t.Errorf("Set() incorrectly modified users[1].active: expected true, got %v", user1["active"])
+		}
+	})
+
+	t.Run("setting on root should fail", func(t *testing.T) {
+		data := map[string]interface{}{
+			"name": "John",
+		}
+
+		success := Set(data, "$", "Jane")
+
+		if success {
+			t.Errorf("Set() returned true, expected false for setting on root")
+		}
+
+		// Data should be unchanged
+		if data["name"] != "John" {
+			t.Errorf("Data was modified when setting on root")
+		}
+	})
+}
+
+func TestSetAll(t *testing.T) {
+	t.Run("simple property", func(t *testing.T) {
+		data := map[string]interface{}{
+			"name": "John",
+			"age":  30,
+		}
+		success := SetAll(data, "$.name", "Jane")
+
+		if !success {
+			t.Errorf("SetAll() returned false, expected true")
+		}
+		if data["name"] != "Jane" {
+			t.Errorf("SetAll() failed: expected name to be 'Jane', got %v", data["name"])
+		}
+	})
+
+	t.Run("all array elements", func(t *testing.T) {
+		data := map[string]interface{}{
+			"users": []interface{}{
+				map[string]interface{}{"active": true},
+				map[string]interface{}{"active": true},
+			},
+		}
+
+		success := SetAll(data, "$.users[*].active", false)
+
+		if !success {
+			t.Errorf("SetAll() returned false, expected true")
+		}
+
+		users, ok := data["users"].([]interface{})
+		if !ok {
+			t.Fatalf("Users is not a slice")
+		}
+
+		// Both elements should be updated
+		for i, user := range users {
+			userMap, ok := user.(map[string]interface{})
+			if !ok {
+				t.Fatalf("User%d is not a map", i)
+			}
+
+			if active, exists := userMap["active"]; !exists || active != false {
+				t.Errorf("SetAll() failed: expected users[%d].active to be false, got %v", i, userMap["active"])
+			}
+		}
+	})
+
+	t.Run("recursive descent", func(t *testing.T) {
+		data := map[string]interface{}{
+			"user": map[string]interface{}{
+				"profile": map[string]interface{}{
+					"active": true,
+				},
+			},
+			"admin": map[string]interface{}{
+				"profile": map[string]interface{}{
+					"active": true,
+				},
+			},
+		}
+
+		success := SetAll(data, "$..active", false)
+
+		if !success {
+			t.Errorf("SetAll() returned false, expected true")
+		}
+
+		// Check user profile
+		userProfile, ok := data["user"].(map[string]interface{})["profile"].(map[string]interface{})
+		if !ok {
+			t.Fatalf("Failed to access user.profile")
+		}
+		if active, exists := userProfile["active"]; !exists || active != false {
+			t.Errorf("SetAll() didn't update user.profile.active, got: %v", active)
+		}
+
+		// Check admin profile
+		adminProfile, ok := data["admin"].(map[string]interface{})["profile"].(map[string]interface{})
+		if !ok {
+			t.Fatalf("Failed to access admin.profile")
+		}
+		if active, exists := adminProfile["active"]; !exists || active != false {
+			t.Errorf("SetAll() didn't update admin.profile.active, got: %v", active)
+		}
+	})
+}

processor/jsonpath/jsonpath_test.go

@@ -158,12 +158,7 @@ func TestEvaluator(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			steps, err := ParseJSONPath(tt.path)
-			if err != nil {
-				t.Fatalf("Parse error: %v", err)
-			}
-
-			result := EvaluateJSONPath(testData, steps)
+			result := Get(testData, tt.path)
 
 			// Special handling for wildcard recursive test
 			if tt.name == "wildcard_recursive" {
@@ -187,18 +182,13 @@ func TestEvaluator(t *testing.T) {
 					t.Errorf("Expected %v, got %v", tt.expected, resultItem)
 				}
 			}
-
-			// if !reflect.DeepEqual(result, tt.expected) {
-			// 	t.Errorf("EvaluateJSONPath() = %v, want %v", result, tt.expected)
-			// }
 		})
 	}
 }
 
 func TestEdgeCases(t *testing.T) {
 	t.Run("empty_data", func(t *testing.T) {
-		steps, _ := ParseJSONPath("$.a.b")
-		result := EvaluateJSONPath(nil, steps)
+		result := Get(nil, "$.a.b")
 		if len(result) > 0 {
 			t.Errorf("Expected empty result, got %v", result)
 		}
@@ -215,9 +205,8 @@ func TestEdgeCases(t *testing.T) {
 		data := map[string]interface{}{
 			"42": "answer",
 		}
-		steps, _ := ParseJSONPath("$.42")
-		result := EvaluateJSONPath(data, steps)
-		if result[0] != "answer" {
+		result := Get(data, "$.42")
+		if len(result) == 0 || result[0] != "answer" {
 			t.Errorf("Expected 'answer', got %v", result)
 		}
 	})