From c5970766a35cbc588c3b59b21d626e296894ffe1 Mon Sep 17 00:00:00 2001
From: IamCathal <cathalsocallaghan@gmail.com>
Date: Mon, 24 Aug 2020 20:41:39 +0100
Subject: [PATCH] Add init to es.go

---
 pkg/logs/es/es.go | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/pkg/logs/es/es.go b/pkg/logs/es/es.go
index 06dfece1..4dfc4160 100644
--- a/pkg/logs/es/es.go
+++ b/pkg/logs/es/es.go
@@ -31,8 +31,10 @@ func NewES() logs.Logger {
 // import _ "github.com/astaxie/beego/logs/es"
 type esLogger struct {
 	*elasticsearch.Client
-	DSN   string `json:"dsn"`
-	Level int    `json:"level"`
+	DSN                string `json:"dsn"`
+	Level              int    `json:"level"`
+	UseCustomFormatter bool
+	CustomFormatter    func(*logs.LogMsg) string
 }
 
 func (el *esLogger) Format(lm *logs.LogMsg) string {
@@ -40,7 +42,14 @@ func (el *esLogger) Format(lm *logs.LogMsg) string {
 }
 
 // {"dsn":"http://localhost:9200/","level":1}
-func (el *esLogger) Init(jsonconfig string) error {
+func (el *esLogger) Init(jsonconfig string, LogFormatter ...func(*logs.LogMsg) string) error {
+	for _, elem := range LogFormatter {
+		if elem != nil {
+			el.UseCustomFormatter = true
+			el.CustomFormatter = elem
+		}
+	}
+
 	err := json.Unmarshal([]byte(jsonconfig), el)
 	if err != nil {
 		return err
@@ -69,9 +78,16 @@ func (el *esLogger) WriteMsg(lm *logs.LogMsg) error {
 		return nil
 	}
 
+	msg := ""
+	if el.UseCustomFormatter {
+		msg = el.CustomFormatter(lm)
+	} else {
+		msg = el.Format(lm)
+	}
+
 	idx := LogDocument{
 		Timestamp: lm.When.Format(time.RFC3339),
-		Msg:       el.Format(lm),
+		Msg:       msg,
 	}
 
 	body, err := json.Marshal(idx)