support customized formatter

README.md

@@ -5,23 +5,10 @@ Shipping log entries from [logrus](https://github.com/sirupsen/logrus) to Datado
 ## Example
 
 ```golang
-    // Sending log in JSON format
     hostName, _ := os.Hostname()
-    // When failure, retry up to 3 times with 5s interval
-    hook := datadog.NewHook(datadog.DatadogUSHost, apiKey, true, 3, 5*time.Second)
-    hook.Hostname = hostName
+    // Sending log in JSON, batch log every 5 sec and when failure, retry up to 3 times
+    hook := NewHook(host, apiKey, 5*time.Second, 3, logrus.TraceLevel, &logrus.JSONFormatter{}, Options{Hostname: hostName})
     l := logrus.New()
     l.Hooks.Add(hook)
     l.WithField("from", "unitest").Infof("TestSendingJSON - %d", i)
 ```
-
-```golang
-    // Sending log in plain text
-    hostName, _ := os.Hostname()
-    // When failure, retry up to 3 times with 5s interval
-    hook := datadog.NewHook(datadog.DatadogUSHost, apiKey, false, 3, 5*time.Second)
-    hook.Hostname = hostName
-    l := logrus.New()
-    l.Hooks.Add(hook)
-    l.WithField("from", "unitest").Infof("TestSendingText - %d", i)
-```

go.mod

@@ -2,4 +2,7 @@ module github.com/bin3377/logrus-datadog-hook
 
 go 1.13
 
-require github.com/sirupsen/logrus v1.4.2
+require (
+	github.com/sirupsen/logrus v1.4.2
+	gopkg.in/natefinch/lumberjack.v2 v2.0.0
+)

go.sum

@@ -7,3 +7,5 @@ github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 golang.org/x/sys v0.0.0-20190422165155-953cdadca894 h1:Cz4ceDQGXuKRnVBDTS23GTn/pU5OE2C0WrNTOYK1Uuc=
 golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+gopkg.in/natefinch/lumberjack.v2 v2.0.0 h1:1Lc07Kr7qY4U2YPouBjpCLxpiyxIVoxqXgkXLknAOE8=
+gopkg.in/natefinch/lumberjack.v2 v2.0.0/go.mod h1:l0ndWWf7gzL7RNwBG7wST/UCcT4T24xpD6X8LsfU/+k=

hook.go

@@ -12,21 +12,27 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-// Hook is the struct holding connect information to Datadog backend
-type Hook struct {
+// Options define the options for Datadog log stream
+type Options struct {
 	Source   string
 	Service  string
 	Hostname string
 	Tags     []string
+}
 
-	host     string
-	apiKey   string
-	isJSON   bool
-	maxRetry int
-	buffer   [][]byte
-	m        sync.Mutex
-	ch       chan []byte
-	err      error
+// Hook is the struct holding connect information to Datadog backend
+type Hook struct {
+	host      string
+	apiKey    string
+	maxRetry  int
+	formatter logrus.Formatter
+	minLevel  logrus.Level
+	options   Options
+
+	ch     chan []byte
+	buffer [][]byte
+	m      sync.Mutex
+	err    error
 }
 
 const (
@@ -64,19 +70,24 @@ var (
 func NewHook(
 	host string,
 	apiKey string,
-	isJSON bool,
-	maxRetry int,
 	batchTimeout time.Duration,
+	maxRetry int,
+	minLevel logrus.Level,
+	formatter logrus.Formatter,
+	options Options,
 ) *Hook {
+
 	h := &Hook{
-		host:     host,
-		apiKey:   apiKey,
-		isJSON:   isJSON,
-		maxRetry: maxRetry,
+		host:      host,
+		apiKey:    apiKey,
+		maxRetry:  maxRetry,
+		minLevel:  minLevel,
+		formatter: formatter,
+		options:   options,
 	}
 
-	if batchTimeout <= 0 {
-		batchTimeout = defaultTimeout
+	if batchTimeout < 5*time.Second {
+		batchTimeout = 5 * time.Second
 	}
 	h.ch = make(chan []byte, 1)
 	go h.pile(time.Tick(batchTimeout))
@@ -85,25 +96,12 @@ func NewHook(
 
 // Levels - implement Hook interface supporting all levels
 func (h *Hook) Levels() []logrus.Level {
-	return []logrus.Level{
-		logrus.PanicLevel,
-		logrus.FatalLevel,
-		logrus.ErrorLevel,
-		logrus.WarnLevel,
-		logrus.InfoLevel,
-		logrus.DebugLevel,
-	}
+	return logrus.AllLevels[:h.minLevel+1]
 }
 
 // Fire - implement Hook interface fire the entry
 func (h *Hook) Fire(entry *logrus.Entry) error {
-	var fn func(*logrus.Entry) ([]byte, error)
-	if h.isJSON {
-		fn = (&logrus.JSONFormatter{}).Format
-	} else {
-		fn = (&logrus.TextFormatter{DisableColors: true}).Format
-	}
-	line, err := fn(entry)
+	line, err := h.formatter.Format(entry)
 	if err != nil {
 		dbg("Unable to read entry, %v", err)
 		return err
@@ -122,7 +120,7 @@ func (h *Hook) pile(ticker <-chan time.Time) {
 			if str == "" {
 				continue
 			}
-			if h.isJSON {
+			if h.isJSON() {
 				str = strings.TrimRight(str, "\n")
 				str += ","
 			} else if !strings.HasSuffix(str, "\n") {
@@ -145,6 +143,20 @@ func (h *Hook) pile(ticker <-chan time.Time) {
 	}
 }
 
+func (h *Hook) isJSON() bool {
+	if _, ok := h.formatter.(*logrus.JSONFormatter); ok {
+		return true
+	} else if _, ok := h.formatter.(*logrus.TextFormatter); ok {
+		return false
+	}
+	b, err := h.formatter.Format(&logrus.Entry{})
+	if err != nil {
+		return false
+	}
+	str := strings.TrimSpace(string(b))
+	return strings.HasPrefix(str, "{") && strings.HasSuffix(str, "}")
+}
+
 func (h *Hook) send(pile [][]byte) {
 	h.m.Lock()
 	defer h.m.Unlock()
@@ -159,7 +171,7 @@ func (h *Hook) send(pile [][]byte) {
 	if len(buf) == 0 {
 		return
 	}
-	if h.isJSON {
+	if h.isJSON() {
 		if buf[len(buf)-1] == ',' {
 			buf = buf[:len(buf)-1]
 		}
@@ -176,21 +188,28 @@ func (h *Hook) send(pile [][]byte) {
 	}
 	header := http.Header{}
 	header.Add(apiKeyHeader, h.apiKey)
-	if h.isJSON {
+	if h.isJSON() {
 		header.Add("Content-Type", contentTypeJSON)
 	} else {
 		header.Add("Content-Type", contentTypePlain)
 	}
 	header.Add("charset", "UTF-8")
 	req.Header = header
 
-	for i := 0; i <= h.maxRetry; i++ {
+	i := 0
+	for {
 		resp, err := http.DefaultClient.Do(req)
-		if err == nil {
-			dbg("%v", resp)
+		if err == nil && resp.StatusCode < 400 {
+			dbg("Success - %d", resp.StatusCode)
+			return
+		}
+		dbg("err  = %v", err)
+		dbg("resp = %v", resp)
+		i++
+		if h.maxRetry < 0 || i >= h.maxRetry {
+			dbg("Still failed after %d retries", i)
 			return
 		}
-		dbg(err.Error())
 	}
 }
 
@@ -202,17 +221,18 @@ func (h *Hook) datadogURL() string {
 	}
 	u.Path += basePath
 	parameters := url.Values{}
-	if h.Source != "" {
-		parameters.Add("ddsource", h.Source)
+	o := h.options
+	if o.Source != "" {
+		parameters.Add("ddsource", o.Source)
 	}
-	if h.Service != "" {
-		parameters.Add("service", h.Service)
+	if o.Service != "" {
+		parameters.Add("service", o.Service)
 	}
-	if h.Hostname != "" {
-		parameters.Add("hostname", h.Hostname)
+	if o.Hostname != "" {
+		parameters.Add("hostname", o.Hostname)
 	}
-	if h.Tags != nil {
-		tags := strings.Join(h.Tags, ",")
+	if o.Tags != nil {
+		tags := strings.Join(o.Tags, ",")
 		parameters.Add("ddtags", tags)
 	}
 	u.RawQuery = parameters.Encode()

hook_test.go

@@ -1,6 +1,7 @@
 package datadog
 
 import (
+	"fmt"
 	"log"
 	"os"
 	"path/filepath"
@@ -40,7 +41,7 @@ func equals(tb testing.TB, exp, act interface{}) {
 	}
 }
 
-func getLogger(t *testing.T, isJSON bool) (*Hook, *logrus.Logger) {
+func getLogger(t *testing.T, formatter logrus.Formatter) (*Hook, *logrus.Logger) {
 	host := os.Getenv("DATADOG_HOST")
 	apiKey := os.Getenv("DATADOG_APIKEY")
 	Debug = true
@@ -53,47 +54,31 @@ func getLogger(t *testing.T, isJSON bool) (*Hook, *logrus.Logger) {
 	}
 
 	hostName, _ := os.Hostname()
-	hook := NewHook(host, apiKey, isJSON, 3, 5*time.Second)
-	hook.Hostname = hostName
+	hook := NewHook(host, apiKey, 1*time.Second, 3, logrus.TraceLevel, formatter, Options{Hostname: hostName})
 	l := logrus.New()
+	l.Level = logrus.TraceLevel
 	l.Hooks.Add(hook)
 	return hook, l
 }
 
 func TestHook(t *testing.T) {
-	hook, l := getLogger(t, true)
-
+	hook, l := getLogger(t, &logrus.JSONFormatter{})
 	for _, level := range hook.Levels() {
 		if len(l.Hooks[level]) != 1 {
 			t.Errorf("Hook was not added. The length of l.Hooks[%v]: %v", level, len(l.Hooks[level]))
 		}
 	}
 }
-func TestSendingJSON(t *testing.T) {
-	_, l := getLogger(t, true)
-
-	var wg sync.WaitGroup
-	for i := 0; i < 20; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			l.WithField("from", "unitest").Infof("TestSendingJSON - %d", i)
-		}()
-		time.Sleep(1 * time.Second)
-	}
-
-	wg.Wait()
-}
-
-func TestSendingPlain(t *testing.T) {
-	_, l := getLogger(t, false)
+func TestSending(t *testing.T) {
+	_, l := getLogger(t, &logrus.JSONFormatter{TimestampFormat: time.RFC3339})
 
 	var wg sync.WaitGroup
 	for i := 0; i < 20; i++ {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			l.WithField("from", "unitest").Infof("TestSendingPlain - %d", i)
+			w := l.WithField("from", "unitest").WriterLevel(logrus.Level(i%5 + 2))
+			fmt.Fprintf(w, "TestSending - %d\n", i)
 		}()
 		time.Sleep(1 * time.Second)
 	}