From 0fa1fe03109cb5e41aa7a074fe0d837089d65c1e Mon Sep 17 00:00:00 2001
From: fuxiaohei <fuxiaohei@vip.qq.com>
Date: Fri, 29 Jul 2022 15:43:42 +0800
Subject: [PATCH] feat: add logger hook for standalone job logger

---
 pkg/common/logger.go | 21 +++++++++++++++++++++
 pkg/runner/logger.go |  3 +++
 2 files changed, 24 insertions(+)

diff --git a/pkg/common/logger.go b/pkg/common/logger.go
index a9501ce..74fc96d 100644
--- a/pkg/common/logger.go
+++ b/pkg/common/logger.go
@@ -25,3 +25,24 @@ func Logger(ctx context.Context) logrus.FieldLogger {
 func WithLogger(ctx context.Context, logger logrus.FieldLogger) context.Context {
 	return context.WithValue(ctx, loggerContextKeyVal, logger)
 }
+
+type loggerHookKey string
+
+const loggerHookKeyVal = loggerHookKey("logrus.Hook")
+
+// LoggerHook returns the appropriate logger hook for current context
+// the hook affects job logger, not global logger
+func LoggerHook(ctx context.Context) logrus.Hook {
+	val := ctx.Value(loggerHookKeyVal)
+	if val != nil {
+		if hook, ok := val.(logrus.Hook); ok {
+			return hook
+		}
+	}
+	return nil
+}
+
+// WithLoggerHook adds a value to the context for the logger hook
+func WithLoggerHook(ctx context.Context, hook logrus.Hook) context.Context {
+	return context.WithValue(ctx, loggerHookKeyVal, hook)
+}
diff --git a/pkg/runner/logger.go b/pkg/runner/logger.go
index 81cae21..19cdcf9 100644
--- a/pkg/runner/logger.go
+++ b/pkg/runner/logger.go
@@ -79,6 +79,9 @@ func WithJobLogger(ctx context.Context, jobID string, jobName string, config *Co
 	ctx = WithMasks(ctx, masks)
 
 	logger := logrus.New()
+	if hook := common.LoggerHook(ctx); hook != nil {
+		logger.AddHook(hook)
+	}
 	logger.SetFormatter(formatter)
 	logger.SetOutput(os.Stdout)
 	logger.SetLevel(logrus.GetLevel())