From f7ffbdb792fbfd3623ad2b61ade1af7e3e1988ee Mon Sep 17 00:00:00 2001
From: pk910 <philipp@pk910.de>
Date: Mon, 5 Feb 2024 09:14:41 +0100
Subject: [PATCH] use `gojq` to handle variable assignments in a jq like way

---
 go.mod                            |  4 ++-
 go.sum                            |  5 ++++
 pkg/coordinator/types/vars.go     |  1 +
 pkg/coordinator/vars/variables.go | 42 ++++++++++++++++++++++++++++---
 4 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/go.mod b/go.mod
index 38da9af..0908f5f 100644
--- a/go.mod
+++ b/go.mod
@@ -51,6 +51,8 @@ require (
 	github.com/gorilla/websocket v1.4.2 // indirect
 	github.com/huandu/go-clone v1.6.0 // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
+	github.com/itchyny/gojq v0.12.14 // indirect
+	github.com/itchyny/timefmt-go v0.1.5 // indirect
 	github.com/kilic/bls12-381 v0.1.0 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.6 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
@@ -88,4 +90,4 @@ require (
 	google.golang.org/protobuf v1.31.0 // indirect
 	gopkg.in/cenkalti/backoff.v1 v1.1.0 // indirect
 	rsc.io/tmplfunc v0.0.3 // indirect
-)
+)
\ No newline at end of file
diff --git a/go.sum b/go.sum
index ecac48c..128e122 100644
--- a/go.sum
+++ b/go.sum
@@ -136,6 +136,10 @@ github.com/huin/goupnp v1.3.0 h1:UvLUlWDNpoUdYzb2TCn+MuTWtcjXKSza2n6CBdQ0xXc=
 github.com/huin/goupnp v1.3.0/go.mod h1:gnGPsThkYa7bFi/KWmEysQRf48l2dvR5bxr2OFckNX8=
 github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
+github.com/itchyny/gojq v0.12.14 h1:6k8vVtsrhQSYgSGg827AD+PVVaB1NLXEdX+dda2oZCc=
+github.com/itchyny/gojq v0.12.14/go.mod h1:y1G7oO7XkcR1LPZO59KyoCRy08T3j9vDYRV0GgYSS+s=
+github.com/itchyny/timefmt-go v0.1.5 h1:G0INE2la8S6ru/ZI5JecgyzbbJNs5lG1RcBqa7Jm6GE=
+github.com/itchyny/timefmt-go v0.1.5/go.mod h1:nEP7L+2YmAbT2kZ2HfSs1d8Xtw9LY8D2stDBckWakZ8=
 github.com/jackpal/go-nat-pmp v1.0.2 h1:KzKSgb7qkJvOUTqYl9/Hg/me3pWgBmERKrTGD7BdWus=
 github.com/jackpal/go-nat-pmp v1.0.2/go.mod h1:QPH045xvCAeXUZOxsnwmrtiCoxIr9eob+4orBN1SBKc=
 github.com/kilic/bls12-381 v0.1.0 h1:encrdjqKMEvabVQ7qYOKu1OvhqpK4s47wDYtNiPtlp4=
@@ -167,6 +171,7 @@ github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWE
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mattn/go-runewidth v0.0.13 h1:lTGmDsbAYt5DmK6OnoV7EuIF1wEIFAcxld6ypU4OSgU=
 github.com/mattn/go-runewidth v0.0.13/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
+github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U=
 github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo=
 github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4=
 github.com/minio/sha256-simd v0.1.0/go.mod h1:2FMWW+8GMoPweT6+pI63m9YE3Lmw4J71hV56Chs1E/U=
diff --git a/pkg/coordinator/types/vars.go b/pkg/coordinator/types/vars.go
index 994fb5d..b10594f 100644
--- a/pkg/coordinator/types/vars.go
+++ b/pkg/coordinator/types/vars.go
@@ -5,6 +5,7 @@ type Variables interface {
 	LookupVar(name string) (interface{}, bool)
 	SetVar(name string, value interface{})
 	NewScope() Variables
+	GetVarsMap() map[string]any
 	ResolvePlaceholders(str string) string
 	ConsumeVars(config interface{}, consumeMap map[string]string) error
 	CopyVars(source Variables, copyMap map[string]string)
diff --git a/pkg/coordinator/vars/variables.go b/pkg/coordinator/vars/variables.go
index ccc3d1e..c390a41 100644
--- a/pkg/coordinator/vars/variables.go
+++ b/pkg/coordinator/vars/variables.go
@@ -1,11 +1,14 @@
 package vars
 
 import (
+	"context"
 	"fmt"
 	"regexp"
 	"sync"
+	"time"
 
 	"github.com/ethpandaops/assertoor/pkg/coordinator/types"
+	"github.com/itchyny/gojq"
 	"gopkg.in/yaml.v2"
 )
 
@@ -82,16 +85,47 @@ func (v *Variables) ResolvePlaceholders(str string) string {
 	})
 }
 
+func (v *Variables) GetVarsMap() map[string]any {
+	var varsMap map[string]any
+
+	if v.parentScope != nil {
+		varsMap = v.parentScope.GetVarsMap()
+	} else {
+		varsMap = map[string]any{}
+	}
+
+	for varName, varData := range v.varsMap {
+		varsMap[varName] = varData.value
+	}
+
+	return varsMap
+}
+
 func (v *Variables) ConsumeVars(config interface{}, consumeMap map[string]string) error {
 	applyMap := map[string]interface{}{}
 
-	for cfgName, varName := range consumeMap {
-		varValue, varFound := v.LookupVar(varName)
-		if !varFound {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	varsMap := v.GetVarsMap()
+
+	for cfgName, varQuery := range consumeMap {
+		queryStr := fmt.Sprintf(".%v", varQuery)
+
+		query, err := gojq.Parse(queryStr)
+		if err != nil {
+			return fmt.Errorf("could not parse variable query '%v': %v", queryStr, err)
+		}
+
+		iter := query.RunWithContext(ctx, varsMap)
+
+		val, ok := iter.Next()
+		if !ok {
+			// no query result, skip variable assignment
 			continue
 		}
 
-		applyMap[cfgName] = varValue
+		applyMap[cfgName] = val
 	}
 
 	// apply to config by generating a yaml, which is then parsed with the target config types.