From 2173899f8f5b34385d0713b79ad2c4ff2071a850 Mon Sep 17 00:00:00 2001
From: Ryan Canty <jrcanty@gmail.com>
Date: Mon, 15 Aug 2016 11:04:25 -0700
Subject: [PATCH] Added customizable keys to JSON formatter

---
 json_formatter.go      | 24 +++++++++++++++++++++---
 json_formatter_test.go | 41 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/json_formatter.go b/json_formatter.go
index 2ad6dc5..7364fe8 100644
--- a/json_formatter.go
+++ b/json_formatter.go
@@ -8,6 +8,9 @@ import (
 type JSONFormatter struct {
 	// TimestampFormat sets the format used for marshaling timestamps.
 	TimestampFormat string
+	MessageKey      string
+	LevelKey        string
+	TimeKey         string
 }
 
 func (f *JSONFormatter) Format(entry *Entry) ([]byte, error) {
@@ -29,9 +32,24 @@ func (f *JSONFormatter) Format(entry *Entry) ([]byte, error) {
 		timestampFormat = DefaultTimestampFormat
 	}
 
-	data["time"] = entry.Time.Format(timestampFormat)
-	data["msg"] = entry.Message
-	data["level"] = entry.Level.String()
+	timeKey := f.TimeKey
+	if timeKey == "" {
+		timeKey = "time"
+	}
+
+	messageKey := f.MessageKey
+	if messageKey == "" {
+		messageKey = "msg"
+	}
+
+	levelKey := f.LevelKey
+	if levelKey == "" {
+		levelKey = "level"
+	}
+
+	data[timeKey] = entry.Time.Format(timestampFormat)
+	data[messageKey] = entry.Message
+	data[levelKey] = entry.Level.String()
 
 	serialized, err := json.Marshal(data)
 	if err != nil {
diff --git a/json_formatter_test.go b/json_formatter_test.go
index 1d70873..0fa3466 100644
--- a/json_formatter_test.go
+++ b/json_formatter_test.go
@@ -3,7 +3,7 @@ package logrus
 import (
 	"encoding/json"
 	"errors"
-
+	"strings"
 	"testing"
 )
 
@@ -118,3 +118,42 @@ func TestJSONEntryEndsWithNewline(t *testing.T) {
 		t.Fatal("Expected JSON log entry to end with a newline")
 	}
 }
+
+func TestJSONMessageKey(t *testing.T) {
+	formatter := &JSONFormatter{MessageKey: "message"}
+
+	b, err := formatter.Format(&Entry{Message: "oh hai"})
+	if err != nil {
+		t.Fatal("Unable to format entry: ", err)
+	}
+	s := string(b)
+	if !(strings.Contains(s, "message") && strings.Contains(s, "oh hai")) {
+		t.Fatal("Expected JSON to format message key")
+	}
+}
+
+func TestJSONLevelKey(t *testing.T) {
+	formatter := &JSONFormatter{LevelKey: "somelevel"}
+
+	b, err := formatter.Format(WithField("level", "something"))
+	if err != nil {
+		t.Fatal("Unable to format entry: ", err)
+	}
+	s := string(b)
+	if !strings.Contains(s, "somelevel") {
+		t.Fatal("Expected JSON to format level key")
+	}
+}
+
+func TestJSONTimeKey(t *testing.T) {
+	formatter := &JSONFormatter{TimeKey: "timeywimey"}
+
+	b, err := formatter.Format(WithField("level", "something"))
+	if err != nil {
+		t.Fatal("Unable to format entry: ", err)
+	}
+	s := string(b)
+	if !strings.Contains(s, "timeywimey") {
+		t.Fatal("Expected JSON to format time key")
+	}
+}