From 01d31ac7486fdb8558ba616525343ee70f7da1d6 Mon Sep 17 00:00:00 2001
From: Jeremie <jeremie.bordier@gmail.com>
Date: Sun, 5 May 2024 09:59:56 +0200
Subject: [PATCH] Add -log-format argument and custom log formatter.

---
 cronexpr/cronexpr/go.mod                |  4 ++-
 log/formatter/custom_field_formatter.go | 48 +++++++++++++++++++++++++
 main.go                                 |  9 ++++-
 3 files changed, 59 insertions(+), 2 deletions(-)
 create mode 100644 log/formatter/custom_field_formatter.go

diff --git a/cronexpr/cronexpr/go.mod b/cronexpr/cronexpr/go.mod
index 1b1121a..610ce8e 100644
--- a/cronexpr/cronexpr/go.mod
+++ b/cronexpr/cronexpr/go.mod
@@ -1,6 +1,8 @@
 module github.com/aptible/supercronic/cronexpr/cronexpr
 
-go 1.18
+go 1.21.4
+
+toolchain go1.22.2
 
 replace github.com/aptible/supercronic => ../../
 
diff --git a/log/formatter/custom_field_formatter.go b/log/formatter/custom_field_formatter.go
new file mode 100644
index 0000000..eff1cbe
--- /dev/null
+++ b/log/formatter/custom_field_formatter.go
@@ -0,0 +1,48 @@
+package formatter
+
+import (
+	"github.com/sirupsen/logrus"
+	"regexp"
+	"strings"
+	"time"
+)
+
+type CustomFieldFormatter struct {
+	LogFormat string
+}
+
+func (f *CustomFieldFormatter) getFieldValue(entry *logrus.Entry, field string) (string, bool) {
+	switch strings.ToLower(field) {
+	case "level":
+		return entry.Level.String(), true
+	case "time":
+		return entry.Time.Format(time.RFC3339Nano), true
+	case "message":
+		return entry.Message, true
+	default:
+		val, ok := entry.Data[field]
+
+		if ok {
+			return val.(string), true
+		}
+
+		return "", false
+	}
+}
+
+func (f *CustomFieldFormatter) Format(entry *logrus.Entry) ([]byte, error) {
+	re := regexp.MustCompile(`%[\w.]+`)
+
+	replaced := re.ReplaceAllStringFunc(f.LogFormat, func(match string) string {
+		// Remove the $ prefix to get the key for valuesMap
+		key := strings.TrimPrefix(match, "%")
+		// If the key exists in the entry, replace with the value from the map
+		if value, ok := f.getFieldValue(entry, key); ok {
+			return value
+		}
+
+		return ""
+	})
+
+	return []byte(strings.TrimSpace(replaced) + "\n"), nil
+}
diff --git a/main.go b/main.go
index 9a09d2d..e2edc25 100644
--- a/main.go
+++ b/main.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/aptible/supercronic/cron"
 	"github.com/aptible/supercronic/crontab"
+	"github.com/aptible/supercronic/log/formatter"
 	"github.com/aptible/supercronic/log/hook"
 	"github.com/aptible/supercronic/prometheus_metrics"
 	"github.com/evalphobia/logrus_sentry"
@@ -37,6 +38,7 @@ func main() {
 			prometheus_metrics.DefaultPort,
 		),
 	)
+	customLogFormat := flag.String("log-format", "", "custom log format. Available fields are %level %time %message %job.command %job.schedule %job.position")
 	splitLogs := flag.Bool("split-logs", false, "split log output into stdout/stderr")
 	passthroughLogs := flag.Bool("passthrough-logs", false, "passthrough logs from commands, do not wrap them in Supercronic logging")
 	sentry := flag.String("sentry-dsn", "", "enable Sentry error logging, using provided DSN")
@@ -80,11 +82,16 @@ func main() {
 		logrus.SetLevel(logrus.WarnLevel)
 	}
 
-	if *json {
+	if *customLogFormat != "" {
+		customFormatter := new(formatter.CustomFieldFormatter)
+		customFormatter.LogFormat = *customLogFormat
+		logrus.SetFormatter(customFormatter)
+	} else if *json {
 		logrus.SetFormatter(&logrus.JSONFormatter{})
 	} else {
 		logrus.SetFormatter(&logrus.TextFormatter{FullTimestamp: true})
 	}
+
 	if *splitLogs {
 		hook.RegisterSplitLogger(
 			logrus.StandardLogger(),