edge/pkg/module/context.go

package module

import (
	"context"

	"forge.cadoles.com/arcad/edge/pkg/app"
	"forge.cadoles.com/arcad/edge/pkg/module/util"
	"github.com/dop251/goja"
	"github.com/pkg/errors"
)

type ContextKey string

type ContextModule struct{}

func (m *ContextModule) Name() string {
	return "context"
}

func (m *ContextModule) new(call goja.FunctionCall, rt *goja.Runtime) goja.Value {
	return rt.ToValue(context.Background())
}

func (m *ContextModule) with(call goja.FunctionCall, rt *goja.Runtime) goja.Value {
	ctx := util.AssertContext(call.Argument(0), rt)
	rawValues := util.AssertObject(call.Argument(1), rt)

	values := make(map[ContextKey]any)
	for k, v := range rawValues {
		values[ContextKey(k)] = v
	}

	ctx = WithContext(ctx, values)

	return rt.ToValue(ctx)
}

func (m *ContextModule) get(call goja.FunctionCall, rt *goja.Runtime) goja.Value {
	ctx := util.AssertContext(call.Argument(0), rt)
	rawKey := util.AssertString(call.Argument(1), rt)

	value := ctx.Value(ContextKey(rawKey))

	return rt.ToValue(value)
}

func (m *ContextModule) Export(export *goja.Object) {
	if err := export.Set("new", m.new); err != nil {
		panic(errors.Wrap(err, "could not set 'new' function"))
	}

	if err := export.Set("get", m.get); err != nil {
		panic(errors.Wrap(err, "could not set 'get' function"))
	}

	if err := export.Set("with", m.with); err != nil {
		panic(errors.Wrap(err, "could not set 'with' function"))
	}
}

func ContextModuleFactory() app.ServerModuleFactory {
	return func(server *app.Server) app.ServerModule {
		return &ContextModule{}
	}
}

func ContextValue[T any](ctx context.Context, key ContextKey) T {
	value, ok := ctx.Value(key).(T)
	if !ok {
		return *new(T)
	}

	return value
}

func WithContext(ctx context.Context, values map[ContextKey]any) context.Context {
	for k, v := range values {
		ctx = context.WithValue(ctx, k, v)
	}

	return ctx
}