Merge pull request #11 from mariussturm/master

Add log rotation through io.Writer interface

README.md

@@ -29,5 +29,27 @@ func NewLogger( config map[string]interface{} ) *log.Logger {
 ### Formatters
 lfshook will strip colors from any TextFormatter type formatters when writing to local file, because the color codes don't look so great.
 
+### Log rotation
+In order to enable automatic log rotation it's possible to provide an io.Writer instead of the path string of a log file.
+In combination with pakages like [go-file-rotatelogs](https://github.com/lestrrat/go-file-rotatelogs) log rotation can easily be achieved.
+
+```go
+  import "github.com/lestrrat/go-file-rotatelogs"
+
+  path := "/var/log/go.log"
+  writer := rotatelogs.NewRotateLogs(
+    path + ".%Y%m%d%H%M", // rotation pattern
+  )
+  writer.LinkName = path
+  writer.RotationTime = time.Duration(86400) * time.Second // rotate once a day
+  writer.MaxAge = time.Duration(604800) * time.Second // keep one week of log files
+  writer.Offset = 0
+
+  log.Hooks.Add(lfshook.NewHook(lfshook.WriterMap{
+    logrus.InfoLevel: writer,
+    logrus.ErrorLevel: writer,
+  }))
+```
+
 ### Note:
 Whichever user is running the go application must have read/write permissions to the log files selected, or if the files do not yet exists, then to the directory in which the files will be created.

lfshook.go

@@ -3,10 +3,13 @@ package lfshook
 
 import (
 	"fmt"
-	"github.com/Sirupsen/logrus"
+	"io"
 	"log"
 	"os"
+	"reflect"
 	"sync"
+
+	"github.com/Sirupsen/logrus"
 )
 
 // We are logging to file, strip colors to make the output more readable
@@ -16,42 +19,109 @@ var txtFormatter = &logrus.TextFormatter{DisableColors: true}
 // Multiple levels may share a file, but multiple files may not be used for one level
 type PathMap map[logrus.Level]string
 
+// Alternatively map a log level to an io.Writer
+type WriterMap map[logrus.Level]io.Writer
+
 // Hook to handle writing to local log files.
 type lfsHook struct {
-	paths  PathMap
-	levels []logrus.Level
-	lock   *sync.Mutex
+	paths     PathMap
+	writer    WriterMap
+	levels    []logrus.Level
+	lock      *sync.Mutex
 	formatter logrus.Formatter
 }
 
 // Given a map with keys equal to log levels.
 // We can generate our levels handled on the fly, and write to a specific file for each level.
 // We can also write to the same file for all levels. They just need to be specified.
-func NewHook(levelMap PathMap) *lfsHook {
+func NewHook(levelMap interface{}) *lfsHook {
 	hook := &lfsHook{
-		paths: levelMap,
-		lock:  new(sync.Mutex),
+		lock:      new(sync.Mutex),
 		formatter: txtFormatter,
 	}
-	for level, _ := range levelMap {
-		hook.levels = append(hook.levels, level)
+
+	switch levelMap.(type) {
+	case PathMap:
+		hook.paths = levelMap.(PathMap)
+		for level := range levelMap.(PathMap) {
+			hook.levels = append(hook.levels, level)
+		}
+		break
+	case WriterMap:
+		hook.writer = levelMap.(WriterMap)
+		for level := range levelMap.(WriterMap) {
+			hook.levels = append(hook.levels, level)
+		}
+		break
+	default:
+		panic(fmt.Sprintf("unsupported level map type: %s", reflect.TypeOf(levelMap)))
 	}
+
 	return hook
 }
 
+// Replace the color stripped default formatter with a custom formatter
 func (hook *lfsHook) SetFormatter(formatter logrus.Formatter) {
 	hook.formatter = formatter
 
 	switch hook.formatter.(type) {
-		case *logrus.TextFormatter:
-			textFormatter := hook.formatter.(*logrus.TextFormatter)
-			textFormatter.DisableColors = true
+	case *logrus.TextFormatter:
+		textFormatter := hook.formatter.(*logrus.TextFormatter)
+		textFormatter.DisableColors = true
 	}
 }
 
 // Open the file, write to the file, close the file.
 // Whichever user is running the function needs write permissions to the file or directory if the file does not yet exist.
 func (hook *lfsHook) Fire(entry *logrus.Entry) error {
+	// only modify Formatter if we are using a TextFormatter so we can strip colors
+	switch entry.Logger.Formatter.(type) {
+	case *logrus.TextFormatter:
+		// swap to colorless TextFormatter
+		formatter := entry.Logger.Formatter
+		entry.Logger.Formatter = txtFormatter
+		defer func() {
+			// assign back original formatter
+			entry.Logger.Formatter = formatter
+		}()
+	}
+
+	if hook.writer != nil {
+		return hook.ioWrite(entry)
+	} else {
+		return hook.fileWrite(entry)
+	}
+}
+
+// Write a log line to an io.Writer
+func (hook *lfsHook) ioWrite(entry *logrus.Entry) error {
+	var (
+		msg string
+		err error
+		ok  bool
+	)
+
+	hook.lock.Lock()
+	defer hook.lock.Unlock()
+
+	if _, ok = hook.writer[entry.Level]; !ok {
+		err = fmt.Errorf("no writer provided for loglevel: %d", entry.Level)
+		log.Println(err.Error())
+		return err
+	}
+
+	msg, err = entry.String()
+
+	if err != nil {
+		log.Println("failed to generate string for entry:", err)
+		return err
+	}
+	_, err = hook.writer[entry.Level].Write([]byte(msg))
+	return err
+}
+
+// Write a log line directly to a file
+func (hook *lfsHook) fileWrite(entry *logrus.Entry) error {
 	var (
 		fd   *os.File
 		path string
@@ -75,18 +145,6 @@ func (hook *lfsHook) Fire(entry *logrus.Entry) error {
 	}
 	defer fd.Close()
 
-	// only modify Formatter if we are using a TextFormatter so we can strip colors
-	switch entry.Logger.Formatter.(type) {
-	case *logrus.TextFormatter:
-		// swap to colorless TextFormatter
-		formatter := entry.Logger.Formatter
-		entry.Logger.Formatter = txtFormatter
-		defer func() {
-			// assign back original formatter
-			entry.Logger.Formatter = formatter
-		}()
-	}
-
 	msg, err = entry.String()
 
 	if err != nil {
@@ -97,6 +155,7 @@ func (hook *lfsHook) Fire(entry *logrus.Entry) error {
 	return nil
 }
 
+// Return configured log levels
 func (hook *lfsHook) Levels() []logrus.Level {
 	return hook.levels
 }