logging.go raw

   1  package middleware
   2  
   3  import (
   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.
  10  type 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.
  15  func 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.
  25  func SetLogger(ctx context.Context, logger logging.Logger) context.Context {
  26  	return context.WithValue(ctx, loggerKey{}, logger)
  27  }
  28  
  29  type setLogger struct {
  30  	Logger logging.Logger
  31  }
  32  
  33  // AddSetLoggerMiddleware adds a middleware that will add the provided logger to the middleware context.
  34  func AddSetLoggerMiddleware(stack *Stack, logger logging.Logger) error {
  35  	return stack.Initialize.Add(&setLogger{Logger: logger}, After)
  36  }
  37  
  38  func (a *setLogger) ID() string {
  39  	return "SetLogger"
  40  }
  41  
  42  func (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  }
  47