Merge remote-tracking branch 'ikedam/feature/pretimestamp'

10gen · Aug 21, 2023 · 22af876 · 22af876
2 parents b81cc57 + 474986a
commit 22af876
Show file tree

Hide file tree

Showing 4 changed files with 172 additions and 8 deletions.
diff --git a/context.go b/context.go
@@ -349,6 +349,16 @@ func (c Context) Timestamp() Context {
 	return c
 }
 
+// Pretimestamp adds the timestamp to the logger like `Timestamp` does.
+// It adds the "time" key before "level" key. Useful when sorting logs as text in time-order.
+//
+// NOTE: It adds timestamp not when sending log, but when starting a new message.
+// NOTE: Applying both `Timestamp()` and `PreTimestamp()`results duplicated "time" keys.
+func (c Context) Pretimestamp() Context {
+	c.l = c.l.Prehook(th)
+	return c
+}
+
 // Time adds the field key with t formated as string using zerolog.TimeFieldFormat.
 func (c Context) Time(key string, t time.Time) Context {
 	c.l.context = enc.AppendTime(enc.AppendKey(c.l.context, key), t, TimeFieldFormat)

diff --git a/hook_test.go b/hook_test.go
@@ -171,6 +171,112 @@ func TestHook(t *testing.T) {
 	}
 }
 
