From 6253de5599ad4af71fdb979a35e0e0b27cd95ea2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E5=AF=BB=E6=AC=A2?= <lxh@cxh.cn>
Date: Wed, 29 Dec 2021 14:33:48 +0800
Subject: [PATCH] =?UTF-8?q?:sparkles:=20=E6=94=AF=E6=8C=81=E5=88=9D?=
 =?UTF-8?q?=E5=A7=8B=E5=8C=96=E7=9A=84=E6=97=B6=E5=80=99=E4=BB=8E=E7=8E=AF?=
 =?UTF-8?q?=E5=A2=83=E5=8F=98=E9=87=8F=E8=AF=BB=E5=8F=96=E9=85=8D=E7=BD=AE?=
 =?UTF-8?q?=E4=BF=A1=E6=81=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 config.go      | 13 +++++++------
 go.mod         |  1 +
 go.sum         |  2 ++
 logger.go      | 10 +++++++++-
 logger_test.go | 10 +++++++++-
 loki.go        |  7 +++----
 6 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/config.go b/config.go
index 3492dcf..9c9768d 100644
--- a/config.go
+++ b/config.go
@@ -11,12 +11,13 @@ var (
 
 // LogConfig 日志配置
 type LogConfig struct {
-	Mode       mode // dev, prod
-	LokiEnable bool
-	FileEnable bool
-	LokiHost   string
-	LokiPort   int
-	LokiName   string // Loki的job和source名称
+	Mode       mode   `env:"LOG_MODE"`             // dev, prod
+	LokiEnable bool   `env:"LOG_LOKI_ENABLE"`      // 是否启用Loki
+	FileEnable bool   `env:"LOG_FILE_ENABLE"`      // 是否输出到文件
+	LokiHost   string `env:"LOG_LOKI_HOST"`        // Loki地址
+	LokiPort   int    `env:"LOG_LOKI_PORT"`        // Loki端口
+	LokiSource string `env:"LOG_LOKI_SOURCE_NAME"` // Loki的source名称
+	LokiJob    string `env:"LOG_LOKI_JOB_NAME"`    // Loki的job名称
 }
 
 func (c LogConfig) getLokiPushURL() string {
diff --git a/go.mod b/go.mod
index 25244f3..8b434b0 100644
--- a/go.mod
+++ b/go.mod
@@ -13,6 +13,7 @@ require (
 require (
 	github.com/BurntSushi/toml v0.4.1 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/caarlos0/env/v6 v6.8.0 // indirect
 	github.com/cespare/xxhash/v2 v2.1.1 // indirect
 	github.com/go-logfmt/logfmt v0.5.0 // indirect
 	github.com/gogo/protobuf v1.3.1 // indirect
diff --git a/go.sum b/go.sum
index 0e48471..5b5f94a 100644
--- a/go.sum
+++ b/go.sum
@@ -114,6 +114,8 @@ github.com/blang/semver v3.5.1+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnweb
 github.com/bmizerany/pat v0.0.0-20170815010413-6226ea591a40/go.mod h1:8rLXio+WjiTceGBHIoTvn60HIbs7Hm7bcHjyrSqYB9c=
 github.com/boltdb/bolt v1.3.1/go.mod h1:clJnj/oiGkjum5o1McbSZDSLxVThjynRyGBgiAx27Ps=
 github.com/c-bata/go-prompt v0.2.2/go.mod h1:VzqtzE2ksDBcdln8G7mk2RX9QyGjH+OVqOCSiVIqS34=
+github.com/caarlos0/env/v6 v6.8.0 h1:abF9JinEXaibthiOowf4uSnRBWN66aJOxSpHLH67jeI=
+github.com/caarlos0/env/v6 v6.8.0/go.mod h1:FE0jGiAnQqtv2TenJ4KTa8+/T2Ss8kdS5s1VEjasoN0=
 github.com/casbin/casbin/v2 v2.1.2/go.mod h1:YcPU1XXisHhLzuxH9coDNf2FbKpjGlbCg3n9yuLkIJQ=
 github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM=
 github.com/cenkalti/backoff/v4 v4.0.2/go.mod h1:eEew/i+1Q6OrCDZh3WiXYv3+nJwBASZ8Bog/87DQnVg=
diff --git a/logger.go b/logger.go
index 34c2447..4259776 100644
--- a/logger.go
+++ b/logger.go
@@ -1,6 +1,8 @@
 package logger
 
 import (
+	"fmt"
+	"github.com/caarlos0/env/v6"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
@@ -11,7 +13,13 @@ var Say *zap.SugaredLogger
 // 避免异常，在第一次调用时初始化一个只打印到控制台的logger
 func init() {
 	if Say == nil {
-		InitLogger(LogConfig{Mode: Dev, LokiEnable: false, FileEnable: false})
+		// 从环境变量读取配置
+		var c LogConfig
+		if err := env.Parse(&c); err != nil {
+			fmt.Println("日志配置解析错误: " + err.Error())
+			c = LogConfig{Mode: Dev, LokiEnable: false, FileEnable: false}
+		}
+		InitLogger(c)
 	}
 }
 
diff --git a/logger_test.go b/logger_test.go
index 98c5ee1..f8a0318 100644
--- a/logger_test.go
+++ b/logger_test.go
@@ -1,8 +1,16 @@
 package logger
 
-import "testing"
+import (
+	"testing"
+	"time"
+)
 
 func TestLogger(t *testing.T) {
 	InitLogger(LogConfig{Mode: Dev, LokiEnable: false, FileEnable: true})
 	Say.Debug("芜湖")
 }
+
+func TestLogger1(t *testing.T) {
+	Say.Info("我是测试消息")
+	time.Sleep(5 * time.Second)
+}
diff --git a/loki.go b/loki.go
index b82116a..897ed58 100644
--- a/loki.go
+++ b/loki.go
@@ -69,14 +69,13 @@ func (c lokiWriter) Write(p []byte) (int, error) {
 		return 0, err
 	}
 
-	label := model.LabelSet{"job": model.LabelValue(config.LokiName)}
-	label["source"] = model.LabelValue(config.LokiName)
+	label := model.LabelSet{"job": model.LabelValue(config.LokiJob)}
+	label["source"] = model.LabelValue(config.LokiSource)
 	label["level"] = model.LabelValue(li.Level)
 	label["caller"] = model.LabelValue(li.Caller)
 	// 异步推送消息到服务器
 	go func() {
-		err = lokiClient.Handle(label, time.Now().Local(), li.Msg)
-		if err != nil {
+		if err = lokiClient.Handle(label, time.Now().Local(), li.Msg); err != nil {
 			fmt.Printf("日志推送到Loki失败: %v\n", err.Error())
 		}
 	}()