From 9709b356da474cb4e043b75c294c7dd90ddca2ef Mon Sep 17 00:00:00 2001
From: 576470954 <576470954@qq.com>
Date: Sat, 3 Feb 2024 17:52:45 +0800
Subject: [PATCH] fix: AT mode queryContext error

---
 pkg/datasource/sql/conn_at.go                        |  1 +
 pkg/datasource/sql/datasource/base/meta_cache.go     |  2 ++
 .../sql/exec/at/select_for_update_executor.go        | 12 ++++++------
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/pkg/datasource/sql/conn_at.go b/pkg/datasource/sql/conn_at.go
index 1ddfba517..e640254ac 100644
--- a/pkg/datasource/sql/conn_at.go
+++ b/pkg/datasource/sql/conn_at.go
@@ -62,6 +62,7 @@ func (c *ATConn) QueryContext(ctx context.Context, query string, args []driver.N
 			Query:                query,
 			NamedValues:          args,
 			Conn:                 c.targetConn,
+			DBName:               c.dbName,
 			IsSupportsSavepoints: true,
 			IsAutoCommit:         c.GetAutoCommit(),
 		}
diff --git a/pkg/datasource/sql/datasource/base/meta_cache.go b/pkg/datasource/sql/datasource/base/meta_cache.go
index 1edbf38e4..b7e4e15e9 100644
--- a/pkg/datasource/sql/datasource/base/meta_cache.go
+++ b/pkg/datasource/sql/datasource/base/meta_cache.go
@@ -138,6 +138,8 @@ func (c *BaseTableMetaCache) GetTableMeta(ctx context.Context, dbName, tableName
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
+	defer conn.Close()
+
 	v, ok := c.cache[tableName]
 	if !ok {
 		meta, err := c.trigger.LoadOne(ctx, dbName, tableName, conn)
diff --git a/pkg/datasource/sql/exec/at/select_for_update_executor.go b/pkg/datasource/sql/exec/at/select_for_update_executor.go
index 3e51a97b1..289b666b6 100644
--- a/pkg/datasource/sql/exec/at/select_for_update_executor.go
+++ b/pkg/datasource/sql/exec/at/select_for_update_executor.go
@@ -174,12 +174,6 @@ func (s *selectForUpdateExecutor) doExecContext(ctx context.Context, f exec.Call
 		return nil, fmt.Errorf("not support savepoint. please check your db version")
 	}
 
-	// execute business SQL, try to get local lock
-	result, err = f(ctx, s.execContext.Query, s.execContext.NamedValues)
-	if err != nil {
-		return nil, err
-	}
-
 	// query primary key values
 	var lockKey string
 	_, err = s.exec(ctx, s.selectPKSQL, s.execContext.NamedValues, func(rows driver.Rows) {
@@ -194,6 +188,12 @@ func (s *selectForUpdateExecutor) doExecContext(ctx context.Context, f exec.Call
 		return nil, nil
 	}
 
+	// execute business SQL, try to get local lock
+	result, err = f(ctx, s.execContext.Query, s.execContext.NamedValues)
+	if err != nil {
+		return nil, err
+	}
+
 	// check global lock
 	lockable, err := datasource.GetDataSourceManager(branch.BranchTypeAT).LockQuery(ctx, rm.LockQueryParam{
 		Xid:        s.execContext.TxCtx.XID,