lfshook.go 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. // Package lfshook is hook for sirupsen/logrus that used for writing the logs to local files.
  2. package lfshook
  3. import (
  4. "fmt"
  5. "github.com/sirupsen/logrus"
  6. "io"
  7. "log"
  8. "os"
  9. "path/filepath"
  10. "reflect"
  11. "sync"
  12. )
  13. // We are logging to file, strip colors to make the output more readable.
  14. var defaultFormatter = &logrus.TextFormatter{DisableColors: true}
  15. // PathMap is map for mapping a log level to a file's path.
  16. // Multiple levels may share a file, but multiple files may not be used for one level.
  17. type PathMap map[logrus.Level]string
  18. // WriterMap is map for mapping a log level to an io.Writer.
  19. // Multiple levels may share a writer, but multiple writers may not be used for one level.
  20. type WriterMap map[logrus.Level]io.Writer
  21. // LfsHook is a hook to handle writing to local log files.
  22. type LfsHook struct {
  23. paths PathMap
  24. writers WriterMap
  25. levels []logrus.Level
  26. lock *sync.Mutex
  27. formatter logrus.Formatter
  28. defaultPath string
  29. defaultWriter io.Writer
  30. hasDefaultPath bool
  31. hasDefaultWriter bool
  32. }
  33. // NewHook returns new LFS hook.
  34. // Output can be a string, io.Writer, WriterMap or PathMap.
  35. // If using io.Writer or WriterMap, user is responsible for closing the used io.Writer.
  36. func NewHook(output interface{}, formatter logrus.Formatter) *LfsHook {
  37. hook := &LfsHook{
  38. lock: new(sync.Mutex),
  39. }
  40. hook.SetFormatter(formatter)
  41. switch output.(type) {
  42. case string:
  43. hook.SetDefaultPath(output.(string))
  44. break
  45. case io.Writer:
  46. hook.SetDefaultWriter(output.(io.Writer))
  47. break
  48. case PathMap:
  49. hook.paths = output.(PathMap)
  50. for level := range output.(PathMap) {
  51. hook.levels = append(hook.levels, level)
  52. }
  53. break
  54. case WriterMap:
  55. hook.writers = output.(WriterMap)
  56. for level := range output.(WriterMap) {
  57. hook.levels = append(hook.levels, level)
  58. }
  59. break
  60. default:
  61. panic(fmt.Sprintf("unsupported level map type: %v", reflect.TypeOf(output)))
  62. }
  63. return hook
  64. }
  65. // SetFormatter sets the format that will be used by hook.
  66. // If using text formatter, this method will disable color output to make the log file more readable.
  67. func (hook *LfsHook) SetFormatter(formatter logrus.Formatter) {
  68. hook.lock.Lock()
  69. defer hook.lock.Unlock()
  70. if formatter == nil {
  71. formatter = defaultFormatter
  72. } else {
  73. switch formatter.(type) {
  74. case *logrus.TextFormatter:
  75. textFormatter := formatter.(*logrus.TextFormatter)
  76. textFormatter.DisableColors = true
  77. }
  78. }
  79. hook.formatter = formatter
  80. }
  81. // SetDefaultPath sets default path for levels that don't have any defined output path.
  82. func (hook *LfsHook) SetDefaultPath(defaultPath string) {
  83. hook.lock.Lock()
  84. defer hook.lock.Unlock()
  85. hook.defaultPath = defaultPath
  86. hook.hasDefaultPath = true
  87. }
  88. // SetDefaultWriter sets default writer for levels that don't have any defined writer.
  89. func (hook *LfsHook) SetDefaultWriter(defaultWriter io.Writer) {
  90. hook.lock.Lock()
  91. defer hook.lock.Unlock()
  92. hook.defaultWriter = defaultWriter
  93. hook.hasDefaultWriter = true
  94. }
  95. // Fire writes the log file to defined path or using the defined writer.
  96. // User who run this function needs write permissions to the file or directory if the file does not yet exist.
  97. func (hook *LfsHook) Fire(entry *logrus.Entry) error {
  98. hook.lock.Lock()
  99. defer hook.lock.Unlock()
  100. if hook.writers != nil || hook.hasDefaultWriter {
  101. return hook.ioWrite(entry)
  102. } else if hook.paths != nil || hook.hasDefaultPath {
  103. return hook.fileWrite(entry)
  104. }
  105. return nil
  106. }
  107. // Write a log line to an io.Writer.
  108. func (hook *LfsHook) ioWrite(entry *logrus.Entry) error {
  109. var (
  110. writer io.Writer
  111. msg []byte
  112. err error
  113. ok bool
  114. )
  115. if writer, ok = hook.writers[entry.Level]; !ok {
  116. if hook.hasDefaultWriter {
  117. writer = hook.defaultWriter
  118. } else {
  119. return nil
  120. }
  121. }
  122. // use our formatter instead of entry.String()
  123. msg, err = hook.formatter.Format(entry)
  124. if err != nil {
  125. log.Println("failed to generate string for entry:", err)
  126. return err
  127. }
  128. _, err = writer.Write(msg)
  129. return err
  130. }
  131. // Write a log line directly to a file.
  132. func (hook *LfsHook) fileWrite(entry *logrus.Entry) error {
  133. var (
  134. fd *os.File
  135. path string
  136. msg []byte
  137. err error
  138. ok bool
  139. )
  140. if path, ok = hook.paths[entry.Level]; !ok {
  141. if hook.hasDefaultPath {
  142. path = hook.defaultPath
  143. } else {
  144. return nil
  145. }
  146. }
  147. dir := filepath.Dir(path)
  148. os.MkdirAll(dir, os.ModePerm)
  149. fd, err = os.OpenFile(path, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
  150. if err != nil {
  151. log.Println("failed to open logfile:", path, err)
  152. return err
  153. }
  154. defer fd.Close()
  155. // use our formatter instead of entry.String()
  156. msg, err = hook.formatter.Format(entry)
  157. if err != nil {
  158. log.Println("failed to generate string for entry:", err)
  159. return err
  160. }
  161. fd.Write(msg)
  162. return nil
  163. }
  164. // Levels returns configured log levels.
  165. func (hook *LfsHook) Levels() []logrus.Level {
  166. return logrus.AllLevels
  167. }