feat!: support returning result in custom context key or environment

README.md

@@ -31,7 +31,7 @@ spec:
       spec:
         extraResources:
           - kind: XCluster
-            into: XCluster
+            toFieldPath: XCluster
             apiVersion: example.crossplane.io/v1
             type: Selector
             selector:

example/composition.yaml

@@ -17,7 +17,7 @@ spec:
       spec:
         extraResources:
           - kind: EnvironmentConfig
-            into: envConfs
+            toFieldPath: envConfs
             apiVersion: apiextensions.crossplane.io/v1alpha1
             type: Selector
             selector:
@@ -28,7 +28,7 @@ spec:
                   type: Value
                   value: cluster
           - kind: XCluster
-            into: XCluster
+            toFieldPath: XCluster
             apiVersion: example.crossplane.io/v1
             type: Selector
             selector:

fn.go

@@ -2,13 +2,13 @@ package main
 
 import (
 	"context"
-	"encoding/json"
+	"fmt"
 	"reflect"
 	"sort"
 
-	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/types/known/structpb"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 
 	"github.com/crossplane/crossplane-runtime/v2/pkg/errors"
 	"github.com/crossplane/crossplane-runtime/v2/pkg/fieldpath"
@@ -21,9 +21,10 @@ import (
 	"github.com/crossplane-contrib/function-extra-resources/input/v1beta1"
 )
 
-// Key to retrieve extras at.
 const (
-	FunctionContextKeyExtraResources = "apiextensions.crossplane.io/extra-resources"
+	// FunctionContextKeyEnvironment is a well-known Context key where the computed Environment
+	// will be stored, so that Crossplane v1 and other functions can access it, e.g. function-patch-and-transform.
+	FunctionContextKeyEnvironment = "apiextensions.crossplane.io/environment"
 )
 
 // Function returns whatever response you ask it to.
@@ -33,6 +34,11 @@ type Function struct {
 	log logging.Logger
 }
 
+type FetchedResult struct {
+	source    v1beta1.ResourceSource
+	resources []interface{}
+}
+
 // RunFunction runs the Function.
 func (f *Function) RunFunction(_ context.Context, req *fnv1.RunFunctionRequest) (*fnv1.RunFunctionResponse, error) {
 	f.log.Info("Running function", "tag", req.GetMeta().GetTag())
@@ -88,38 +94,100 @@ func (f *Function) RunFunction(_ context.Context, req *fnv1.RunFunctionRequest)
 		return rsp, nil
 	}
 
-	// For now cheaply convert to JSON for serializing.
-	//
-	// TODO(reedjosh): look into resources.AsStruct or simlar since unsturctured k8s objects are already almost json.
-	//    structpb.NewList(v []interface{}) should create an array like.
-	//    Combining this and similar structures from the structpb lib should should be done to create
-	//    a map[string][object] container into which the found extra resources can be dumped.
-	//
-	//    The found extra resources should then be directly marhsal-able via:
-	//    obj := &unstructured.Unstructured{}
-	//    obj.MarshalJSON()
-	b, err := json.Marshal(verifiedExtras)
+	var out *unstructured.Unstructured
+	var key string
+
+	t := in.Spec.Into.GetIntoType()
+	switch t {
+	case v1beta1.IntoTypeContext:
+		out, err = f.intoContext(verifiedExtras)
+		key = in.Spec.Into.GetIntoContextKey()
+	case v1beta1.IntoTypeEnvironment:
+		out, err = f.intoEnvironment(req, verifiedExtras)
+		key = FunctionContextKeyEnvironment
+	default:
+		err = errors.Errorf("unknown into type: %q", t)
+	}
+
 	if err != nil {
-		response.Fatal(rsp, errors.Errorf("cannot marshal %T: %w", verifiedExtras, err))
+		response.Fatal(rsp, err)
 		return rsp, nil
 	}
-	s := &structpb.Struct{}
-	err = protojson.Unmarshal(b, s)
+
+	s, err := resource.AsStruct(out)
 	if err != nil {
-		response.Fatal(rsp, errors.Errorf("cannot unmarshal %T into %T: %w", extraResources, s, err))
+		response.Fatal(rsp, errors.Wrap(err, "cannot convert unstructured to protobuf Struct well-known type"))
 		return rsp, nil
 	}
-	response.SetContextKey(rsp, FunctionContextKeyExtraResources, structpb.NewStructValue(s))
+
+	response.SetContextKey(rsp, key, structpb.NewStructValue(s))
 
 	return rsp, nil
 }
 
