fix: AT mode queryContext error

Author: 576470954

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(),
 		}

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)

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,