github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/go-grpc-middleware/logging/logrus/ctxlogrus/context.go (about)

     1  package ctxlogrus
     2  
     3  import (
     4  	grpc_ctxtags "github.com/hxx258456/ccgo/go-grpc-middleware/tags"
     5  	"github.com/hxx258456/ccgo/net/context"
     6  	"github.com/sirupsen/logrus"
     7  )
     8  
     9  type ctxLoggerMarker struct{}
    10  
    11  type ctxLogger struct {
    12  	logger *logrus.Entry
    13  	fields logrus.Fields
    14  }
    15  
    16  var (
    17  	ctxLoggerKey = &ctxLoggerMarker{}
    18  )
    19  
    20  // AddFields adds logrus fields to the logger.
    21  func AddFields(ctx context.Context, fields logrus.Fields) {
    22  	l, ok := ctx.Value(ctxLoggerKey).(*ctxLogger)
    23  	if !ok || l == nil {
    24  		return
    25  	}
    26  	for k, v := range fields {
    27  		l.fields[k] = v
    28  	}
    29  }
    30  
    31  // Extract takes the call-scoped logrus.Entry from ctx_logrus middleware.
    32  //
    33  // If the ctx_logrus middleware wasn't used, a no-op `logrus.Entry` is returned. This makes it safe to
    34  // use regardless.
    35  func Extract(ctx context.Context) *logrus.Entry {
    36  	l, ok := ctx.Value(ctxLoggerKey).(*ctxLogger)
    37  	if !ok || l == nil {
    38  		return logrus.NewEntry(nullLogger)
    39  	}
    40  
    41  	fields := logrus.Fields{}
    42  
    43  	// Add grpc_ctxtags tags metadata until now.
    44  	tags := grpc_ctxtags.Extract(ctx)
    45  	for k, v := range tags.Values() {
    46  		fields[k] = v
    47  	}
    48  
    49  	// Add logrus fields added until now.
    50  	for k, v := range l.fields {
    51  		fields[k] = v
    52  	}
    53  
    54  	return l.logger.WithFields(fields)
    55  }
    56  
    57  // ToContext adds the logrus.Entry to the context for extraction later.
    58  // Returning the new context that has been created.
    59  func ToContext(ctx context.Context, entry *logrus.Entry) context.Context {
    60  	l := &ctxLogger{
    61  		logger: entry,
    62  		fields: logrus.Fields{},
    63  	}
    64  	return context.WithValue(ctx, ctxLoggerKey, l)
    65  }