This is an automated email from the ASF dual-hosted git repository.
tew pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-seata-go.git
The following commit(s) were added to refs/heads/master by this push:
new 2cce81bc feat: add QueryContext for AT mode (#1039)
2cce81bc is described below
commit 2cce81bc00446d0849c12f6f0742829eaae72138
Author: simple-xair <[email protected]>
AuthorDate: Sat Mar 14 19:35:38 2026 +0800
feat: add QueryContext for AT mode (#1039)
* feat: add QueryContext for AT mode
* refactor: split createNewTxOnExec
* fix: add missing /v2 suffix in import paths
---------
Co-authored-by: lxfeng1997 <[email protected]>
---
pkg/datasource/sql/conn.go | 14 +-
pkg/datasource/sql/conn_at.go | 121 +++++++++-
.../sql/exec/select_for_update_executor.go | 4 +-
pkg/datasource/sql/rows_commit_on_close.go | 91 +++++++
pkg/datasource/sql/rows_commit_on_close_test.go | 265 +++++++++++++++++++++
pkg/datasource/sql/types/types.go | 2 +
6 files changed, 470 insertions(+), 27 deletions(-)
diff --git a/pkg/datasource/sql/conn.go b/pkg/datasource/sql/conn.go
index 7a57c829..f40e8153 100644
--- a/pkg/datasource/sql/conn.go
+++ b/pkg/datasource/sql/conn.go
@@ -114,11 +114,7 @@ func (c *Conn) Exec(query string, args []driver.Value)
(driver.Result, error) {
func (c *Conn) ExecContext(ctx context.Context, query string, args
[]driver.NamedValue) (driver.Result, error) {
targetConn, ok := c.targetConn.(driver.ExecerContext)
if !ok {
- values := make([]driver.Value, 0, len(args))
- for i := range args {
- values = append(values, args[i].Value)
- }
- return c.Exec(query, values)
+ return c.Exec(query, util.NamedValueToValue(args))
}
ret, err := targetConn.ExecContext(ctx, query, args)
@@ -166,13 +162,7 @@ func (c *Conn) Query(query string, args []driver.Value)
(driver.Rows, error) {
func (c *Conn) QueryContext(ctx context.Context, query string, args
[]driver.NamedValue) (driver.Rows, error) {
conn, ok := c.targetConn.(driver.QueryerContext)
if !ok {
- values := make([]driver.Value, 0, len(args))
-
- for i := range args {
- values = append(values, args[i].Value)
- }
-
- return c.Query(query, values)
+ return c.Query(query, util.NamedValueToValue(args))
}
ret, err := conn.QueryContext(ctx, query, args)
diff --git a/pkg/datasource/sql/conn_at.go b/pkg/datasource/sql/conn_at.go
index 576d22f4..e19aa017 100644
--- a/pkg/datasource/sql/conn_at.go
+++ b/pkg/datasource/sql/conn_at.go
@@ -51,7 +51,7 @@ func (c *ATConn) ExecContext(ctx context.Context, query
string, args []driver.Na
}()
}
- ret, err := c.createNewTxOnExecIfNeed(ctx, func() (types.ExecResult,
error) {
+ ret, err := c.createTxAndExecIfNeeded(ctx, func() (types.ExecResult,
error) {
executor, err := exec.BuildExecutor(c.res.dbType,
c.txCtx.TransactionMode, query)
if err != nil {
return nil, err
@@ -85,6 +85,48 @@ func (c *ATConn) ExecContext(ctx context.Context, query
string, args []driver.Na
return ret.GetResult(), nil
}
+// QueryContext
+func (c *ATConn) QueryContext(ctx context.Context, query string, args
[]driver.NamedValue) (driver.Rows, error) {
+ if c.createOnceTxContext(ctx) {
+ defer func() {
+ c.txCtx = types.NewTxCtx()
+ }()
+ }
+
+ ret, err := c.createTxAndQueryIfNeeded(ctx, func() (types.ExecResult,
error) {
+ executor, err := exec.BuildExecutor(c.res.dbType,
c.txCtx.TransactionMode, query)
+ if err != nil {
+ return nil, err
+ }
+
+ execCtx := &types.ExecContext{
+ TxCtx: c.txCtx,
+ Query: query,
+ NamedValues: args,
+ Conn: c.targetConn,
+ DBName: c.dbName,
+ DbVersion: c.GetDbVersion(),
+ IsSupportsSavepoints: true,
+ IsAutoCommit: c.GetAutoCommit(),
+ }
+
+ ret, err := executor.ExecWithNamedValue(ctx, execCtx,
+ func(ctx context.Context, query string, args
[]driver.NamedValue) (types.ExecResult, error) {
+ ret, err := c.Conn.QueryContext(ctx, query,
args)
+ if err != nil {
+ return nil, err
+ }
+ return types.NewResult(types.WithRows(ret)), nil
+ })
+
+ return ret, err
+ })
+ if err != nil {
+ return nil, err
+ }
+ return ret.GetRows(), nil
+}
+
// BeginTx
func (c *ATConn) BeginTx(ctx context.Context, opts driver.TxOptions)
(driver.Tx, error) {
c.autoCommit = false
@@ -122,7 +164,8 @@ func (c *ATConn) createOnceTxContext(ctx context.Context)
bool {
return onceTx
}
-func (c *ATConn) createNewTxOnExecIfNeed(ctx context.Context, f func()
(types.ExecResult, error)) (types.ExecResult, error) {
+// createTxAndExecIfNeeded creates a transaction for execution context and
commits it after execution
+func (c *ATConn) createTxAndExecIfNeeded(ctx context.Context, f func()
(types.ExecResult, error)) (types.ExecResult, error) {
var (
tx driver.Tx
err error
@@ -133,25 +176,26 @@ func (c *ATConn) createNewTxOnExecIfNeed(ctx
context.Context, f func() (types.Ex
if err != nil {
return nil, err
}
- }
- defer func() {
- recoverErr := recover()
- if recoverErr != nil {
- log.Errorf("at exec panic, recoverErr:%v", recoverErr)
- if tx != nil {
- rollbackErr := tx.Rollback()
- if rollbackErr != nil {
- log.Errorf("conn at rollback error:%v",
rollbackErr)
+ defer func() {
+ recoverErr := recover()
+ if recoverErr != nil {
+ log.Errorf("at exec panic, recoverErr:%v",
recoverErr)
+ if tx != nil {
+ rollbackErr := tx.Rollback()
+ if rollbackErr != nil {
+ log.Errorf("conn at rollback
error:%v", rollbackErr)
+ }
}
}
- }
- }()
+ }()
+ }
ret, err := f()
if err != nil {
return nil, err
}
+ // For ExecContext, commit the transaction if it was created
if tx != nil {
if err := tx.Commit(); err != nil {
return nil, err
@@ -160,3 +204,54 @@ func (c *ATConn) createNewTxOnExecIfNeed(ctx
context.Context, f func() (types.Ex
return ret, nil
}
+
+// createTxAndQueryIfNeeded creates a transaction for query context and wraps
the rows to commit on close
+func (c *ATConn) createTxAndQueryIfNeeded(ctx context.Context, f func()
(types.ExecResult, error)) (types.ExecResult, error) {
+ var (
+ tx driver.Tx
+ err error
+ )
+
+ if c.txCtx.TransactionMode != types.Local && tm.IsGlobalTx(ctx) &&
c.autoCommit {
+ tx, err = c.BeginTx(ctx, driver.TxOptions{Isolation:
driver.IsolationLevel(gosql.LevelDefault)})
+ if err != nil {
+ return nil, err
+ }
+ defer func() {
+ recoverErr := recover()
+ if recoverErr != nil {
+ log.Errorf("at exec panic, recoverErr:%v",
recoverErr)
+ if tx != nil {
+ rollbackErr := tx.Rollback()
+ if rollbackErr != nil {
+ log.Errorf("conn at rollback
error:%v", rollbackErr)
+ }
+ }
+ }
+ }()
+ }
+
+ ret, err := f()
+ if err != nil {
+ return nil, err
+ }
+
+ // For QueryContext, wrap rows to commit on close
+ var activeTx driver.Tx
+ if c.txCtx.LocalTx != nil {
+ activeTx = c.txCtx.LocalTx
+ } else if tx != nil {
+ activeTx = tx
+ }
+
+ if activeTx != nil {
+ if rows, ok := ret.(types.ExecResult); ok {
+ if dr := rows.GetRows(); dr != nil {
+ wrappedRows := &RowsCommitOnClose{rows: dr, tx:
activeTx}
+ return
types.NewResult(types.WithRows(wrappedRows)), nil
+ }
+ }
+ }
+
+ return ret, nil
+}
diff --git a/pkg/datasource/sql/exec/select_for_update_executor.go
b/pkg/datasource/sql/exec/select_for_update_executor.go
index 30f4086c..fd82b2fa 100644
--- a/pkg/datasource/sql/exec/select_for_update_executor.go
+++ b/pkg/datasource/sql/exec/select_for_update_executor.go
@@ -268,8 +268,8 @@ func (s SelectForUpdateExecutor) ExecWithValue(ctx
context.Context, execCtx *typ
}
if originalAutoCommit {
- if err = tx.Commit(); err != nil {
- return nil, err
+ if tx != nil {
+ execCtx.TxCtx.LocalTx = tx
}
execCtx.IsAutoCommit = true
}
diff --git a/pkg/datasource/sql/rows_commit_on_close.go
b/pkg/datasource/sql/rows_commit_on_close.go
new file mode 100644
index 00000000..0e20b78b
--- /dev/null
+++ b/pkg/datasource/sql/rows_commit_on_close.go
@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package sql
+
+import (
+ "database/sql/driver"
+ "errors"
+ "io"
+ "sync/atomic"
+
+ "seata.apache.org/seata-go/v2/pkg/util/log"
+)
+
+type closeState int32
+
+const (
+ stateOpen closeState = iota
+ stateClosing
+ stateClosed
+)
+
+// RowsCommitOnClose wraps driver.Rows and commits the transaction
+type RowsCommitOnClose struct {
+ rows driver.Rows
+ tx driver.Tx
+
+ state int32 // atomic closeState
+ err error
+}
+
+// Close implements driver.Rows.Close.
+func (r *RowsCommitOnClose) Close() error {
+ if !atomic.CompareAndSwapInt32(&r.state, int32(stateOpen),
int32(stateClosing)) {
+ return r.err
+ }
+
+ var rowErr error
+ if r.rows != nil {
+ rowErr = r.rows.Close()
+ }
+
+ var txErr error
+ if r.tx != nil {
+ txErr = r.tx.Commit()
+ if txErr != nil {
+ log.Errorf("RowsCommitOnClose: commit failed: %v",
txErr)
+ }
+ r.tx = nil
+ }
+
+ var errs []error
+ if rowErr != nil {
+ errs = append(errs, rowErr)
+ }
+ if txErr != nil {
+ errs = append(errs, txErr)
+ }
+ r.err = errors.Join(errs...)
+
+ atomic.StoreInt32(&r.state, int32(stateClosed))
+ return r.err
+}
+
+// Columns implements driver.Rows.Columns.
+func (r *RowsCommitOnClose) Columns() []string {
+ return r.rows.Columns()
+}
+
+// Next implements driver.Rows.Next.
+func (r *RowsCommitOnClose) Next(dest []driver.Value) error {
+ err := r.rows.Next(dest)
+ if err == io.EOF {
+ _ = r.Close()
+ }
+ return err
+}
diff --git a/pkg/datasource/sql/rows_commit_on_close_test.go
b/pkg/datasource/sql/rows_commit_on_close_test.go
new file mode 100644
index 00000000..e49f5995
--- /dev/null
+++ b/pkg/datasource/sql/rows_commit_on_close_test.go
@@ -0,0 +1,265 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package sql
+
+import (
+ "database/sql/driver"
+ "errors"
+ "io"
+ "testing"
+
+ "github.com/golang/mock/gomock"
+ "github.com/stretchr/testify/assert"
+
+ "seata.apache.org/seata-go/v2/pkg/datasource/sql/mock"
+)
+
+/*
+ * ------------------------------------------------------------------------
+ * Close()
+ * ------------------------------------------------------------------------
+ */
+
+func TestRowsCommitOnClose_Close(t *testing.T) {
+ tests := []struct {
+ name string
+ setup func(ctrl *gomock.Controller) *RowsCommitOnClose
+ wantErr error
+ wantState int32
+ }{
+ {
+ name: "rows and tx success",
+ setup: func(ctrl *gomock.Controller) *RowsCommitOnClose
{
+ rows := mock.NewMockTestDriverRows(ctrl)
+ tx := mock.NewMockTestDriverTx(ctrl)
+ rows.EXPECT().Close().Return(nil)
+ tx.EXPECT().Commit().Return(nil)
+
+ return &RowsCommitOnClose{
+ rows: rows,
+ tx: tx,
+ state: int32(stateOpen),
+ }
+ },
+ wantErr: nil,
+ wantState: int32(stateClosed),
+ },
+ {
+ name: "rows close error only",
+ setup: func(ctrl *gomock.Controller) *RowsCommitOnClose
{
+ rows := mock.NewMockTestDriverRows(ctrl)
+ tx := mock.NewMockTestDriverTx(ctrl)
+ rowErr := errors.New("close failed")
+
+ rows.EXPECT().Close().Return(rowErr)
+ tx.EXPECT().Commit().Return(nil)
+
+ return &RowsCommitOnClose{
+ rows: rows,
+ tx: tx,
+ state: int32(stateOpen),
+ }
+ },
+ wantErr: errors.New("close failed"),
+ wantState: int32(stateClosed),
+ },
+ {
+ name: "tx commit error only",
+ setup: func(ctrl *gomock.Controller) *RowsCommitOnClose
{
+ rows := mock.NewMockTestDriverRows(ctrl)
+ tx := mock.NewMockTestDriverTx(ctrl)
+ txErr := errors.New("commit failed")
+
+ rows.EXPECT().Close().Return(nil)
+ tx.EXPECT().Commit().Return(txErr)
+
+ return &RowsCommitOnClose{
+ rows: rows,
+ tx: tx,
+ state: int32(stateOpen),
+ }
+ },
+ wantErr: errors.New("commit failed"),
+ wantState: int32(stateClosed),
+ },
+ {
+ name: "rows and tx both error",
+ setup: func(ctrl *gomock.Controller) *RowsCommitOnClose
{
+ rows := mock.NewMockTestDriverRows(ctrl)
+ tx := mock.NewMockTestDriverTx(ctrl)
+ rowErr := errors.New("close failed")
+ txErr := errors.New("commit failed")
+
+ rows.EXPECT().Close().Return(rowErr)
+ tx.EXPECT().Commit().Return(txErr)
+
+ return &RowsCommitOnClose{
+ rows: rows,
+ tx: tx,
+ state: int32(stateOpen),
+ }
+ },
+ wantErr: errors.New("close failed\ncommit failed"),
+ wantState: int32(stateClosed),
+ },
+ {
+ name: "already closed",
+ setup: func(ctrl *gomock.Controller) *RowsCommitOnClose
{
+ return &RowsCommitOnClose{
+ state: int32(stateClosed),
+ }
+ },
+ wantErr: nil,
+ wantState: int32(stateClosed),
+ },
+ {
+ name: "closing in progress",
+ setup: func(ctrl *gomock.Controller) *RowsCommitOnClose
{
+ return &RowsCommitOnClose{
+ state: int32(stateClosing),
+ }
+ },
+ wantErr: nil,
+ wantState: int32(stateClosing),
+ },
+ {
+ name: "nil rows",
+ setup: func(ctrl *gomock.Controller) *RowsCommitOnClose
{
+ tx := mock.NewMockTestDriverTx(ctrl)
+ tx.EXPECT().Commit().Return(nil)
+ return &RowsCommitOnClose{
+ tx: tx,
+ state: int32(stateOpen),
+ }
+ },
+ wantErr: nil,
+ wantState: int32(stateClosed),
+ },
+ {
+ name: "nil tx",
+ setup: func(ctrl *gomock.Controller) *RowsCommitOnClose
{
+ rows := mock.NewMockTestDriverRows(ctrl)
+ rows.EXPECT().Close().Return(nil)
+ return &RowsCommitOnClose{
+ rows: rows,
+ state: int32(stateOpen),
+ }
+ },
+ wantErr: nil,
+ wantState: int32(stateClosed),
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ ctrl := gomock.NewController(t)
+ defer ctrl.Finish()
+
+ r := tt.setup(ctrl)
+ err := r.Close()
+
+ if tt.wantErr == nil {
+ assert.NoError(t, err)
+ } else {
+ assert.Error(t, err)
+ assert.Equal(t, tt.wantErr.Error(), err.Error())
+ }
+
+ assert.Equal(t, tt.wantState, r.state)
+ if r.state == int32(stateClosed) {
+ assert.Nil(t, r.tx)
+ }
+ })
+ }
+}
+
+/*
+ * ------------------------------------------------------------------------
+ * Columns()
+ * ------------------------------------------------------------------------
+ */
+
+func TestRowsCommitOnClose_Columns_WithRows(t *testing.T) {
+ ctrl := gomock.NewController(t)
+ defer ctrl.Finish()
+
+ rowsMock := mock.NewMockTestDriverRows(ctrl)
+ rowsMock.EXPECT().Columns().Return([]string{"id", "name"})
+
+ r := &RowsCommitOnClose{rows: rowsMock}
+ cols := r.Columns()
+
+ assert.Equal(t, []string{"id", "name"}, cols)
+}
+
+func TestRowsCommitOnClose_Columns_NilRows(t *testing.T) {
+ r := &RowsCommitOnClose{}
+ assert.Panics(t, func() {
+ _ = r.Columns()
+ })
+}
+
+/*
+ * ------------------------------------------------------------------------
+ * Next()
+ * ------------------------------------------------------------------------
+ */
+
+func TestRowsCommitOnClose_Next_Error(t *testing.T) {
+ ctrl := gomock.NewController(t)
+ defer ctrl.Finish()
+
+ rowsMock := mock.NewMockTestDriverRows(ctrl)
+ dest := make([]driver.Value, 1)
+ rowsMock.EXPECT().Next(dest).Return(errors.New("next error"))
+
+ r := &RowsCommitOnClose{rows: rowsMock}
+ err := r.Next(dest)
+
+ assert.EqualError(t, err, "next error")
+}
+
+func TestRowsCommitOnClose_Next_NilRows(t *testing.T) {
+ r := &RowsCommitOnClose{}
+ assert.Panics(t, func() {
+ _ = r.Next(make([]driver.Value, 1))
+ })
+}
+
+func TestRowsCommitOnClose_Next_EOF_Close(t *testing.T) {
+ ctrl := gomock.NewController(t)
+ defer ctrl.Finish()
+
+ rowsMock := mock.NewMockTestDriverRows(ctrl)
+ txMock := mock.NewMockTestDriverTx(ctrl)
+
+ rowsMock.EXPECT().Next(gomock.Any()).Return(io.EOF)
+ rowsMock.EXPECT().Close().Return(nil)
+ txMock.EXPECT().Commit().Return(nil)
+
+ r := &RowsCommitOnClose{
+ rows: rowsMock,
+ tx: txMock,
+ state: int32(stateOpen),
+ }
+
+ err := r.Next(make([]driver.Value, 1))
+ assert.Equal(t, io.EOF, err)
+ assert.Equal(t, int32(stateClosed), r.state)
+ assert.Nil(t, r.tx)
+}
diff --git a/pkg/datasource/sql/types/types.go
b/pkg/datasource/sql/types/types.go
index 127bb925..d85cea32 100644
--- a/pkg/datasource/sql/types/types.go
+++ b/pkg/datasource/sql/types/types.go
@@ -143,6 +143,8 @@ type TransactionContext struct {
GlobalLockRequire bool
// RoundImages when run in AT mode, record before and after Row image
RoundImages *RoundRecordImage
+ // LocalTx local transaction instance, managed by at connection
+ LocalTx driver.Tx
}
// ExecContext
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
