From 2d420a6dfdf0d716fdd514b806e5779eb46ca00b Mon Sep 17 00:00:00 2001
From: niuyufu <niuyufu@didichuxing.com>
Date: Sat, 3 Nov 2018 01:53:47 +0800
Subject: [PATCH] Log adds context features

---
 logger.go      | 14 ++++++++++++-
 logger_test.go | 54 ++++++++++++++++++++++++++++++++++++++++++++++++++
 main.go        | 31 ++++++++++++++++++++++++++++-
 3 files changed, 97 insertions(+), 2 deletions(-)
 create mode 100644 logger_test.go

diff --git a/logger.go b/logger.go
index 4324a2e4..bad45857 100644
--- a/logger.go
+++ b/logger.go
@@ -43,7 +43,6 @@ var LogFormatter = func(values ...interface{}) (messages []interface{}) {
 			// duration
 			messages = append(messages, fmt.Sprintf(" \033[36;1m[%.2fms]\033[0m ", float64(values[2].(time.Duration).Nanoseconds()/1e4)/100.0))
 			// sql
-
 			for _, value := range values[4].([]interface{}) {
 				indirectValue := reflect.Indirect(reflect.ValueOf(value))
 				if indirectValue.IsValid() {
@@ -101,6 +100,7 @@ var LogFormatter = func(values ...interface{}) (messages []interface{}) {
 
 type logger interface {
 	Print(v ...interface{})
+	CtxPrint(s *DB,v ...interface{})
 }
 
 // LogWriter log writer interface
@@ -117,3 +117,15 @@ type Logger struct {
 func (logger Logger) Print(values ...interface{}) {
 	logger.Println(LogFormatter(values...)...)
 }
+
+// Print format & print log
+func (logger Logger) CtxPrint(s *DB,values ...interface{}) {
+	ctx,ok:=s.GetCtx()
+	if ok{
+		logMessage:=[]interface{}{}
+		logMessage=append(LogFormatter(values...), ctx)
+		logger.Println(logMessage...)
+	}else{
+		logger.Println(LogFormatter(values...)...)
+	}
+}
diff --git a/logger_test.go b/logger_test.go
new file mode 100644
index 00000000..96d97f4e
--- /dev/null
+++ b/logger_test.go
@@ -0,0 +1,54 @@
+package gorm_test
+
+import (
+	"testing"
+	"os"
+	"crypto/md5"
+	"fmt"
+	"time"
+	"github.com/jinzhu/gorm"
+	"log"
+)
+
+func TestLoggerCtx(t *testing.T) {
+	DB.SetLogger(gorm.Logger{log.New(os.Stdout, "\r\n", 0)})
+	if debug := os.Getenv("DEBUG"); debug == "true" {
+		DB.LogMode(true)
+	} else if debug == "false" {
+		DB.LogMode(false)
+	}
+
+	if logCtx := os.Getenv("LOGCTX"); logCtx == "true" {
+		DB.LogCtx(true)
+	} else if logCtx == "false" {
+		DB.LogCtx(false)
+	}
+
+	i := 0
+	for i < 10 {
+		i++
+
+		//Generating context information
+		unixTime := fmt.Sprint(time.Now().Unix())
+		traceId := fmt.Sprintf("%x", md5.Sum([]byte(unixTime)))
+		ctxInfo:= "\n[context] trace_id="+traceId
+		builder := DB.SetCtx(ctxInfo)
+		if i > 5 {
+			builder = builder.Where("Age = ?", i)
+		} else {
+			builder = builder.Where("Name = ?", i)
+		}
+
+		if builder.Find(&User{}).Error == nil {
+			t.Errorf("Should got error with invalid SQL")
+		}
+
+		//Verify context information
+		ctxTmp,_:=builder.GetCtx()
+		ctxInfo2,_:=ctxTmp.(string)
+		if ctxInfo!=ctxInfo2{
+			t.Fatal("get context error")
+		}
+	}
+}
+
diff --git a/main.go b/main.go
index 17c75ed3..e7b3abb2 100644
--- a/main.go
+++ b/main.go
@@ -20,6 +20,7 @@ type DB struct {
 	db                SQLCommon
 	blockGlobalUpdate bool
 	logMode           int
+	logCtx            int
 	logger            logger
 	search            *search
 	values            sync.Map
@@ -148,6 +149,30 @@ func (s *DB) LogMode(enable bool) *DB {
 	return s
 }
 
+// LogMode set log ctx, `true` for get context when printing logs, `false` for no context, default, false
+func (s *DB) LogCtx(enable bool) *DB {
+	s.logCtx = 2
+	return s
+	fmt.Println(enable)
+	if enable {
+		s.logCtx = 2
+	} else {
+		s.logCtx = 1
+	}
+	return s
+}
+
+// Set query context
+func (s *DB) SetCtx(v interface{}) *DB {
+	s=s.Set("trace_context",v)
+	return s
+}
+
+// Get query context
+func (s *DB) GetCtx() (value interface{}, ok bool)  {
+	return s.Get("trace_context")
+}
+
 // BlockGlobalUpdate if true, generates an error on update/delete without where clause.
 // This is to prevent eventual error with empty objects updates/deletions
 func (s *DB) BlockGlobalUpdate(enable bool) *DB {
@@ -776,7 +801,11 @@ func (s *DB) clone() *DB {
 }
 
 func (s *DB) print(v ...interface{}) {
-	s.logger.Print(v...)
+	if s.logCtx == 2{
+		s.logger.CtxPrint(s, v...)
+	}else{
+		s.logger.Print(v...)
+	}
 }
 
 func (s *DB) log(v ...interface{}) {