Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support custom functions #375

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions api/v1alpha1/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,10 @@ type ResourceGraphDefinitionSpec struct {
//
// +kubebuilder:validation:Optional
Resources []*Resource `json:"resources,omitempty"`

// Functions defines custom CEL functions that can be used in expressions
// +kubebuilder:validation:Optional
Functions []Function `json:"functions,omitempty"`
// ServiceAccount configuration for controller impersonation.
// Key is the namespace, value is the service account name to use.
// Special key "*" defines the default service account for any
Expand All @@ -45,6 +49,18 @@ type ResourceGraphDefinitionSpec struct {
DefaultServiceAccounts map[string]string `json:"defaultServiceAccounts,omitempty"`
}

type Function struct {
// The unique identifier for the function
ID string `json:"id"`

// Inputs defines the expected input types
// +kubebuilder:validation:MinItems=1
Inputs []string `json:"inputs"`

// The CEL expression template that defines the function logic
Template string `json:"template"`
}

// Schema represents the attributes that define an instance of
// a resourcegraphdefinition.
type Schema struct {
Expand Down
4 changes: 4 additions & 0 deletions pkg/cel/environment.go
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,12 @@ func DefaultEnvironment(options ...EnvOption) (*cel.Env, error) {
ext.Strings(),
}

// Add resource IDs
for _, name := range opts.resourceIDs {
declarations = append(declarations, cel.Variable(name, cel.AnyType))
}

declarations = append(declarations, opts.customDeclarations...)

return cel.NewEnv(declarations...)
}
108 changes: 108 additions & 0 deletions pkg/cel/function/function.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
package function

import (
"fmt"

"github.com/google/cel-go/cel"
"github.com/google/cel-go/common/types"
"github.com/google/cel-go/common/types/ref"
xv1alpha1 "github.com/kro-run/kro/api/v1alpha1"
)

// FunctionRegistry manages custom CEL functions
type FunctionRegistry struct {
functions map[string]*xv1alpha1.Function
}

// NewFunctionRegistry creates a new function registry
func NewFunctionRegistry() *FunctionRegistry {
return &FunctionRegistry{
functions: make(map[string]*xv1alpha1.Function),
}
}

// RegisterFunction adds a function to the registry
func (r *FunctionRegistry) RegisterFunction(fn *xv1alpha1.Function) error {
if _, exists := r.functions[fn.ID]; exists {
return fmt.Errorf("function %s already registered", fn.ID)
}
r.functions[fn.ID] = fn
return nil
}

// GetCELFunction converts a Function to a CEL function declaration
func (r *FunctionRegistry) GetCELFunction(fn *xv1alpha1.Function) (cel.EnvOption, error) {

//convert input types to CEL types
celTypes := make([]*cel.Type, len((fn.Inputs)))
envOpts := make([]cel.EnvOption, len(fn.Inputs))

for i, inputType := range fn.Inputs {
celType, err := getCelType(inputType)
if err != nil {
return nil, err
}
celTypes[i] = celType
envOpts[i] = cel.Variable(fmt.Sprintf("_%d", i), celType)
}

// Create CEL environment for function template
env, err := cel.NewEnv(envOpts...)
if err != nil {
return nil, fmt.Errorf("environment creation failed: %w", err)
}

// Compile function template
ast, iss := env.Compile(fn.Template)
if iss.Err() != nil {
return nil, fmt.Errorf("template compilation failed: %w", iss.Err())
}

prg, err := env.Program(ast)
if err != nil {
return nil, fmt.Errorf("program creation failed: %w", err)
}

// Return CEL function declaration
return cel.Function(fn.ID,
cel.Overload(fn.ID,
celTypes,
ast.OutputType(),
cel.FunctionBinding(func(args ...ref.Val) ref.Val {
vars := map[string]any{}
for i, arg := range args {
vars[fmt.Sprintf("_%d", i)] = arg
}
out, _, err := prg.Eval(vars)
if err != nil {
return types.WrapErr(err)
}
return out
}),
),
), nil

}

func getCelType(typeName string) (*cel.Type, error) {
switch typeName {
case "string":
return cel.StringType, nil
case "int":
return cel.IntType, nil
case "bool":
return cel.BoolType, nil
case "double", "float":
return cel.DoubleType, nil
case "bytes":
return cel.BytesType, nil
case "any":
return cel.AnyType, nil
case "list":
return cel.ListType(cel.AnyType), nil
case "map":
return cel.MapType(cel.StringType, cel.AnyType), nil
Comment on lines +101 to +104
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we're ok with making backwards incompatible changes to this (@a-hilaly your input please), we can start including list and map. If not, I suggest we start with just the primitives and add support for lists and maps later, probably when the design for complex types #144 is finished

default:
return nil, fmt.Errorf("unsupported type: %s", typeName)
}
}