From 5e411852958021e507decb9cca6b37344aa103a3 Mon Sep 17 00:00:00 2001
From: Tamal Saha <tamal@appscode.com>
Date: Sun, 8 Sep 2024 12:47:50 -0700
Subject: [PATCH] Support jmespath for references with jmes: prefix

Signed-off-by: Tamal Saha <tamal@appscode.com>
---
 pkg/graph/lib.go | 96 +++++++++++++++++++++++++++++-------------------
 1 file changed, 59 insertions(+), 37 deletions(-)

diff --git a/pkg/graph/lib.go b/pkg/graph/lib.go
index 3681063c5..23f903440 100644
--- a/pkg/graph/lib.go
+++ b/pkg/graph/lib.go
@@ -26,6 +26,7 @@ import (
 	"sort"
 	"strings"
 
+	"github.com/jmespath/go-jmespath"
 	"gomodules.xyz/jsonpath"
 	core "k8s.io/api/core/v1"
 	kerr "k8s.io/apimachinery/pkg/api/errors"
@@ -304,25 +305,19 @@ func (finder ObjectFinder) ResourcesFor(src *unstructured.Unstructured, e *Edge)
 			// TODO: check that namespacePath must be empty
 
 			var out []*unstructured.Unstructured
-
 			for _, reference := range e.Connection.References {
-				j := jsonpath.New("jsonpath")
-				j.AllowMissingKeys(true)
-				err := j.Parse(reference)
-				if err != nil {
-					return nil, fmt.Errorf("fails to parse reference %q between %s -> %s. err:%v", e.Connection.References, e.Src, e.Dst, err)
-				}
-				buf := new(bytes.Buffer)
-				err = j.Execute(buf, src.Object)
-				if err != nil {
-					return nil, fmt.Errorf("fails to execute reference %q between %s -> %s. err:%v", e.Connection.References, e.Src, e.Dst, err)
-				}
-				r := csv.NewReader(buf)
-				// Mapper.Comma = ';'
-				r.Comment = '#'
-				records, err := r.ReadAll()
-				if err != nil {
-					return nil, err
+				var records [][]string
+				var err error
+				if strings.HasPrefix(reference, "jmes:") {
+					records, err = execJmesPath(src.Object, reference[5:], e.Src, e.Dst)
+					if err != nil {
+						return nil, err
+					}
+				} else {
+					records, err = execJsonPath(src.Object, reference, e.Src, e.Dst)
+					if err != nil {
+						return nil, err
+					}
 				}
 				refs, err := ParseResourceRefs(records)
 				if err != nil {
@@ -509,25 +504,18 @@ func (finder ObjectFinder) ResourcesFor(src *unstructured.Unstructured, e *Edge)
 				rs := result.Items[i]
 
 				for _, reference := range e.Connection.References {
-
-					j := jsonpath.New("jsonpath")
-					j.AllowMissingKeys(true)
-					err := j.Parse(reference)
-					if err != nil {
-						return nil, fmt.Errorf("fails to parse reference %q between %s -> %s. err:%v", e.Connection.References, e.Src, e.Dst, err)
-					}
-
-					buf := new(bytes.Buffer)
-					err = j.Execute(buf, rs.Object)
-					if err != nil {
-						return nil, fmt.Errorf("fails to execute reference %q between %s -> %s. err:%v", e.Connection.References, e.Src, e.Dst, err)
-					}
-					r := csv.NewReader(buf)
-					// Mapper.Comma = ';'
-					r.Comment = '#'
-					records, err := r.ReadAll()
-					if err != nil {
-						return nil, err
+					var records [][]string
+					var err error
+					if strings.HasPrefix(reference, "jmes:") {
+						records, err = execJmesPath(rs.Object, reference[5:], e.Src, e.Dst)
+						if err != nil {
+							return nil, err
+						}
+					} else {
+						records, err = execJsonPath(rs.Object, reference, e.Src, e.Dst)
+						if err != nil {
+							return nil, err
+						}
 					}
 					refs, err := ParseResourceRefs(records)
 					if err != nil {
@@ -581,6 +569,40 @@ func (finder ObjectFinder) ResourcesFor(src *unstructured.Unstructured, e *Edge)
 	return nil, nil
 }
 
+func execJmesPath(data any, reference string, src, dst schema.GroupVersionKind) ([][]string, error) {
+	result, err := jmespath.Search(reference, data)
+	if err != nil {
+		return nil, fmt.Errorf("fails to execute jmes reference %q between %s -> %s. err:%w", reference, src, dst, err)
+	}
+	switch v := result.(type) {
+	case string:
+		return [][]string{{v}}, nil
+	case []string:
+		return [][]string{v}, nil
+	case [][]string:
+		return v, nil
+	}
+	return nil, fmt.Errorf("invalid result type %T for jmes reference %q between %s -> %s", result, reference, src, dst)
+}
+
+func execJsonPath(data any, reference string, src, dst schema.GroupVersionKind) ([][]string, error) {
+	j := jsonpath.New("jsonpath")
+	j.AllowMissingKeys(true)
+	err := j.Parse(reference)
+	if err != nil {
+		return nil, fmt.Errorf("fails to parse reference %q between %s -> %s. err:%v", reference, src, dst, err)
+	}
+	buf := new(bytes.Buffer)
+	err = j.Execute(buf, data)
+	if err != nil {
+		return nil, fmt.Errorf("fails to execute reference %q between %s -> %s. err:%v", reference, src, dst, err)
+	}
+	r := csv.NewReader(buf)
+	// Mapper.Comma = ';'
+	r.Comment = '#'
+	return r.ReadAll()
+}
+
 func isConnected(conn rsapi.OwnershipLevel, obj *unstructured.Unstructured, owner *unstructured.Unstructured) bool {
 	switch conn {
 	case rsapi.Controller: