From da55c35b25c38f46d9b541be9d1970c04636c4fb Mon Sep 17 00:00:00 2001
From: "yong.an" <yong.an@shopee.com>
Date: Thu, 9 Jul 2020 16:37:17 +0800
Subject: [PATCH] Used to monitor which DB executed this SQL

---
 session.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/session.go b/session.go
index 761b1415..d82c7cad 100644
--- a/session.go
+++ b/session.go
@@ -101,13 +101,15 @@ func newSessionID() string {
 
 func newSession(engine *Engine) *Session {
 	var ctx context.Context
+	var sessionID string
 	if engine.logSessionID {
-		ctx = context.WithValue(engine.defaultContext, log.SessionIDKey, newSessionID())
+		sessionID = newSessionID()
+		ctx = context.WithValue(engine.defaultContext, log.SessionIDKey, sessionID)
 	} else {
 		ctx = engine.defaultContext
 	}
 
-	return &Session{
+	session := &Session{
 		ctx:    ctx,
 		engine: engine,
 		tx:     nil,
@@ -136,6 +138,10 @@ func newSession(engine *Engine) *Session {
 
 		sessionType: engineSession,
 	}
+	if engine.logSessionID {
+		ctx = context.WithValue(engine.defaultContext, sessionID, session)
+	}
+	return session
 }
 
 // Close release the connection from pool