You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
111 lines
2.9 KiB
111 lines
2.9 KiB
package DouTokContext
|
|
|
|
import (
|
|
"context"
|
|
"go.uber.org/zap"
|
|
"go.uber.org/zap/zapcore"
|
|
"reflect"
|
|
)
|
|
|
|
var (
|
|
DefaultNamespace = zap.Namespace("DouTokDefault")
|
|
)
|
|
|
|
var DefaultLogger = zap.NewNop()
|
|
|
|
type loggerKey struct{}
|
|
type namespaceKey struct{}
|
|
type globalFieldKey struct{}
|
|
|
|
// Extract return a logger from the given context
|
|
func Extract(ctx context.Context) *zap.Logger {
|
|
v := ctx.Value(loggerKey{})
|
|
if v == nil {
|
|
logger := DefaultLogger.With(zap.String("logger", "default"))
|
|
return logger
|
|
}
|
|
|
|
logger, ok := v.(*zap.Logger)
|
|
if !ok {
|
|
logger := DefaultLogger.With(zap.String("logger", "default"))
|
|
return logger
|
|
}
|
|
|
|
logger = logger.With(zap.String("logger", "ctx"))
|
|
return WithContextFields(ctx, logger)
|
|
}
|
|
|
|
// AddLoggerToContext adds the zap.Logger to a context for extraction later.
|
|
func AddLoggerToContext(ctx context.Context, logger *zap.Logger) context.Context {
|
|
return context.WithValue(ctx, loggerKey{}, logger)
|
|
}
|
|
|
|
// AddFields will add or override fields to context when zap.Namespace is provided.
|
|
// It will nest any following fields until another zap.Namespace is provided.
|
|
// The difference with pure zap.Namespaces is that when you define a namespace, any following attached
|
|
// logger will be nested, not only when you call logger.With
|
|
func AddFields(ctx context.Context, fields ...zap.Field) context.Context {
|
|
namespaces := extractNamespaces(ctx)
|
|
globalFields := extractGlobalfields(ctx)
|
|
|
|
var namespace zap.Field
|
|
for _, field := range fields {
|
|
if field.Type == zapcore.NamespaceType {
|
|
namespace = field
|
|
continue
|
|
}
|
|
|
|
if !reflect.ValueOf(namespace).IsZero() {
|
|
namespaces[namespace] = append(namespaces[namespace], field)
|
|
} else {
|
|
globalFields = append(globalFields, field)
|
|
}
|
|
}
|
|
|
|
ctx = context.WithValue(ctx, namespaceKey{}, namespaces)
|
|
ctx = context.WithValue(ctx, globalFieldKey{}, globalFields)
|
|
|
|
return ctx
|
|
}
|
|
|
|
// WithContextFields adds context fields to the zap.Logger
|
|
func WithContextFields(ctx context.Context, logger *zap.Logger) *zap.Logger {
|
|
ctx = AddFields(
|
|
ctx, zap.String("RequestId", GetRequestID(ctx)),
|
|
)
|
|
|
|
// TODO: Add more public fields here.
|
|
|
|
logger = logger.With(extractGlobalfields(ctx)...)
|
|
for namespace, fields := range extractNamespaces(ctx) {
|
|
logger = logger.With(Nest(namespace.Key, fields...))
|
|
}
|
|
|
|
return logger
|
|
}
|
|
|
|
func extractNamespaces(ctx context.Context) map[zapcore.Field][]zapcore.Field {
|
|
ctxValue := ctx.Value(namespaceKey{})
|
|
namespaces := make(map[zapcore.Field][]zapcore.Field)
|
|
if ctxValue != nil {
|
|
ctxNamespaces, ok := ctxValue.(map[zapcore.Field][]zapcore.Field)
|
|
if ok {
|
|
for k, v := range ctxNamespaces {
|
|
namespaces[k] = v
|
|
}
|
|
}
|
|
}
|
|
return namespaces
|
|
}
|
|
|
|
func extractGlobalfields(ctx context.Context) []zapcore.Field {
|
|
ctxValue := ctx.Value(globalFieldKey{})
|
|
globalFields := make([]zap.Field, 0)
|
|
if ctxValue != nil {
|
|
value, ok := ctxValue.([]zap.Field)
|
|
if ok {
|
|
globalFields = value
|
|
}
|
|
}
|
|
return globalFields
|
|
}
|
|
|