+func (f *Function) intoContext(verifiedExtras []FetchedResult) (*unstructured.Unstructured, error) {
+	out := &unstructured.Unstructured{Object: map[string]interface{}{}}
+	for _, extras := range verifiedExtras {
+		if toFieldPath := extras.source.ToFieldPath; toFieldPath != nil && *toFieldPath != "" {
+			if err := fieldpath.Pave(out.Object).SetValue(*toFieldPath, extras.resources); err != nil {
+				return nil, errors.Wrapf(err, "cannot set nested field path %q", *toFieldPath)
+			}
+		} else {
+			return nil, errors.New("must set toFieldPath for type Context")
+		}
+	}
+
+	return out, nil
+}
+
+func (f *Function) intoEnvironment(req *fnv1.RunFunctionRequest, verifiedExtras []FetchedResult) (*unstructured.Unstructured, error) {
+	var inputEnv *unstructured.Unstructured
+	if v, ok := request.GetContextKey(req, FunctionContextKeyEnvironment); ok {
+		inputEnv = &unstructured.Unstructured{}
+		if err := resource.AsObject(v.GetStructValue(), inputEnv); err != nil {
+			return nil, errors.Wrapf(err, "cannot get Composition environment from %T context key %q", req, FunctionContextKeyEnvironment)
+		}
+		f.log.Debug("Loaded Composition environment from Function context", "context-key", FunctionContextKeyEnvironment)
+	}
+
+	mergedData := map[string]interface{}{}
+	for _, extras := range verifiedExtras {
+		for _, extra := range extras.resources {
+			if toFieldPath := extras.source.ToFieldPath; toFieldPath != nil && *toFieldPath != "" {
+				d := map[string]interface{}{}
+				if err := fieldpath.Pave(d).SetValue(*toFieldPath, extra); err != nil {
+					return nil, errors.Wrapf(err, "cannot set nested field path %q", *toFieldPath)
+				}
+
+				mergedData = mergeMaps(mergedData, d)
+			} else if e, ok := extra.(map[string]interface{}); ok {
+				mergedData = mergeMaps(mergedData, e)
+			} else {
+				return nil, errors.New("must set toFieldPath when extracted value is not an object")
+			}
+		}
+	}
+
+	// merge input env if any
+	if inputEnv != nil {
+		mergedData = mergeMaps(inputEnv.Object, mergedData)
+	}
+
+	// build environment and return it in the response as context
+	out := &unstructured.Unstructured{Object: mergedData}
+	if out.GroupVersionKind().Empty() {
+		out.SetGroupVersionKind(schema.GroupVersionKind{Group: "internal.crossplane.io", Kind: "Environment", Version: "v1alpha1"})
+	}
+
+	return out, nil
+}
+
 // Build requirements takes input and outputs an array of external resoruce requirements to request
 // from Crossplane's external resource API.
 func buildRequirements(in *v1beta1.Input, xr *resource.Composite) (*fnv1.Requirements, error) { //nolint:gocyclo // Adding non-nil validations increases function complexity.
 	extraResources := make(map[string]*fnv1.ResourceSelector, len(in.Spec.ExtraResources))
-	for _, extraResource := range in.Spec.ExtraResources {
-		extraResName := extraResource.Into
+	for i, extraResource := range in.Spec.ExtraResources {
+		extraResName := fmt.Sprintf("resources-%d", i)
 		switch extraResource.Type {
 		case v1beta1.ResourceSourceTypeReference, "":
 			extraResources[extraResName] = &fnv1.ResourceSelector{
@@ -169,16 +237,36 @@ func buildRequirements(in *v1beta1.Input, xr *resource.Composite) (*fnv1.Require
 	return &fnv1.Requirements{Resources: extraResources}, nil
 }
 
+func mergeMaps(a, b map[string]interface{}) map[string]interface{} {
+	out := make(map[string]interface{}, len(a))
+	for k, v := range a {
+		out[k] = v
+	}
+	for k, v := range b {
+		if v, ok := v.(map[string]interface{}); ok {
+			if bv, ok := out[k]; ok {
+				if bv, ok := bv.(map[string]interface{}); ok {
+					out[k] = mergeMaps(bv, v)
+					continue
+				}
+			}
+		}
+		out[k] = v
+	}
+	return out
+}
+
 // Verify Min/Max and sort extra resources by field path within a single kind.
 func verifyAndSortExtras(in *v1beta1.Input, extraResources map[string][]resource.Required, //nolint:gocyclo // TODO(reedjosh): refactor
-) (cleanedExtras map[string][]unstructured.Unstructured, err error) {
-	cleanedExtras = make(map[string][]unstructured.Unstructured)
-	for _, extraResource := range in.Spec.ExtraResources {
-		extraResName := extraResource.Into
+) ([]FetchedResult, error) {
+	results := []FetchedResult{}
+	for i, extraResource := range in.Spec.ExtraResources {
+		extraResName := fmt.Sprintf("resources-%d", i)
 		resources, ok := extraResources[extraResName]
 		if !ok {
 			return nil, errors.Errorf("cannot find expected extra resource %q", extraResName)
 		}
+
 		switch extraResource.GetType() {
 		case v1beta1.ResourceSourceTypeReference:
 			if len(resources) == 0 {
@@ -190,7 +278,6 @@ func verifyAndSortExtras(in *v1beta1.Input, extraResources map[string][]resource
 			if len(resources) > 1 {
 				return nil, errors.Errorf("expected exactly one extra resource %q, got %d", extraResName, len(resources))
 			}
-			cleanedExtras[extraResName] = append(cleanedExtras[extraResName], *resources[0].Resource)
 
 		case v1beta1.ResourceSourceTypeSelector:
 			selector := extraResource.Selector
@@ -203,12 +290,28 @@ func verifyAndSortExtras(in *v1beta1.Input, extraResources map[string][]resource
 			if selector.MaxMatch != nil && uint64(len(resources)) > *selector.MaxMatch {
 				resources = resources[:*selector.MaxMatch]
 			}
-			for _, r := range resources {
-				cleanedExtras[extraResName] = append(cleanedExtras[extraResName], *r.Resource)
+		}
+
+		result := FetchedResult{source: extraResource}
+		for _, r := range resources {
+			if path := extraResource.FromFieldPath; path != nil {
+				if *path == "" {
+					return nil, errors.New("fromFieldPath cannot be empty, omit the field to get the whole object")
+				}
+
+				// Extract part of the object, from `FromFieldPath`.
+				object, err := fieldpath.Pave(r.Resource.Object).GetValue(*path)
+				if err != nil {
+					return nil, err
+				}
+				result.resources = append(result.resources, object)
+			} else {
+				result.resources = append(result.resources, r.Resource.Object)
 			}
 		}
+		results = append(results, result)
 	}
-	return cleanedExtras, nil
+	return results, nil
 }
 
 // Sort extra resources by field path within a single kind.