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

     1  package ctxzap
     2  
     3  import (
     4  	grpc_ctxtags "github.com/hxx258456/ccgo/go-grpc-middleware/tags"
     5  	"github.com/hxx258456/ccgo/net/context"
     6  	"go.uber.org/zap"
     7  	"go.uber.org/zap/zapcore"
     8  )
     9  
    10  type ctxMarker struct{}
    11  
    12  type ctxLogger struct {
    13  	logger *zap.Logger
    14  	fields []zapcore.Field
    15  }
    16  
    17  var (
    18  	ctxMarkerKey = &ctxMarker{}
    19  	nullLogger   = zap.NewNop()
    20  )
    21  
    22  // AddFields adds zap fields to the logger.
    23  func AddFields(ctx context.Context, fields ...zapcore.Field) {
    24  	l, ok := ctx.Value(ctxMarkerKey).(*ctxLogger)
    25  	if !ok || l == nil {
    26  		return
    27  	}
    28  	l.fields = append(l.fields, fields...)
    29  
    30  }
    31  
    32  // Extract takes the call-scoped Logger from grpc_zap middleware.
    33  //
    34  // It always returns a Logger that has all the grpc_ctxtags updated.
    35  func Extract(ctx context.Context) *zap.Logger {
    36  	l, ok := ctx.Value(ctxMarkerKey).(*ctxLogger)
    37  	if !ok || l == nil {
    38  		return nullLogger
    39  	}
    40  	// Add grpc_ctxtags tags metadata until now.
    41  	fields := TagsToFields(ctx)
    42  	// Add zap fields added until now.
    43  	fields = append(fields, l.fields...)
    44  	return l.logger.With(fields...)
    45  }
    46  
    47  // TagsToFields transforms the Tags on the supplied context into zap fields.
    48  func TagsToFields(ctx context.Context) []zapcore.Field {
    49  	fields := []zapcore.Field{}
    50  	tags := grpc_ctxtags.Extract(ctx)
    51  	for k, v := range tags.Values() {
    52  		fields = append(fields, zap.Any(k, v))
    53  	}
    54  	return fields
    55  }
    56  
    57  // ToContext adds the zap.Logger to the context for extraction later.
    58  // Returning the new context that has been created.
    59  func ToContext(ctx context.Context, logger *zap.Logger) context.Context {
    60  	l := &ctxLogger{
    61  		logger: logger,
    62  	}
    63  	return context.WithValue(ctx, ctxMarkerKey, l)
    64  }