Merge pull request #337 from ydb-platform/aborted

added test for Operation error ABORTED (Transaction locks invalidated)

Author: asmyasnikov

internal/xsql/conn.go

@@ -153,7 +153,12 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 		return nil, errClosedConn
 	}
 	m := queryModeFromContext(ctx, c.defaultQueryMode)
-	if c.currentTx != nil && m == DataQueryMode {
+	if c.currentTx != nil {
+		if m != DataQueryMode {
+			return nil, xerrors.WithStackTrace(
+				fmt.Errorf("query mode `%s` not supported with transaction", m.String()),
+			)
+		}
 		return c.currentTx.ExecContext(ctx, query, args)
 	}
 	switch m {
@@ -188,7 +193,12 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 		return nil, errClosedConn
 	}
 	m := queryModeFromContext(ctx, c.defaultQueryMode)
-	if c.currentTx != nil && m == DataQueryMode {
+	if c.currentTx != nil {
+		if m != DataQueryMode {
+			return nil, xerrors.WithStackTrace(
+				fmt.Errorf("query mode `%s` not supported with transaction", m.String()),
+			)
+		}
 		return c.currentTx.QueryContext(ctx, query, args)
 	}
 	switch m {

retry/sql.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"fmt"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/isolation"
@@ -51,19 +52,23 @@ func WithTxSettings(txControl *table.TransactionSettings) DoTxOption {
 
 // DoTx is a shortcut for calling Do(ctx, f) on initialized TxDoer with DB field set to given db.
 func DoTx(ctx context.Context, db *sql.DB, f TxOperationFunc, opts ...DoTxOption) error {
-	options := doTxOptions{
-		txOptions: &sql.TxOptions{
-			Isolation: sql.LevelDefault,
-			ReadOnly:  false,
-		},
-		idempotent: false,
-	}
+	var (
+		options = doTxOptions{
+			txOptions: &sql.TxOptions{
+				Isolation: sql.LevelDefault,
+				ReadOnly:  false,
+			},
+			idempotent: false,
+		}
+		attempts = 0
+	)
 	for _, o := range opts {
 		if err := o(&options); err != nil {
 			return xerrors.WithStackTrace(err)
 		}
 	}
 	err := Retry(ctx, func(ctx context.Context) (err error) {
+		attempts++
 		tx, err := db.BeginTx(ctx, options.txOptions)
 		if err != nil {
 			return xerrors.WithStackTrace(err)
@@ -85,7 +90,9 @@ func DoTx(ctx context.Context, db *sql.DB, f TxOperationFunc, opts ...DoTxOption
 		return nil
 	}, WithIdempotent(options.idempotent))
 	if err != nil {
-		return xerrors.WithStackTrace(err)
+		return xerrors.WithStackTrace(
+			fmt.Errorf("opration failed with %d attempts: %w", attempts, err),
+		)
 	}
 	return nil
 }

retry/sql_test.go

@@ -0,0 +1,76 @@
+package retry
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"testing"
+	"time"
+
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Issue"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+)
+
+type mock struct{}
+
+func (c mock) Commit() error {
+	return nil
+}
+
+func (c mock) Rollback() error {
+	return nil
+}
+
+func (c mock) Open(name string) (driver.Conn, error) {
+	return c, nil
+}
+
+func (c mock) Prepare(query string) (driver.Stmt, error) {
+	return nil, nil
+}
+
+func (c mock) Close() error {
+	return nil
+}
+
+func (c mock) Begin() (driver.Tx, error) {
+	return c, nil
+}
+
+func (c mock) Connect(ctx context.Context) (driver.Conn, error) {
+	return c, nil
+}
+
+func (c mock) Driver() driver.Driver {
+	return c
+}
+
+func (c mock) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	return c, nil
+}
+
+func TestDoTxAbortedTLI(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	db := sql.OpenDB(mock{})
+	var counter int
+	if err := DoTx(ctx, db, func(ctx context.Context, tx *sql.Tx) error {
+		counter++
+		if counter > 10 {
+			return nil
+		}
+		return xerrors.Operation(
+			xerrors.WithStatusCode(Ydb.StatusIds_ABORTED),
+			xerrors.WithIssues([]*Ydb_Issue.IssueMessage{
+				{IssueCode: 2001, Message: "Transaction locks invalidated"},
+			}),
+		)
+	}); err != nil {
+		t.Errorf("retry operation failed: %v", err)
+	}
+	if counter <= 1 {
+		t.Errorf("nothing attempts: %d", counter)
+	}
+}

sql_e2e_test.go

@@ -66,35 +66,42 @@ func TestDatabaseSql(t *testing.T) {
 		t.Fatalf("fill failed: %v\n", err)
 	}
 
-	err = retry.DoTx(ctx, db, func(ctx context.Context, tx *sql.Tx) error {
-		row := tx.QueryRowContext(
-			ydb.WithQueryMode(ctx, ydb.ExplainQueryMode),
-			render(
-				querySelect,
-				templateConfig{
-					TablePathPrefix: path.Join(cc.Name(), folder),
-				},
-			),
-			sql.Named("seriesID", uint64(1)),
-			sql.Named("seasonID", uint64(1)),
-			sql.Named("episodeID", uint64(1)),
-		)
-		var (
-			ast  string
-			plan string
-		)
-		if err = row.Scan(&ast, &plan); err != nil {
-			return fmt.Errorf("cannot explain: %w", err)
+	// getting explain of query
+	row := db.QueryRowContext(
+		ydb.WithQueryMode(ctx, ydb.ExplainQueryMode),
+		render(
+			querySelect,
+			templateConfig{
+				TablePathPrefix: path.Join(cc.Name(), folder),
+			},
+		),
+		sql.Named("seriesID", uint64(1)),
+		sql.Named("seasonID", uint64(1)),
+		sql.Named("episodeID", uint64(1)),
+	)
+	var (
+		ast  string
+		plan string
+	)
+	if err = row.Scan(&ast, &plan); err != nil {
+		t.Fatalf("cannot explain: %v", err)
+	}
+	t.Logf("ast = %v", ast)
+	t.Logf("plan = %v", plan)
+
+	err = retry.DoTx(ctx, db, func(ctx context.Context, tx *sql.Tx) (err error) {
+		var stmt *sql.Stmt
+		stmt, err = tx.PrepareContext(ctx, render(
+			querySelect,
+			templateConfig{
+				TablePathPrefix: path.Join(cc.Name(), folder),
+			},
+		))
+		if err != nil {
+			return fmt.Errorf("cannot prepare query: %w", err)
 		}
-		t.Logf("ast = %v", ast)
-		t.Logf("plan = %v", plan)
-		row = tx.QueryRowContext(ctx,
-			render(
-				querySelect,
-				templateConfig{
-					TablePathPrefix: path.Join(cc.Name(), folder),
-				},
-			),
+
+		row = stmt.QueryRowContext(ctx,
 			sql.Named("seriesID", uint64(1)),
 			sql.Named("seasonID", uint64(1)),
 			sql.Named("episodeID", uint64(1)),