Merge pull request #1365 from ydb-platform/fix-query-tx-commit-flag

Fixed respect to query.WithCommit() flag for tx.Execute in query service

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Fixed `query.WithCommit()` flag behaviour for `tx.Execute` in query service
 * OAuth 2.0 token exchange: allowed multiple resource parameters in according to https://www.rfc-editor.org/rfc/rfc8693
 
 ## v3.76.0

internal/query/options/execute.go

@@ -167,14 +167,28 @@ func (s *commonExecuteSettings) Params() *params.Parameters {
 
 func TxExecuteSettings(id string, opts ...TxExecuteOption) (settings *txExecuteSettings) {
 	settings = &txExecuteSettings{
-		ExecuteSettings: ExecuteSettings(WithTxControl(tx.NewControl(tx.WithTxID(id)))),
+		ExecuteSettings: ExecuteSettings(),
 	}
 	for _, opt := range opts {
 		if opt != nil {
 			opt.applyTxExecuteOption(settings)
 		}
 	}
 
+	var txControlOptions []tx.ControlOption
+	if settings.commitTx {
+		txControlOptions = []tx.ControlOption{
+			tx.WithTxID(id),
+			tx.CommitTx(),
+		}
+	} else {
+		txControlOptions = []tx.ControlOption{
+			tx.WithTxID(id),
+		}
+	}
+
+	settings.ExecuteSettings.SetTxControl(tx.NewControl(txControlOptions...))
+
 	return settings
 }
 

internal/query/transaction_test.go

@@ -218,6 +218,19 @@ func TestTxExecuteSettings(t *testing.T) {
 				params:    params.Builder{}.Param("$a").Text("A").Build(),
 			},
 		},
+		{
+			name: "WithCommitTx",
+			txOpts: []options.TxExecuteOption{
+				options.WithCommit(),
+			},
+			settings: testExecuteSettings{
+				execMode:  options.ExecModeExecute,
+				statsMode: options.StatsModeNone,
+				txControl: query.TxControl(query.WithTxID(""), query.CommitTx()),
+				syntax:    options.SyntaxYQL,
+				params:    nil,
+			},
+		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
 			a := allocator.New()

tests/integration/helpers_test.go

@@ -86,6 +86,20 @@ func (scope *scopeT) Driver(opts ...ydb.Option) *ydb.Driver {
 	return scope.driverNamed("default", opts...)
 }
 
+func (scope *scopeT) DriverWithLogs(opts ...ydb.Option) *ydb.Driver {
+	return scope.driverNamed("logged",
+		ydb.WithTraceQuery(
+			log.Query(
+				log.Default(os.Stdout,
+					log.WithLogQuery(),
+					log.WithMinLevel(log.INFO),
+				),
+				trace.QueryEvents,
+			),
+		),
+	)
+}
+
 func (scope *scopeT) DriverWithGRPCLogging() *ydb.Driver {
 	return scope.driverNamed("grpc-logged", ydb.With(config.WithGrpcOptions(
 		grpc.WithChainUnaryInterceptor(xtest.NewGrpcLogger(scope.t).UnaryClientInterceptor),

tests/integration/query_tx_execute_test.go

@@ -5,46 +5,28 @@ package integration
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 
-	"github.com/ydb-platform/ydb-go-sdk/v3"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/version"
-	"github.com/ydb-platform/ydb-go-sdk/v3/log"
 	"github.com/ydb-platform/ydb-go-sdk/v3/query"
-	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
 func TestQueryTxExecute(t *testing.T) {
 	if version.Lt(os.Getenv("YDB_VERSION"), "24.1") {
 		t.Skip("query service not allowed in YDB version '" + os.Getenv("YDB_VERSION") + "'")
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+	scope := newScope(t)
 
-	db, err := ydb.Open(ctx,
-		os.Getenv("YDB_CONNECTION_STRING"),
-		ydb.WithAccessTokenCredentials(os.Getenv("YDB_ACCESS_TOKEN_CREDENTIALS")),
-		ydb.WithTraceQuery(
-			log.Query(
-				log.Default(os.Stdout,
-					log.WithLogQuery(),
-					log.WithColoring(),
-					log.WithMinLevel(log.INFO),
-				),
-				trace.QueryEvents,
-			),
-		),
-	)
-	require.NoError(t, err)
 	var (
 		columnNames []string
 		columnTypes []string
 	)
-	err = db.Query().DoTx(ctx, func(ctx context.Context, tx query.TxActor) (err error) {
+	err := scope.DriverWithLogs().Query().DoTx(scope.Ctx, func(ctx context.Context, tx query.TxActor) (err error) {
 		res, err := tx.Execute(ctx, "SELECT 1 AS col1")
 		if err != nil {
 			return err
@@ -72,3 +54,41 @@ func TestQueryTxExecute(t *testing.T) {
 	require.Equal(t, []string{"col1"}, columnNames)
 	require.Equal(t, []string{"Int32"}, columnTypes)
 }
+
+func TestQueryWithCommitTxFlag(t *testing.T) {
+	if version.Lt(os.Getenv("YDB_VERSION"), "24.1") {
+		t.Skip("query service not allowed in YDB version '" + os.Getenv("YDB_VERSION") + "'")
+	}
+
+	scope := newScope(t)
+	var count uint64
+	err := scope.DriverWithLogs().Query().Do(scope.Ctx, func(ctx context.Context, s query.Session) error {
+		tableName := scope.TablePath()
+		tx, err := s.Begin(ctx, query.TxSettings(query.WithDefaultTxMode()))
+		if err != nil {
+			return fmt.Errorf("failed start transaction: %w", err)
+		}
+		q := fmt.Sprintf("UPSERT INTO `%v` (id, val) VALUES(1, \"2\")", tableName)
+		res, err := tx.Execute(ctx, q, query.WithCommit())
+		if err != nil {
+			return fmt.Errorf("failed execute insert: %w", err)
+		}
+		if err = res.Close(ctx); err != nil {
+			return err
+		}
+
+		// read row within other (implicit) transaction
+		q2 := fmt.Sprintf("SELECT COUNT(*) FROM `%v`", tableName)
+		row, err := s.ReadRow(ctx, q2)
+		if err != nil {
+			return fmt.Errorf("failed read row: %w", err)
+		}
+
+		if err = row.Scan(&count); err != nil {
+			return fmt.Errorf("failed scan row: %w", err)
+		}
+		return nil
+	})
+	require.NoError(t, err)
+	require.Equal(t, uint64(1), count)
+}