From a4b652d6f288b529d0f0d86c621dd0ba4e44cbda Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Charles-Edouard=20Br=C3=A9t=C3=A9ch=C3=A9?=
 <charles.edouard@nirmata.com>
Date: Tue, 8 Oct 2024 11:36:28 +0200
Subject: [PATCH] feat: add custom cel env support
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Charles-Edouard Brétéché <charles.edouard@nirmata.com>
---
 pkg/core/compilers/cel/cel.go      | 12 ++++++++----
 pkg/core/compilers/cel/cel_test.go | 13 +++++++++++++
 pkg/core/compilers/cel/env.go      | 20 +++++++++++++++++++-
 pkg/core/compilers/compilers.go    |  2 +-
 4 files changed, 41 insertions(+), 6 deletions(-)
 create mode 100644 pkg/core/compilers/cel/cel_test.go

diff --git a/pkg/core/compilers/cel/cel.go b/pkg/core/compilers/cel/cel.go
index e1967564..4f0ba83d 100644
--- a/pkg/core/compilers/cel/cel.go
+++ b/pkg/core/compilers/cel/cel.go
@@ -11,14 +11,18 @@ type Compiler interface {
 	Compile(string) (Program, error)
 }
 
-type compiler struct{}
+type compiler struct {
+	env func() (*cel.Env, error)
+}
 
-func NewCompiler() *compiler {
-	return &compiler{}
+func NewCompiler(env func() (*cel.Env, error)) *compiler {
+	return &compiler{
+		env: env,
+	}
 }
 
 func (c *compiler) Compile(statement string) (Program, error) {
-	env, err := DefaultEnv()
+	env, err := c.env()
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/core/compilers/cel/cel_test.go b/pkg/core/compilers/cel/cel_test.go
new file mode 100644
index 00000000..bed1f085
--- /dev/null
+++ b/pkg/core/compilers/cel/cel_test.go
@@ -0,0 +1,13 @@
+package cel
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_compiler_Compile(t *testing.T) {
+	c := NewCompiler(DefaultEnv)
+	_, err := c.Compile("object.?spec")
+	assert.NoError(t, err)
+}
diff --git a/pkg/core/compilers/cel/env.go b/pkg/core/compilers/cel/env.go
index e3e0be81..a43b8e48 100644
--- a/pkg/core/compilers/cel/env.go
+++ b/pkg/core/compilers/cel/env.go
@@ -6,12 +6,13 @@ import (
 	"github.com/google/cel-go/cel"
 	"github.com/google/cel-go/common/types"
 	"github.com/google/cel-go/common/types/ref"
+	"github.com/google/cel-go/ext"
 	"github.com/jmespath-community/go-jmespath/pkg/binding"
 )
 
 var (
 	BindingsType = cel.OpaqueType("bindings")
-	DefaultEnv   = sync.OnceValues(func() (*cel.Env, error) {
+	BaseEnv      = sync.OnceValues(func() (*cel.Env, error) {
 		return cel.NewEnv(
 			cel.Variable("object", cel.DynType),
 			cel.Variable("bindings", BindingsType),
@@ -38,4 +39,21 @@ var (
 			),
 		)
 	})
+	DefaultEnv = sync.OnceValues(func() (*cel.Env, error) {
+		if env, err := BaseEnv(); err != nil {
+			return nil, err
+		} else if env, err := env.Extend(
+			cel.HomogeneousAggregateLiterals(),
+			cel.EagerlyValidateDeclarations(true),
+			cel.DefaultUTCTimeZone(true),
+			cel.CrossTypeNumericComparisons(true),
+			cel.OptionalTypes(),
+			ext.Strings(),
+			ext.Sets(),
+		); err != nil {
+			return nil, err
+		} else {
+			return env, nil
+		}
+	})
 )
diff --git a/pkg/core/compilers/compilers.go b/pkg/core/compilers/compilers.go
index 2777e85b..a971fbde 100644
--- a/pkg/core/compilers/compilers.go
+++ b/pkg/core/compilers/compilers.go
@@ -13,7 +13,7 @@ const (
 
 var DefaultCompilers = Compilers{
 	Jp:  jp.NewCompiler(),
-	Cel: cel.NewCompiler(),
+	Cel: cel.NewCompiler(cel.DefaultEnv),
 }
 
 type Compilers struct {