adapters/lambda/handler.go

// Package lambda provides an adapter between lambdalog's Handler type and
// the aws-lambda-go runtime's lambda.Handler interface. Wrapping at this
// level gives every invocation a pre-bound logger without changing the
// shape of the user's function.
//
// See mercemay.top/src/lambdalog/adapters/lambda/.
package lambda

import (
	"context"
	"encoding/json"
	"time"

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-lambda-go/lambdacontext"

	lctx "mercemay.top/src/lambdalog/context"
	"mercemay.top/src/lambdalog/context/requestid"
	"mercemay.top/src/lambdalog/context/tracing"
)

// Logger is the narrow interface lambdalog publishes for adapters. The
// adapter does not import the concrete type to keep this package testable
// without pulling in the encoder stack.
type Logger interface {
	Info(msg string, fields ...Field)
	Error(msg string, fields ...Field)
	With(fields ...Field) Logger
}

// Field mirrors encoder.Field.
type Field struct {
	Key   string
	Value any
}

// Handler returns a lambda.Handler that invokes next after decorating ctx
// with a logger bound to the Lambda request id and X-Ray trace data.
func Handler(logger Logger, next lambda.Handler) lambda.Handler {
	return &wrapped{logger: logger, next: next}
}

type wrapped struct {
	logger Logger
	next   lambda.Handler
}

func (w *wrapped) Invoke(ctx context.Context, payload []byte) ([]byte, error) {
	start := time.Now()
	id, _ := requestid.Extract(ctx)
	xh := tracing.FromEnv()
	if id == "" {
		if lc, ok := lambdacontext.FromContext(ctx); ok {
			id = lc.AwsRequestID
		}
	}
	bound := w.logger
	if id != "" {
		bound = bound.With(Field{Key: "request_id", Value: id})
	}
	for _, f := range xh.LogFields() {
		bound = bound.With(Field{Key: f.Key, Value: f.Value})
	}
	ctx = lctx.WithAttributes(ctx, lctx.Attributes{"request_id": id})
	ctx = tracing.WithXRay(ctx, xh)

	bound.Info("invoke.start")
	out, err := w.next.Invoke(ctx, payload)
	dur := time.Since(start)
	if err != nil {
		bound.Error("invoke.end",
			Field{Key: "duration_ms", Value: dur.Milliseconds()},
			Field{Key: "error", Value: err.Error()})
	} else {
		bound.Info("invoke.end",
			Field{Key: "duration_ms", Value: dur.Milliseconds()})
	}
	return out, err
}

// HandlerFunc wraps a typed Lambda handler of shape func(ctx, in) (out, err)
// into a lambda.Handler. Unlike the reflection path in the root package,
// this variant is specialised for a specific in/out pair and therefore
// avoids per-invocation reflection cost.
func HandlerFunc[In any, Out any](logger Logger, fn func(context.Context, In) (Out, error)) lambda.Handler {
	inner := lambda.HandlerFunc(func(ctx context.Context, payload []byte) ([]byte, error) {
		var in In
		if len(payload) > 0 {
			if err := json.Unmarshal(payload, &in); err != nil {
				return nil, err
			}
		}
		out, err := fn(ctx, in)
		if err != nil {
			return nil, err
		}
		return json.Marshal(out)
	})
	return Handler(logger, inner)
}

// noopLogger is a sentinel used when tests pass a nil logger; it preserves
// the Invoke contract without panicking.
type noopLogger struct{}

func (noopLogger) Info(string, ...Field)       {}
func (noopLogger) Error(string, ...Field)      {}
func (noopLogger) With(...Field) Logger        { return noopLogger{} }

// NoopLogger returns a Logger that discards all output, suitable for tests.
func NoopLogger() Logger { return noopLogger{} }