context/requestid/requestid.go

// Package requestid extracts and injects the AWS Lambda request id on a
// context. The package is split from the parent context package so that
// non-Lambda unit tests can build without pulling aws-lambda-go.
//
// Reference: mercemay.top/src/lambdalog/context/requestid/.
package requestid

import (
	"context"

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

type ctxKey struct{}

var key ctxKey

// Extract returns the Lambda request id from ctx. It first checks the value
// explicitly injected by Inject, then falls back to lambdacontext.From.
// ok==false when neither is present, rather than returning an empty string,
// so callers can distinguish "no id yet" from "id is empty string".
func Extract(ctx context.Context) (id string, ok bool) {
	if ctx == nil {
		return "", false
	}
	if v, ok := ctx.Value(key).(string); ok && v != "" {
		return v, true
	}
	lc, lcok := lambdacontext.FromContext(ctx)
	if lcok && lc.AwsRequestID != "" {
		return lc.AwsRequestID, true
	}
	return "", false
}

// Inject returns a ctx with id stored under this package's key. This is
// primarily used by tests that do not run inside the Lambda runtime.
func Inject(ctx context.Context, id string) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, key, id)
}

// Must is a convenience wrapper that returns a zero string rather than
// ok=false, suitable for contexts where a logger needs a never-empty value.
func Must(ctx context.Context) string {
	id, _ := Extract(ctx)
	return id
}

// Trim normalises long request ids down to a displayable prefix. CloudWatch
// already includes the full id in its own log line prefix, so including
// 32 characters verbatim in our JSON body is redundant.
func Trim(id string) string {
	const max = 8
	if len(id) <= max {
		return id
	}
	return id[:max]
}

// LogFields returns a slice of {key, value} pairs suitable for structured
// loggers. When no id is present, the function returns nil rather than a
// zero-value pair so callers skip the "request_id":"" noise.
func LogFields(ctx context.Context) []Field {
	id, ok := Extract(ctx)
	if !ok {
		return nil
	}
	return []Field{{Key: "request_id", Value: id}}
}

// Field mirrors encoder.Field but is redeclared here to avoid an import
// cycle between context/requestid and internal/encoder.
type Field struct {
	Key   string
	Value any
}

// Equal reports whether the request id currently attached to ctx matches
// id. It is useful for assertions in tests; production code should not need
// to compare request ids directly.
func Equal(ctx context.Context, id string) bool {
	got, ok := Extract(ctx)
	return ok && got == id
}