Merge pull request #1106 from ydb-platform/lazy-tx

lazy transactions: query.WithCommit() option + tests

Author: asmyasnikov

internal/query/transaction.go

@@ -22,22 +22,10 @@ func (tx transaction) ID() string {
 	return tx.id
 }
 
-func fromTxOptions(txID string, txOpts ...query.TxExecuteOption) executeSettings {
-	opts := make([]query.ExecuteOption, 0, len(txOpts)+1)
-	for _, opt := range txOpts {
-		if executeOpt, has := opt.(query.ExecuteOption); has {
-			opts = append(opts, executeOpt)
-		}
-	}
-	opts = append(opts, query.WithTxControl(query.TxControl(query.WithTxID(txID))))
-
-	return query.ExecuteSettings(opts...)
-}
-
 func (tx transaction) Execute(ctx context.Context, q string, opts ...query.TxExecuteOption) (
 	r query.Result, err error,
 ) {
-	_, res, err := execute(ctx, tx.s, tx.s.queryClient, q, fromTxOptions(tx.id, opts...))
+	_, res, err := execute(ctx, tx.s, tx.s.queryClient, q, query.TxExecuteSettings(tx.id, opts...).ExecuteSettings)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}

internal/query/transaction_test.go

@@ -137,7 +137,7 @@ func (s testExecuteSettings) CallOptions() []grpc.CallOption {
 
 var _ executeSettings = testExecuteSettings{}
 
-func TestFromTxOptions(t *testing.T) {
+func TestTxExecuteSettings(t *testing.T) {
 	for _, tt := range []struct {
 		name     string
 		txID     string
@@ -224,7 +224,7 @@ func TestFromTxOptions(t *testing.T) {
 	} {
 		t.Run(tt.name, func(t *testing.T) {
 			a := allocator.New()
-			settings := fromTxOptions(tt.txID, tt.txOpts...)
+			settings := query.TxExecuteSettings(tt.txID, tt.txOpts...).ExecuteSettings
 			require.Equal(t, tt.settings.Syntax(), settings.Syntax())
 			require.Equal(t, tt.settings.ExecMode(), settings.ExecMode())
 			require.Equal(t, tt.settings.StatsMode(), settings.StatsMode())

query/execute_options.go

@@ -28,16 +28,23 @@ type (
 		applyExecuteOption(s *executeSettings)
 	}
 	txExecuteSettings struct {
-		commonExecuteSettings
+		ExecuteSettings *executeSettings
+
+		commitTx bool
 	}
 	TxExecuteOption interface {
 		applyTxExecuteOption(s *txExecuteSettings)
 	}
+	txCommitOption   struct{}
 	parametersOption params.Parameters
 )
 
+func (t txCommitOption) applyTxExecuteOption(s *txExecuteSettings) {
+	s.commitTx = true
+}
+
 func (syntax Syntax) applyTxExecuteOption(s *txExecuteSettings) {
-	s.syntax = syntax
+	syntax.applyExecuteOption(s.ExecuteSettings)
 }
 
 func (syntax Syntax) applyExecuteOption(s *executeSettings) {
@@ -50,7 +57,7 @@ const (
 )
 
 func (params parametersOption) applyTxExecuteOption(s *txExecuteSettings) {
-	s.params = append(s.params, params...)
+	params.applyExecuteOption(s.ExecuteSettings)
 }
 
 func (params parametersOption) applyExecuteOption(s *executeSettings) {
@@ -62,19 +69,19 @@ func (opts callOptions) applyExecuteOption(s *executeSettings) {
 }
 
 func (opts callOptions) applyTxExecuteOption(s *txExecuteSettings) {
-	s.callOptions = append(s.callOptions, opts...)
+	opts.applyExecuteOption(s.ExecuteSettings)
 }
 
 func (mode StatsMode) applyTxExecuteOption(s *txExecuteSettings) {
-	s.statsMode = mode
+	mode.applyExecuteOption(s.ExecuteSettings)
 }
 
 func (mode StatsMode) applyExecuteOption(s *executeSettings) {
 	s.statsMode = mode
 }
 
 func (mode ExecMode) applyTxExecuteOption(s *txExecuteSettings) {
-	s.execMode = mode
+	mode.applyExecuteOption(s.ExecuteSettings)
 }
 
 func (mode ExecMode) applyExecuteOption(s *executeSettings) {
@@ -144,9 +151,9 @@ func (s *commonExecuteSettings) Params() *params.Parameters {
 	return &s.params
 }
 
-func TxExecuteSettings(opts ...TxExecuteOption) (settings *txExecuteSettings) {
+func TxExecuteSettings(id string, opts ...TxExecuteOption) (settings *txExecuteSettings) {
 	settings = &txExecuteSettings{
-		commonExecuteSettings: defaultCommonExecuteSettings(),
+		ExecuteSettings: ExecuteSettings(WithTxControl(TxControl(WithTxID(id)))),
 	}
 	for _, opt := range opts {
 		opt.applyTxExecuteOption(settings)
@@ -168,8 +175,13 @@ var (
 	_ ExecuteOption   = StatsMode(0)
 	_ TxExecuteOption = ExecMode(0)
 	_ TxExecuteOption = StatsMode(0)
+	_ TxExecuteOption = txCommitOption{}
 )
 
+func WithCommit() txCommitOption {
+	return txCommitOption{}
+}
+
 func WithExecMode(mode ExecMode) ExecMode {
 	return mode
 }

tests/integration/query_execute_test.go

@@ -5,6 +5,7 @@ package integration
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"testing"
 	"time"
@@ -168,4 +169,89 @@ func TestQueryExecute(t *testing.T) {
 		require.EqualValues(t, time.Duration(100500000000), data.P3)
 		require.Nil(t, data.P4)
 	})
+	t.Run("Transaction", func(t *testing.T) {
+		t.Run("Explicit", func(t *testing.T) {
+			err = db.Query().Do(ctx, func(ctx context.Context, s query.Session) (err error) {
+				tx, err := s.Begin(ctx, query.TxSettings(query.WithSerializableReadWrite()))
+				if err != nil {
+					return err
+				}
+				res, err := tx.Execute(ctx, `SELECT 1`)
+				if err != nil {
+					return err
+				}
+				rs, err := res.NextResultSet(ctx)
+				if err != nil {
+					return err
+				}
+				row, err := rs.NextRow(ctx)
+				if err != nil {
+					return err
+				}
+				var v int32
+				err = row.Scan(&v)
+				if err != nil {
+					return err
+				}
+				if v != 1 {
+					return fmt.Errorf("unexpected value from database: %d", v)
+				}
+				if err = res.Err(); err != nil {
+					return err
+				}
+				return tx.CommitTx(ctx)
+			}, query.WithIdempotent())
+			require.NoError(t, err)
+		})
+		t.Run("Lazy", func(t *testing.T) {
+			err = db.Query().Do(ctx, func(ctx context.Context, s query.Session) (err error) {
+				tx, res, err := s.Execute(ctx, `SELECT 1`,
+					query.WithTxControl(query.TxControl(query.BeginTx(query.WithSerializableReadWrite()))),
+				)
+				if err != nil {
+					return err
+				}
+				rs, err := res.NextResultSet(ctx)
+				if err != nil {
+					return err
+				}
+				row, err := rs.NextRow(ctx)
+				if err != nil {
+					return err
+				}
+				var v int32
+				err = row.Scan(&v)
+				if err != nil {
+					return err
+				}
+				if v != 1 {
+					return fmt.Errorf("unexpected value from database: %d", v)
+				}
+				if err = res.Err(); err != nil {
+					return err
+				}
+				res, err = tx.Execute(ctx, `SELECT 2`, query.WithCommit())
+				if err != nil {
+					return err
+				}
+				rs, err = res.NextResultSet(ctx)
+				if err != nil {
+					return err
+				}
+				row, err = rs.NextRow(ctx)
+				if err != nil {
+					return err
+				}
+				err = row.Scan(&v)
+				if err != nil {
+					return err
+				}
+				if v != 2 {
+					return fmt.Errorf("unexpected value from database: %d", v)
+				}
+				return res.Err()
+			}, query.WithIdempotent())
+			require.NoError(t, err)
+		})
+	})
 }