From 0fae64bb3b4c280dc2f8ab6fe6d15e80068354bc Mon Sep 17 00:00:00 2001
From: Thomas_An <thomas_an@noreply.gitea.io>
Date: Mon, 13 Jul 2020 13:29:38 +0000
Subject: [PATCH] Support get dataSourceName on ContextHook for  monitor which
 DB executed SQL (#1740)

Merge branch 'log_context_add_db_info'

add session to context

Revert "accept lunny's suggestion"

This reverts commit 57fd6699425bdb97baf6bbeddc3099ffcb4576d8.

Merge branch 'master' of https://gitea.com/Thomas_An/thomasan_xorm

accept lunny's suggestion

Merge branch 'master' into master

add test code

session add Engine func

Used to monitor which DB executed this SQL

Co-authored-by: yong.an <yong.an@shopee.com>
Reviewed-on: https://gitea.com/xorm/xorm/pulls/1740
Reviewed-by: Lunny Xiao <xiaolunwen@gmail.com>
---
 integrations/session_test.go |  8 ++++++++
 interface.go                 |  1 +
 log/logger_context.go        |  1 +
 session.go                   | 12 ++++++++++--
 4 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/integrations/session_test.go b/integrations/session_test.go
index bdf3278d..af8b0375 100644
--- a/integrations/session_test.go
+++ b/integrations/session_test.go
@@ -54,3 +54,11 @@ func TestMustLogSQL(t *testing.T) {
 	_, err := testEngine.Table("userinfo").MustLogSQL(true).Get(new(Userinfo))
 	assert.NoError(t, err)
 }
+
+func TestEnableSessionId(t *testing.T) {
+	assert.NoError(t, PrepareEngine())
+	testEngine.EnableSessionID(true)
+	assertSync(t, new(Userinfo))
+	_, err := testEngine.Table("userinfo").MustLogSQL(true).Get(new(Userinfo))
+	assert.NoError(t, err)
+}
\ No newline at end of file
diff --git a/interface.go b/interface.go
index 6aac4ae8..0fe9cbe1 100644
--- a/interface.go
+++ b/interface.go
@@ -120,6 +120,7 @@ type EngineInterface interface {
 	TableInfo(bean interface{}) (*schemas.Table, error)
 	TableName(interface{}, ...bool) string
 	UnMapType(reflect.Type)
+	EnableSessionID(bool)
 }
 
 var (
diff --git a/log/logger_context.go b/log/logger_context.go
index 6b7252ef..46802576 100644
--- a/log/logger_context.go
+++ b/log/logger_context.go
@@ -42,6 +42,7 @@ var (
 // enumerate all the context keys
 var (
 	SessionIDKey      = "__xorm_session_id"
+	SessionKey        = "__xorm_session_key"
 	SessionShowSQLKey = "__xorm_show_sql"
 )
 
diff --git a/session.go b/session.go
index 761b1415..48b3779e 100644
--- a/session.go
+++ b/session.go
@@ -102,12 +102,12 @@ func newSessionID() string {
 func newSession(engine *Engine) *Session {
 	var ctx context.Context
 	if engine.logSessionID {
-		ctx = context.WithValue(engine.defaultContext, log.SessionIDKey, newSessionID())
+		ctx = context.WithValue(engine.defaultContext, log.SessionIDKey,  newSessionID())
 	} else {
 		ctx = engine.defaultContext
 	}
 
-	return &Session{
+	session := &Session{
 		ctx:    ctx,
 		engine: engine,
 		tx:     nil,
@@ -136,6 +136,10 @@ func newSession(engine *Engine) *Session {
 
 		sessionType: engineSession,
 	}
+	if engine.logSessionID {
+		session.ctx = context.WithValue(session.ctx, log.SessionKey, session)
+	}
+	return session
 }
 
 // Close release the connection from pool
@@ -165,6 +169,10 @@ func (session *Session) db() *core.DB {
 	return session.engine.db
 }
 
+func (session *Session) Engine() *Engine {
+	return session.engine
+}
+
 func (session *Session) getQueryer() core.Queryer {
 	if session.tx != nil {
 		return session.tx