+func TestPrehook(t *testing.T) {
+	tests := []struct {
+		name string
+		want string
+		test func(log Logger)
+	}{
+		{"Message", `{"level_name":"nolevel","message":"test message"}` + "\n", func(log Logger) {
+			log = log.Prehook(levelNameHook)
+			log.Log().Msg("test message")
+		}},
+		{"NoLevel", `{"level_name":"nolevel"}` + "\n", func(log Logger) {
+			log = log.Prehook(levelNameHook)
+			log.Log().Msg("")
+		}},
+		// Note: Prehooks come before "level"
+		{"Print", `{"level_name":"debug","level":"debug"}` + "\n", func(log Logger) {
+			log = log.Prehook(levelNameHook)
+			log.Print("")
+		}},
+		{"Error", `{"level_name":"error","level":"error"}` + "\n", func(log Logger) {
+			log = log.Prehook(levelNameHook)
+			log.Error().Msg("")
+		}},
+		{"Copy/1", `{"copy_has_level":false,"copy_msg":""}` + "\n", func(log Logger) {
+			log = log.Prehook(copyHook)
+			log.Log().Msg("")
+		}},
+		// Note: Prehooks receives message always as empty
+		{"Copy/2", `{"copy_has_level":true,"copy_level":"info","copy_msg":"","level":"info","message":"a message"}` + "\n", func(log Logger) {
+			log = log.Prehook(copyHook)
+			log.Info().Msg("a message")
+		}},
+		{"Multi", `{"level_name":"error","has_level":true,"test":"logged","level":"error"}` + "\n", func(log Logger) {
+			log = log.Prehook(levelNameHook).Prehook(simpleHook)
+			log.Error().Msg("")
+		}},
+		{"Multi/Message", `{"level_name":"error","has_level":true,"test":"logged","level":"error","message":"a message"}` + "\n", func(log Logger) {
+			log = log.Prehook(levelNameHook).Prehook(simpleHook)
+			log.Error().Msg("a message")
+		}},
+		{"Output/single/pre", `{"level_name":"error","level":"error"}` + "\n", func(log Logger) {
+			ignored := &bytes.Buffer{}
+			log = New(ignored).Prehook(levelNameHook).Output(log.w)
+			log.Error().Msg("")
+		}},
+		{"Output/single/post", `{"level_name":"error","level":"error"}` + "\n", func(log Logger) {
+			ignored := &bytes.Buffer{}
+			log = New(ignored).Output(log.w).Prehook(levelNameHook)
+			log.Error().Msg("")
+		}},
+		{"Output/multi/pre", `{"level_name":"error","has_level":true,"test":"logged","level":"error"}` + "\n", func(log Logger) {
+			ignored := &bytes.Buffer{}
+			log = New(ignored).Prehook(levelNameHook).Prehook(simpleHook).Output(log.w)
+			log.Error().Msg("")
+		}},
+		{"Output/multi/post", `{"level_name":"error","has_level":true,"test":"logged","level":"error"}` + "\n", func(log Logger) {
+			ignored := &bytes.Buffer{}
+			log = New(ignored).Output(log.w).Prehook(levelNameHook).Prehook(simpleHook)
+			log.Error().Msg("")
+		}},
+		{"Output/mixed", `{"level_name":"error","has_level":true,"test":"logged","level":"error"}` + "\n", func(log Logger) {
+			ignored := &bytes.Buffer{}
+			log = New(ignored).Prehook(levelNameHook).Output(log.w).Prehook(simpleHook)
+			log.Error().Msg("")
+		}},
+		{"With/single/pre", `{"level_name":"error","level":"error","with":"pre"}` + "\n", func(log Logger) {
+			log = log.Prehook(levelNameHook).With().Str("with", "pre").Logger()
+			log.Error().Msg("")
+		}},
+		{"With/single/post", `{"level_name":"error","level":"error","with":"post"}` + "\n", func(log Logger) {
+			log = log.With().Str("with", "post").Logger().Prehook(levelNameHook)
+			log.Error().Msg("")
+		}},
+		{"With/multi/pre", `{"level_name":"error","has_level":true,"test":"logged","level":"error","with":"pre"}` + "\n", func(log Logger) {
+			log = log.Prehook(levelNameHook).Prehook(simpleHook).With().Str("with", "pre").Logger()
+			log.Error().Msg("")
+		}},
+		{"With/multi/post", `{"level_name":"error","has_level":true,"test":"logged","level":"error","with":"post"}` + "\n", func(log Logger) {
+			log = log.With().Str("with", "post").Logger().Prehook(levelNameHook).Prehook(simpleHook)
+			log.Error().Msg("")
+		}},
+		{"With/mixed", `{"level_name":"error","has_level":true,"test":"logged","level":"error","with":"mixed"}` + "\n", func(log Logger) {
+			log = log.Prehook(levelNameHook).With().Str("with", "mixed").Logger().Prehook(simpleHook)
+			log.Error().Msg("")
+		}},
+		{"Discard", "", func(log Logger) {
+			log = log.Prehook(discardHook)
+			log.Log().Msg("test message")
+		}},
+		{"None", `{"level":"error"}` + "\n", func(log Logger) {
+			log.Error().Msg("")
+		}},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			out := &bytes.Buffer{}
+			log := New(out)
+			tt.test(log)
+			if got, want := decodeIfBinaryToString(out.Bytes()), tt.want; got != want {
+				t.Errorf("invalid log output:\ngot:  %v\nwant: %v", got, want)
+			}
+		})
+	}
+}
+
 func BenchmarkHooks(b *testing.B) {
 	logger := New(ioutil.Discard)
 	b.ResetTimer()

diff --git a/log.go b/log.go
@@ -82,9 +82,17 @@
 //     log.Warn().Msg("")
 //     // Output: {"level":"warn","severity":"warn"}
 //
-// # Caveats
+// Prehooks are also available. It's called just before "level" is added to a log.
+// `msg` always gets empty in `Run()`:
+//
+//     // Use as prehook:
+//     var h SeverityHook
+//     log := zerolog.New(os.Stdout).Prehook(h)
+//     log.Warn().Msg("")
+//     // Output: {"severity":"warn", "level":"warn"}
 //
-// Field duplication:
+//
+// Caveats
 //
 // There is no fields deduplication out-of-the-box.
 // Using the same key multiple times creates new key in final JSON each time.
@@ -228,12 +236,13 @@ func (l Level) MarshalText() ([]byte, error) {
 // serialization to the Writer. If your Writer is not thread safe,
 // you may consider a sync wrapper.
 type Logger struct {
-	w       LevelWriter
-	level   Level
-	sampler Sampler
-	context []byte
-	hooks   []Hook
-	stack   bool
+	w        LevelWriter
+	level    Level
+	sampler  Sampler
+	context  []byte
+	hooks    []Hook
+	prehooks []Hook
+	stack    bool
 	ctx     context.Context
 }
 
@@ -269,6 +278,9 @@ func (l Logger) Output(w io.Writer) Logger {
 	if len(l.hooks) > 0 {
 		l2.hooks = append(l2.hooks, l.hooks...)
 	}
+	if len(l.prehooks) > 0 {
+		l2.prehooks = append(l2.prehooks, l.prehooks...)
+	}
 	if l.context != nil {
 		l2.context = make([]byte, len(l.context), cap(l.context))
 		copy(l2.context, l.context)
@@ -333,6 +345,13 @@ func (l Logger) Hook(h Hook) Logger {
 	return l
 }
 
+// Hook returns a logger with the h Prehook.
+// The `message` parameter for `Run()` gets always an empty string.
+func (l Logger) Prehook(h Hook) Logger {
+	l.prehooks = append(l.prehooks, h)
+	return l
+}
+
 // Trace starts a new message with trace level.
 //
 // You must call Msg on the returned event in order to send the event.
@@ -474,6 +493,9 @@ func (l *Logger) newEvent(level Level, done func(string)) *Event {
 	e.done = done
 	e.ch = l.hooks
 	e.ctx = l.ctx
+	for _, hook := range l.prehooks {
+		hook.Run(e, e.level, "")
+	}
 	if level != NoLevel && LevelFieldName != "" {
 		e.Str(LevelFieldName, LevelFieldMarshalFunc(level))
 	}

diff --git a/log_test.go b/log_test.go
@@ -723,6 +723,32 @@ func TestOutputWithTimestamp(t *testing.T) {
 	}
 }
 
+func TestContextPretimestamp(t *testing.T) {
+	now := time.Date(2001, time.February, 3, 4, 5, 6, 7, time.UTC)
+	TimestampFunc = func() time.Time {
+		return now
+	}
+	defer func() {
+		TimestampFunc = time.Now
+	}()
+	out := &bytes.Buffer{}
+	log := New(out).With().Pretimestamp().Str("foo", "bar").Logger()
+	log.Debug().Msg("hello world")
+
+	// timestamp comes first
+	if got, want := decodeIfBinaryToString(out.Bytes()), `{"time":"2001-02-03T04:05:06Z","level":"debug","foo":"bar","message":"hello world"}`+"\n"; got != want {
+		t.Errorf("invalid log output:\ngot:  %v\nwant: %v", got, want)
+	}
+
+	// of course, timestamp changes
+	now = time.Date(2002, time.March, 4, 5, 6, 7, 8, time.UTC)
+	out.Reset()
+	log.Debug().Msg("hello world")
+	if got, want := decodeIfBinaryToString(out.Bytes()), `{"time":"2002-03-04T05:06:07Z","level":"debug","foo":"bar","message":"hello world"}`+"\n"; got != want {
+		t.Errorf("invalid log output:\ngot:  %v\nwant: %v", got, want)
+	}
+}
+
 type loggableError struct {
 	error
 }