From 0851ae825d82aa7426003fc4d4c2ab335795bfc0 Mon Sep 17 00:00:00 2001
From: "shipei.liu" <shipei.liu@rdc-west.com>
Date: Thu, 15 Jul 2021 18:36:48 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=F0=9F=8E=B8=20set=20the=20default=20ti?=
 =?UTF-8?q?meout=20for=20every=20session?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 engine.go  | 11 +++++++++++
 session.go | 10 +++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/engine.go b/engine.go
index d3ee8a8c..7aeca0dc 100644
--- a/engine.go
+++ b/engine.go
@@ -46,6 +46,7 @@ type Engine struct {
 	DatabaseTZ *time.Location // The timezone of the database
 
 	logSessionID bool // create session id
+	timeoutSecond uint // if > 0 set default timeout for session
 }
 
 // NewEngine new a db manager according to the parameter. Currently support four
@@ -273,6 +274,16 @@ func (engine *Engine) SetDefaultCacher(cacher caches.Cacher) {
 	engine.cacherMgr.SetDefaultCacher(cacher)
 }
 
+// SetSessionTimeout set the default timeout for every session
+func (engine *Engine) SetSessionTimeout(timeout uint) {
+	engine.timeoutSecond = timeout
+}
+
+// GetDefaultTimeout return the default timeout of new session
+func (engine *Engine) GetDefaultTimeout() uint {
+	return engine.timeoutSecond
+}
+
 // GetDefaultCacher returns the default cacher
 func (engine *Engine) GetDefaultCacher() caches.Cacher {
 	return engine.cacherMgr.GetDefaultCacher()
diff --git a/session.go b/session.go
index 486911a5..a4292362 100644
--- a/session.go
+++ b/session.go
@@ -85,6 +85,7 @@ type Session struct {
 	lastSQLArgs []interface{}
 
 	ctx         context.Context
+	cancel		context.CancelFunc
 	sessionType sessionType
 }
 
@@ -106,9 +107,13 @@ func newSession(engine *Engine) *Session {
 	} else {
 		ctx = engine.defaultContext
 	}
-
+	var cancel context.CancelFunc
+	if engine.timeoutSecond > 0 {
+		ctx, cancel = context.WithTimeout(ctx, time.Second*time.Duration(engine.timeoutSecond))
+	}
 	session := &Session{
 		ctx:    ctx,
+		cancel: cancel,
 		engine: engine,
 		tx:     nil,
 		statement: statements.NewStatement(
@@ -144,6 +149,9 @@ func newSession(engine *Engine) *Session {
 
 // Close release the connection from pool
 func (session *Session) Close() error {
+	if session.cancel != nil{
+		session.cancel()
+	}
 	for _, v := range session.stmtCache {
 		if err := v.Close(); err != nil {
 			return err