resolver/wasm: Add NewWithContext to allow passing context (#8499)

Previously, initializing a new WASM resolver always used a background
context. This prevented callers from passing down an existing context
for timeouts, cancellation, or tracing.

This change introduces `NewWithContext` in `v1/resolver/wasm` which accepts
a context and propagates it to `Entrypoints()`. The existing `New`
function has been updated to wrap `NewWithContext` using a background
context to preserve backwards compatibility. `LoadWasmResolversFromStore`
has been updated to pass the provided context appropriately.

Signed-off-by: Dominik Schulz <dschulz@google.com>

internal/bundle/utils.go

@@ -78,7 +78,7 @@ func LoadWasmResolversFromStore(ctx context.Context, store storage.Store, txn st
 		}
 
 		for _, wmf := range resolversToLoad {
-			resolver, err := wasm.New(wmf.Entrypoints, wmf.Raw, data)
+			resolver, err := wasm.NewWithContext(ctx, wmf.Entrypoints, wmf.Raw, data)
 			if err != nil {
 				return nil, fmt.Errorf("failed to initialize wasm module for entrypoints '%s': %s", wmf.Entrypoints, err)
 			}

resolver/wasm/wasm.go

@@ -5,6 +5,8 @@
 package wasm
 
 import (
+	"context"
+
 	"github.com/open-policy-agent/opa/ast"
 	v1 "github.com/open-policy-agent/opa/v1/resolver/wasm"
 )
@@ -12,7 +14,7 @@ import (
 // New creates a new Resolver instance which is using the Wasm module
 // policy for the given entrypoint ref.
 func New(entrypoints []ast.Ref, policy []byte, data any) (*Resolver, error) {
-	return v1.New(entrypoints, policy, data)
+	return v1.NewWithContext(context.TODO(), entrypoints, policy, data)
 }
 
 // Resolver implements the resolver.Resolver interface

v1/resolver/wasm/wasm.go

@@ -16,8 +16,16 @@ import (
 )
 
 // New creates a new Resolver instance which is using the Wasm module
-// policy for the given entrypoint ref.
+// policy for the given entrypoint ref. This method creates a new
+// background context. If you need to pass an existing context use
+// NewWithContext instead.
 func New(entrypoints []ast.Ref, policy []byte, data any) (*Resolver, error) {
+	return NewWithContext(context.Background(), entrypoints, policy, data)
+}
+
+// NewWithContext creates a new Resolver instance which is using the Wasm module
+// policy for the given entrypoint ref. This method accepts a context.
+func NewWithContext(ctx context.Context, entrypoints []ast.Ref, policy []byte, data any) (*Resolver, error) {
 	e, err := opa.LookupEngine("wasm")
 	if err != nil {
 		return nil, err
@@ -37,7 +45,7 @@ func New(entrypoints []ast.Ref, policy []byte, data any) (*Resolver, error) {
 	// only the configured ones will be used when Eval() is
 	// called.
 	entrypointRefToID := ast.NewValueMap()
-	epIDs, err := o.Entrypoints(context.Background())
+	epIDs, err := o.Entrypoints(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -137,7 +145,6 @@ func (r *Resolver) RemoveDataPath(ctx context.Context, path []string) error {
 }
 
 func getResult(evalResult *opa.Result) (ast.Value, error) {
-
 	parsed, err := ast.ParseTerm(string(evalResult.Result))
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse wasm result: %s", err)