Add Config (#21)

README.md

@@ -34,7 +34,10 @@ import (
 )
 
 func main() {
-	hook, err := logrus_fluent.New("localhost", 24224)
+	hook, err := logrus_fluent.NewWithConfig(logrus_fluent.Config{
+		Host: "localhost",
+		Port: 24224,
+	})
 	if err != nil {
 		panic(err)
 	}

config.go

@@ -0,0 +1,53 @@
+package logrus_fluent
+
+import (
+	"time"
+
+	"github.com/fluent/fluent-logger-golang/fluent"
+	"github.com/sirupsen/logrus"
+)
+
+// Config is settings for FluentHook.
+type Config struct {
+	Port                  int
+	Host                  string
+	LogLevels             []logrus.Level
+	DisableConnectionPool bool // Fluent client will be created every logging if true.
+	DefaultTag            string
+	DefaultMessageField   string
+	DefaultIgnoreFields   map[string]struct{}
+	DefaultFilters        map[string]func(interface{}) interface{}
+
+	// from fluent.Config
+	// see https://github.com/fluent/fluent-logger-golang/blob/master/fluent/fluent.go
+	FluentNetwork      string
+	FluentSocketPath   string
+	Timeout            time.Duration
+	WriteTimeout       time.Duration
+	BufferLimit        int
+	RetryWait          int
+	MaxRetry           int
+	TagPrefix          string
+	AsyncConnect       bool
+	MarshalAsJSON      bool
+	SubSecondPrecision bool
+}
+
+// FluentConfig converts data to fluent.Config.
+func (c Config) FluentConfig() fluent.Config {
+	return fluent.Config{
+		FluentPort:         c.Port,
+		FluentHost:         c.Host,
+		FluentNetwork:      c.FluentNetwork,
+		FluentSocketPath:   c.FluentSocketPath,
+		Timeout:            c.Timeout,
+		WriteTimeout:       c.WriteTimeout,
+		BufferLimit:        c.BufferLimit,
+		RetryWait:          c.RetryWait,
+		MaxRetry:           c.MaxRetry,
+		TagPrefix:          c.TagPrefix,
+		AsyncConnect:       c.AsyncConnect,
+		MarshalAsJSON:      c.MarshalAsJSON,
+		SubSecondPrecision: c.SubSecondPrecision,
+	}
+}

fluent.go

@@ -36,9 +36,8 @@ type FluentHook struct {
 	// If set, this logger is used for logging.
 	// otherwise new logger is created every time.
 	Fluent *fluent.Fluent
+	conf   Config
 
-	host   string
-	port   int
 	levels []logrus.Level
 	tag    *string
 
@@ -49,32 +48,59 @@ type FluentHook struct {
 
 // New returns initialized logrus hook for fluentd with persistent fluentd logger.
 func New(host string, port int) (*FluentHook, error) {
-	fd, err := fluent.New(fluent.Config{FluentHost: host, FluentPort: port})
-	if err != nil {
-		return nil, err
+	return NewWithConfig(Config{
+		Host:                host,
+		Port:                port,
+		DefaultMessageField: MessageField,
+	})
+}
+
+// NewWithConfig returns initialized logrus hook by config setting.
+func NewWithConfig(conf Config) (*FluentHook, error) {
+	var fd *fluent.Fluent
+	if !conf.DisableConnectionPool {
+		var err error
+		fd, err = fluent.New(conf.FluentConfig())
+		if err != nil {
+			return nil, err
+		}
 	}
 
-	return &FluentHook{
-		levels:       defaultLevels,
-		Fluent:       fd,
-		messageField: MessageField,
-		ignoreFields: make(map[string]struct{}),
-		filters:      make(map[string]func(interface{}) interface{}),
-	}, nil
+	hook := &FluentHook{
+		Fluent: fd,
+		conf:   conf,
+		levels: conf.LogLevels,
+	}
+	// set default values
+	if len(hook.levels) == 0 {
+		hook.levels = defaultLevels
+	}
+	if conf.DefaultTag != "" {
+		tag := conf.DefaultTag
+		hook.tag = &tag
+	}
+	if conf.DefaultMessageField != "" {
+		hook.messageField = conf.DefaultMessageField
+	}
+	if hook.ignoreFields == nil {
+		hook.ignoreFields = make(map[string]struct{})
+	}
+	if hook.filters == nil {
+		hook.filters = make(map[string]func(interface{}) interface{})
+	}
+	return hook, nil
 }
 
 // NewHook returns initialized logrus hook for fluentd.
-// (** deperecated: use New() **)
+// (** deperecated: use New() or NewWithConfig() **)
 func NewHook(host string, port int) *FluentHook {
-	return &FluentHook{
-		host:         host,
-		port:         port,
-		levels:       defaultLevels,
-		tag:          nil,
-		messageField: MessageField,
-		ignoreFields: make(map[string]struct{}),
-		filters:      make(map[string]func(interface{}) interface{}),
-	}
+	hook, _ := NewWithConfig(Config{
+		Host:                  host,
+		Port:                  port,
+		DefaultMessageField:   MessageField,
+		DisableConnectionPool: true,
+	})
+	return hook
 }
 
 // Levels returns logging level to fire this hook.
@@ -125,10 +151,7 @@ func (hook *FluentHook) Fire(entry *logrus.Entry) error {
 	case hook.Fluent != nil:
 		logger = hook.Fluent
 	default:
-		logger, err = fluent.New(fluent.Config{
-			FluentHost: hook.host,
-			FluentPort: hook.port,
-		})
+		logger, err = fluent.New(hook.conf.FluentConfig())
 		if err != nil {
 			return err
 		}

fluent_test.go

@@ -53,6 +53,36 @@ func TestNew(t *testing.T) {
 		t.Errorf("hook should not be nil")
 	case len(hook.levels) != len(defaultLevels):
 		t.Errorf("hook.levels should be defaultLevels")
+	case hook.Fluent == nil:
+		t.Errorf("hook.Fluent should not be nil")
+	case hook.messageField != MessageField:
+		t.Errorf("hook.messageField should be %s", MessageField)
+	}
+}
+
+func TestNewWithConfig(t *testing.T) {
+	_, port := newMockServer(t, nil)
+	conf := Config{
+		Host:                testHOST,
+		Port:                port,
+		DefaultMessageField: "DefaultMessageField",
+	}
+	hook, err := NewWithConfig(conf)
+	switch {
+	case err != nil:
+		t.Errorf("error on New: %s", err.Error())
+	case hook == nil:
+		t.Errorf("hook should not be nil")
+	case hook.conf.Host != testHOST:
+		t.Errorf("hook.host should be %s, but %s", testHOST, hook.conf.Host)
+	case hook.conf.Port != port:
+		t.Errorf("hook.port should be %d, but %d", port, hook.conf.Port)
+	case len(hook.levels) != len(defaultLevels):
+		t.Errorf("hook.levels should be defaultLevels")
+	case hook.Fluent == nil:
+		t.Errorf("hook.Fluent should not be nil")
+	case hook.messageField != "DefaultMessageField":
+		t.Errorf("hook.messageField should be DefaultMessageField")
 	}
 }
 
@@ -62,12 +92,16 @@ func TestNewHook(t *testing.T) {
 	switch {
 	case hook == nil:
 		t.Errorf("hook should not be nil")
-	case hook.host != testHOST:
-		t.Errorf("hook.host should be %s, but %s", testHOST, hook.host)
-	case hook.port != testPort:
-		t.Errorf("hook.port should be %d, but %d", testPort, hook.port)
+	case hook.conf.Host != testHOST:
+		t.Errorf("hook.host should be %s, but %s", testHOST, hook.conf.Host)
+	case hook.conf.Port != testPort:
+		t.Errorf("hook.port should be %d, but %d", testPort, hook.conf.Port)
 	case len(hook.levels) != len(defaultLevels):
 		t.Errorf("hook.levels should be defaultLevels")
+	case hook.Fluent != nil:
+		t.Errorf("hook.Fluent should be nil")
+	case hook.messageField != MessageField:
+		t.Errorf("hook.messageField should be %s", MessageField)
 	}
 }
 
@@ -244,6 +278,30 @@ func TestLogEntryMessageReceivedWithTag(t *testing.T) {
 	assertLogHook(t, f, entryMessage, assertion)
 }
 
+func TestLogEntryMessageReceivedWithCustomMessageField(t *testing.T) {
+	conf := Config{
+		DefaultTag:          fieldTag,
+		DefaultMessageField: "my-message-field",
+	}
+
+	f := logrus.Fields{
+		"value": fieldValue,
+	}
+
+	assertion := func(result string) {
+		switch {
+		case !strings.Contains(result, assertFieldTagAsFluentTag):
+			t.Errorf("message should contain fluent-tag from field")
+		case !strings.Contains(result, assertFieldValue):
+			t.Errorf("message should contain value from field")
+		case !strings.Contains(result, "\xb0my-message-field\xaeMyEntryMessage"):
+			t.Errorf("message should contain message from entry.Message")
+		}
+	}
+
+	assertLogMessageWithConfig(t, conf, f, entryMessage, assertion)
+}
+
 func TestLogEntryMessageReceivedWithMessage(t *testing.T) {
 	f := logrus.Fields{
 		"message": fieldMessage,
@@ -382,33 +440,41 @@ func assertLogMessage(t *testing.T, f logrus.Fields, message string, tag string,
 		if tag != "" {
 			hook.SetTag(tag)
 		}
-		logger := logrus.New()
-		logger.Hooks.Add(hook)
-
-		for i := 0; i < defaultLoopCount; i++ {
-			logger.WithFields(f).Error(message)
-			assertFunc(<-localData)
-		}
+		assertHook(t, hook, f, message, assertFunc, localData)
 	}
 
 	// assert persistent logger
 	{
 		port := getOrCreateMockServer(t, data)
 		hook, err := New(testHOST, port)
 		if err != nil {
-			t.Errorf("Error on NewHookWithLogger: %s", err.Error())
+			t.Errorf("Error on New: %s", err.Error())
 		}
 		if tag != "" {
 			hook.SetTag(tag)
 		}
+		assertHook(t, hook, f, message, assertFunc, data)
+	}
+}
 
-		logger := logrus.New()
-		logger.Hooks.Add(hook)
+func assertLogMessageWithConfig(t *testing.T, conf Config, f logrus.Fields, message string, assertFunc func(string)) {
+	port := getOrCreateMockServer(t, data)
+	conf.Port = port
+	conf.Host = testHOST
+	hook, err := NewWithConfig(conf)
+	if err != nil {
+		t.Errorf("Error on NewWithConfig: %s", err.Error())
+	}
+	assertHook(t, hook, f, message, assertFunc, data)
+}
 
-		for i := 0; i < defaultLoopCount; i++ {
-			logger.WithFields(f).Error(message)
-			assertFunc(<-data)
-		}
+func assertHook(t *testing.T, hook *FluentHook, f logrus.Fields, message string, assertFunc func(string), data chan string) {
+	logger := logrus.New()
+	logger.Hooks.Add(hook)
+
+	for i := 0; i < defaultLoopCount; i++ {
+		logger.WithFields(f).Error(message)
+		assertFunc(<-data)
 	}
 }