github.com/GuanceCloud/cliutils@v1.1.21/logger/opt.go (about)

     1  package logger
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  
     7  	"go.uber.org/zap"
     8  	"go.uber.org/zap/zapcore"
     9  )
    10  
    11  const (
    12  	defaultTraceLoggerName = "trace_id"
    13  	defaultSpanLoggerName  = "span_id"
    14  )
    15  
    16  type CtxOption interface {
    17  	apply(config) config
    18  }
    19  
    20  type optionFunc func(config) config
    21  
    22  func (fn optionFunc) apply(cfg config) config {
    23  	return fn(cfg)
    24  }
    25  
    26  type Trace struct {
    27  	TraceID string
    28  	SpanID  string
    29  }
    30  
    31  type ExtractTrace func(ctx context.Context) Trace
    32  
    33  type config struct {
    34  	enabledTrace bool
    35  	traceKey     string
    36  	spanKey      string
    37  	parseTrace   ExtractTrace
    38  	core         zapcore.Core
    39  	opts         []zap.Option
    40  }
    41  
    42  func newDefaultConfig() config {
    43  	return config{
    44  		traceKey: defaultTraceLoggerName,
    45  		spanKey:  defaultSpanLoggerName,
    46  		core:     zapcore.NewCore(zapcore.NewJSONEncoder(zapcore.EncoderConfig{MessageKey: "message"}), os.Stdout, zap.DebugLevel),
    47  	}
    48  }
    49  
    50  func EnableTrace() CtxOption {
    51  	return optionFunc(func(cfg config) config {
    52  		cfg.enabledTrace = true
    53  		return cfg
    54  	})
    55  }
    56  
    57  func WithZapCore(core zapcore.Core) CtxOption {
    58  	return optionFunc(func(cfg config) config {
    59  		cfg.core = core
    60  		return cfg
    61  	})
    62  }
    63  
    64  func WithZapOpts(opts ...zap.Option) CtxOption {
    65  	return optionFunc(func(cfg config) config {
    66  		cfg.opts = opts
    67  		return cfg
    68  	})
    69  }
    70  
    71  func WithTraceKey(tn string, sn string) CtxOption {
    72  	return optionFunc(func(cfg config) config {
    73  		cfg.traceKey = tn
    74  		cfg.spanKey = sn
    75  		return cfg
    76  	})
    77  }
    78  
    79  func WithParseTrace(f ExtractTrace) CtxOption {
    80  	return optionFunc(func(cfg config) config {
    81  		cfg.parseTrace = f
    82  		return cfg
    83  	})
    84  }