feat: configurable context key in which to put resolved extra resources

fn.go

@@ -21,11 +21,6 @@ import (
 	"github.com/crossplane-contrib/function-extra-resources/input/v1beta1"
 )
 
-// Key to retrieve extras at.
-const (
-	FunctionContextKeyExtraResources = "apiextensions.crossplane.io/extra-resources"
-)
-
 // Function returns whatever response you ask it to.
 type Function struct {
 	fnv1.UnimplementedFunctionRunnerServiceServer
@@ -109,7 +104,7 @@ func (f *Function) RunFunction(_ context.Context, req *fnv1.RunFunctionRequest)
 		response.Fatal(rsp, errors.Errorf("cannot unmarshal %T into %T: %w", extraResources, s, err))
 		return rsp, nil
 	}
-	response.SetContextKey(rsp, FunctionContextKeyExtraResources, structpb.NewStructValue(s))
+	response.SetContextKey(rsp, in.Spec.Context.GetKey(), structpb.NewStructValue(s))
 
 	return rsp, nil
 }

fn_test.go

@@ -17,6 +17,8 @@ import (
 	fnv1 "github.com/crossplane/function-sdk-go/proto/v1"
 	"github.com/crossplane/function-sdk-go/resource"
 	"github.com/crossplane/function-sdk-go/response"
+
+	"github.com/crossplane-contrib/function-extra-resources/input/v1beta1"
 )
 
 func TestRunFunction(t *testing.T) {
@@ -458,7 +460,7 @@ func TestRunFunction(t *testing.T) {
 					},
 					Context: &structpb.Struct{
 						Fields: map[string]*structpb.Value{
-							FunctionContextKeyExtraResources: structpb.NewStructValue(resource.MustStructJSON(`{
+							v1beta1.FunctionContextKeyExtraResources: structpb.NewStructValue(resource.MustStructJSON(`{
 									"obj-0": [
         							    {
         							        "apiVersion": "apiextensions.crossplane.io/v1beta1",
@@ -600,6 +602,92 @@ func TestRunFunction(t *testing.T) {
 				},
 			},
 		},
+		"CustomContextKey": {
+			reason: "The Function should put resolved extra resources into custom context key when specified.",
+			args: args{
+				req: &fnv1.RunFunctionRequest{
+					Meta: &fnv1.RequestMeta{Tag: "hello"},
+					Observed: &fnv1.State{
+						Composite: &fnv1.Resource{
+							Resource: resource.MustStructJSON(`{
+								"apiVersion": "test.crossplane.io/v1alpha1",
+								"kind": "XR",
+								"metadata": {
+									"name": "my-xr"
+								}
+							}`),
+						},
+					},
+					RequiredResources: map[string]*fnv1.Resources{
+						"obj-0": {
+							Items: []*fnv1.Resource{
+								{
+									Resource: resource.MustStructJSON(`{
+									"apiVersion": "apiextensions.crossplane.io/v1beta1",
+									"kind": "EnvironmentConfig",
+									"metadata": {
+										"name": "my-env-config"
+									}
+								}`),
+								},
+							},
+						},
+					},
+					Input: resource.MustStructJSON(`{
+						"apiVersion": "extra-resources.fn.crossplane.io/v1beta1",
+						"kind": "Input",
+						"spec": {
+							"context": {
+								"key": "apiextensions.crossplane.io/environment"
+							},
+							"extraResources": [
+								{
+									"kind": "EnvironmentConfig",
+									"apiVersion": "apiextensions.crossplane.io/v1beta1",
+									"type": "Reference",
+									"into": "obj-0",
+									"ref": {	
+										"name": "my-env-config"
+									}
+								}
+							]
+						}
+					}`),
+				},
+			},
+			want: want{
+				rsp: &fnv1.RunFunctionResponse{
+					Meta:    &fnv1.ResponseMeta{Tag: "hello", Ttl: durationpb.New(response.DefaultTTL)},
+					Results: []*fnv1.Result{},
+					Requirements: &fnv1.Requirements{
+						Resources: map[string]*fnv1.ResourceSelector{
+							"obj-0": {
+								ApiVersion: "apiextensions.crossplane.io/v1beta1",
+								Kind:       "EnvironmentConfig",
+								Match: &fnv1.ResourceSelector_MatchName{
+									MatchName: "my-env-config",
+								},
+							},
+						},
+					},
+					Context: &structpb.Struct{
+						Fields: map[string]*structpb.Value{
+							"apiextensions.crossplane.io/environment": structpb.NewStructValue(resource.MustStructJSON(`{
+								"obj-0": [
+									{
+										"apiVersion": "apiextensions.crossplane.io/v1beta1",
+										"kind": "EnvironmentConfig",
+										"metadata": {
+											"name": "my-env-config"
+										}
+									}
+								]
+							}`)),
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for name, tc := range cases {

input/v1beta1/resource_select.go

@@ -20,8 +20,16 @@ import (
 	xpv1 "github.com/crossplane/crossplane-runtime/v2/apis/common/v1"
 )
 
+const (
+	// FunctionContextKeyExtraResources is the default context key.
+	FunctionContextKeyExtraResources = "apiextensions.crossplane.io/extra-resources"
+)
+
 // An InputSpec specifies extra resource(s) for rendering composed resources.
 type InputSpec struct {
+	// Context specifies how the function uses the response context.
+	Context *Context `json:"context,omitempty"`
+
 	// ExtraResources selects a list of `ExtraResource`s. The resolved
 	// resources are stored in the composite resource at
 	// `spec.extraResourceRefs` and is only updated if it is null.
@@ -33,6 +41,24 @@ type InputSpec struct {
 	Policy *Policy `json:"policy,omitempty"`
 }
 
+// A Context specifies how the function uses the response context.
+type Context struct {
+	// Key specifies the context key in which to put resolved extra resources.
+	// E.g. 'apiextensions.crossplane.io/environment', the environment used in
+	// standard functions such as Function Patch and Transform.
+	// +kubebuilder:default=apiextensions.crossplane.io/extra-resources
+	Key *string `json:"key,omitempty"`
+}
+
+// GetKey returns the key of the context, defaulting to
+// FunctionContextKeyExtraResources if not specified.
+func (i *Context) GetKey() string {
+	if i == nil || i.Key == nil {
+		return FunctionContextKeyExtraResources
+	}
+	return *i.Key
+}
+
 // Policy represents the Resolution policy of Reference instance.
 type Policy struct {
 	// Resolution specifies whether resolution of this reference is required.

package/input/extra-resources.fn.crossplane.io_inputs.yaml

@@ -41,6 +41,18 @@ spec:
           spec:
             description: Spec is the input to this function.
             properties:
+              context:
+                description: Context specifies how the function uses the response
+                  context.
+                properties:
+                  key:
+                    default: apiextensions.crossplane.io/extra-resources
+                    description: |-
+                      Key specifies the context key in which to put resolved extra resources.
+                      E.g. 'apiextensions.crossplane.io/environment', the environment used in
+                      standard functions such as Function Patch and Transform.
+                    type: string
+                type: object
               extraResources:
                 description: |-
                   ExtraResources selects a list of `ExtraResource`s. The resolved