From 637ac3856b43e1aa812ee7a7fa544ccc4a47f51f Mon Sep 17 00:00:00 2001
From: Kuwon Sebastian Na <laggu91@gmail.com>
Date: Fri, 29 Nov 2024 09:33:50 +0900
Subject: [PATCH] feat: logging data from gin.Context (#99)

---
 _example/main.go | 12 ++++++++++++
 logger.go        |  7 +++++++
 options.go       |  6 ++++++
 3 files changed, 25 insertions(+)

diff --git a/_example/main.go b/_example/main.go
index 089e27d..b9f150f 100644
--- a/_example/main.go
+++ b/_example/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"math/rand"
 	"net/http"
 	"regexp"
 	"time"
@@ -117,6 +118,17 @@ func main() {
 		c.String(http.StatusOK, "pong "+fmt.Sprint(time.Now().Unix()))
 	})
 
+	// Example of logging data on gin.Context
+	r.GET("/context", logger.SetLogger(
+		logger.WithContext(func(c *gin.Context, e *zerolog.Event) *zerolog.Event {
+			return e.Any("data1", c.MustGet("data1")).Any("data2", c.MustGet("data2"))
+		}),
+	), func(c *gin.Context) {
+		c.Set("data1", rand.Intn(100))
+		c.Set("data2", rand.Intn(100))
+		c.String(http.StatusOK, "pong "+fmt.Sprint(time.Now().Unix()))
+	})
+
 	// Example of skipper usage
 	v1 := r.Group("/v1", logger.SetLogger(
 		logger.WithSkipper(func(c *gin.Context) bool {
diff --git a/logger.go b/logger.go
index 87ff4b9..b4fc9d0 100644
--- a/logger.go
+++ b/logger.go
@@ -17,6 +17,8 @@ import (
 // within the context of a Gin HTTP request.
 type Fn func(*gin.Context, zerolog.Logger) zerolog.Logger
 
+type EventFn func(*gin.Context, *zerolog.Event) *zerolog.Event
+
 // Skipper defines a function to skip middleware. It takes a gin.Context as input
 // and returns a boolean indicating whether to skip the middleware for the given context.
 type Skipper func(c *gin.Context) bool
@@ -25,6 +27,8 @@ type Skipper func(c *gin.Context) bool
 type config struct {
 	// logger is a function that defines the logging behavior.
 	logger Fn
+	// context is a function that defines the logging behavior of gin.Context data
+	context EventFn
 	// utc is a boolean stating whether to use UTC time zone or local.
 	utc bool
 	// skipPath is a list of paths to be skipped from logging.
@@ -175,6 +179,9 @@ func SetLogger(opts ...Option) gin.HandlerFunc {
 			default:
 				evt = rl.WithLevel(cfg.defaultLevel).Ctx(c)
 			}
+			if cfg.context != nil {
+				evt = cfg.context(c, evt)
+			}
 			evt.
 				Int("status", c.Writer.Status()).
 				Str("method", c.Request.Method).
diff --git a/options.go b/options.go
index 7aa91ac..eec355f 100644
--- a/options.go
+++ b/options.go
@@ -105,3 +105,9 @@ func WithSkipper(s Skipper) Option {
 		c.skip = s
 	})
 }
+
+func WithContext(fn func(*gin.Context, *zerolog.Event) *zerolog.Event) Option {
+	return optionFunc(func(c *config) {
+		c.context = fn
+	})
+}