1package middleware
 2
 3import (
 4	"context"
 5
 6	"github.com/aws/smithy-go/logging"
 7)
 8
 9// loggerKey is the context value key for which the logger is associated with.
10type loggerKey struct{}
11
12// GetLogger takes a context to retrieve a Logger from. If no logger is present on the context a logging.Nop logger
13// is returned. If the logger retrieved from context supports the ContextLogger interface, the context will be passed
14// to the WithContext method and the resulting logger will be returned. Otherwise the stored logger is returned as is.
15func GetLogger(ctx context.Context) logging.Logger {
16	logger, ok := ctx.Value(loggerKey{}).(logging.Logger)
17	if !ok || logger == nil {
18		return logging.Nop{}
19	}
20
21	return logging.WithContext(ctx, logger)
22}
23
24// SetLogger sets the provided logger value on the provided ctx.
25func SetLogger(ctx context.Context, logger logging.Logger) context.Context {
26	return context.WithValue(ctx, loggerKey{}, logger)
27}
28
29type setLogger struct {
30	Logger logging.Logger
31}
32
33// AddSetLoggerMiddleware adds a middleware that will add the provided logger to the middleware context.
34func AddSetLoggerMiddleware(stack *Stack, logger logging.Logger) error {
35	return stack.Initialize.Add(&setLogger{Logger: logger}, After)
36}
37
38func (a *setLogger) ID() string {
39	return "SetLogger"
40}
41
42func (a *setLogger) HandleInitialize(ctx context.Context, in InitializeInput, next InitializeHandler) (
43	out InitializeOutput, metadata Metadata, err error,
44) {
45	return next.HandleInitialize(SetLogger(ctx, a.Logger), in)
46}