From fcf6d6a02b201dd2660b9b97036f7d0ec8af1124 Mon Sep 17 00:00:00 2001
From: wineandchord <guoqizhou19@gmail.com>
Date: Thu, 14 Dec 2023 17:12:41 +0800
Subject: [PATCH] log: provide registration for custom format encoder (#146)

Fixes #145
---
 log/zaplogger.go      | 26 ++++++++++++++++-------
 log/zaplogger_test.go | 49 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/log/zaplogger.go b/log/zaplogger.go
index 63b48af..974618f 100644
--- a/log/zaplogger.go
+++ b/log/zaplogger.go
@@ -119,14 +119,26 @@ func newEncoder(c *OutputConfig) zapcore.Encoder {
 	if c.EnableColor {
 		encoderCfg.EncodeLevel = zapcore.CapitalColorLevelEncoder
 	}
-	switch c.Formatter {
-	case "console":
-		return zapcore.NewConsoleEncoder(encoderCfg)
-	case "json":
-		return zapcore.NewJSONEncoder(encoderCfg)
-	default:
-		return zapcore.NewConsoleEncoder(encoderCfg)
+	if newFormatEncoder, ok := formatEncoders[c.Formatter]; ok {
+		return newFormatEncoder(encoderCfg)
 	}
+	// Defaults to console encoder.
+	return zapcore.NewConsoleEncoder(encoderCfg)
+}
+
+var formatEncoders = map[string]NewFormatEncoder{
+	"console": zapcore.NewConsoleEncoder,
+	"json":    zapcore.NewJSONEncoder,
+}
+
+// NewFormatEncoder is the function type for creating a format encoder out of an encoder config.
+type NewFormatEncoder func(zapcore.EncoderConfig) zapcore.Encoder
+
+// RegisterFormatEncoder registers a NewFormatEncoder with the specified formatName key.
+// The existing formats include "console" and "json", but you can override these format encoders
+// or provide a new custom one.
+func RegisterFormatEncoder(formatName string, newFormatEncoder NewFormatEncoder) {
+	formatEncoders[formatName] = newFormatEncoder
 }
 
 // GetLogEncoderKey gets user defined log output name, uses defKey if empty.
diff --git a/log/zaplogger_test.go b/log/zaplogger_test.go
index 58eca32..5846fec 100644
--- a/log/zaplogger_test.go
+++ b/log/zaplogger_test.go
@@ -17,12 +17,14 @@ import (
 	"errors"
 	"fmt"
 	"runtime"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
+	"go.uber.org/zap/buffer"
 	"go.uber.org/zap/zapcore"
 	"go.uber.org/zap/zaptest/observer"
 
@@ -334,3 +336,50 @@ func TestLogEnableColor(t *testing.T) {
 	l.Warn("hello")
 	l.Error("hello")
 }
+
+func TestLogNewFormatEncoder(t *testing.T) {
+	const myFormatter = "myformatter"
+	log.RegisterFormatEncoder(myFormatter, func(ec zapcore.EncoderConfig) zapcore.Encoder {
+		return &consoleEncoder{
+			Encoder: zapcore.NewJSONEncoder(zapcore.EncoderConfig{}),
+			pool:    buffer.NewPool(),
+			cfg:     ec,
+		}
+	})
+	cfg := []log.OutputConfig{{Writer: "console", Level: "trace", Formatter: myFormatter}}
+	l := log.NewZapLog(cfg).With(log.Field{Key: "trace-id", Value: "xx"})
+	l.Trace("hello")
+	l.Debug("hello")
+	l.Info("hello")
+	l.Warn("hello")
+	l.Error("hello")
+	// 2023/12/14 10:54:55 {"trace-id":"xx"} DEBUG hello
+	// 2023/12/14 10:54:55 {"trace-id":"xx"} DEBUG hello
+	// 2023/12/14 10:54:55 {"trace-id":"xx"} INFO hello
+	// 2023/12/14 10:54:55 {"trace-id":"xx"} WARN hello
+	// 2023/12/14 10:54:55 {"trace-id":"xx"} ERROR hello
+}
+
+type consoleEncoder struct {
+	zapcore.Encoder
+	pool buffer.Pool
+	cfg  zapcore.EncoderConfig
+}
+
+func (c consoleEncoder) Clone() zapcore.Encoder {
+	return consoleEncoder{Encoder: c.Encoder.Clone(), pool: buffer.NewPool(), cfg: c.cfg}
+}
+
+func (c consoleEncoder) EncodeEntry(entry zapcore.Entry, fields []zapcore.Field) (*buffer.Buffer, error) {
+	buf, err := c.Encoder.EncodeEntry(zapcore.Entry{}, nil)
+	if err != nil {
+		return nil, err
+	}
+	buffer := c.pool.Get()
+	buffer.AppendString(entry.Time.Format("2006/01/02 15:04:05"))
+	field := buf.String()
+	buffer.AppendString(" " + field[:len(field)-1] + " ")
+	buffer.AppendString(strings.ToUpper(entry.Level.String()) + " ")
+	buffer.AppendString(entry.Message + "\n")
+	return buffer, nil
+}