fix call tx

Author: asmyasnikov

internal/xsql/conn.go

@@ -54,6 +54,8 @@ type conn struct {
 
 	dataOpts []options.ExecuteDataQueryOption
 	scanOpts []options.ExecuteScanQueryOption
+
+	currentTx *tx
 }
 
 var (
@@ -134,17 +136,22 @@ func (c *conn) BeginTx(ctx context.Context, txOptions driver.TxOptions) (driver.
 	if err != nil {
 		return nil, c.checkClosed(err)
 	}
-	return &tx{
+	c.currentTx = &tx{
 		conn:        c,
 		transaction: t,
-	}, nil
+	}
+	return c.currentTx, nil
 }
 
 func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
 	if c.isClosed() {
 		return nil, errClosedConn
 	}
-	switch m := queryModeFromContext(ctx, c.defaultQueryMode); m {
+	m := queryModeFromContext(ctx, c.defaultQueryMode)
+	if c.currentTx != nil && m == DataQueryMode {
+		return c.currentTx.ExecContext(ctx, query, args)
+	}
+	switch m {
 	case DataQueryMode:
 		_, res, err := c.session.Execute(ctx, txControl(ctx, c.defaultTxControl), query, toQueryParams(args))
 		if err != nil {
@@ -175,7 +182,11 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 	if c.isClosed() {
 		return nil, errClosedConn
 	}
-	switch m := queryModeFromContext(ctx, c.defaultQueryMode); m {
+	m := queryModeFromContext(ctx, c.defaultQueryMode)
+	if c.currentTx != nil && m == DataQueryMode {
+		return c.currentTx.QueryContext(ctx, query, args)
+	}
+	switch m {
 	case DataQueryMode:
 		_, res, err := c.session.Execute(ctx, txControl(ctx, c.defaultTxControl), query, toQueryParams(args))
 		if err != nil {

internal/xsql/tx.go

@@ -59,6 +59,9 @@ func (tx *tx) Commit() (err error) {
 	if tx.conn.isClosed() {
 		return errClosedConn
 	}
+	defer func() {
+		tx.conn.currentTx = nil
+	}()
 	_, err = tx.transaction.CommitTx(context.Background())
 	if err != nil {
 		return tx.conn.checkClosed(err)
@@ -70,6 +73,9 @@ func (tx *tx) Rollback() (err error) {
 	if tx.conn.isClosed() {
 		return errClosedConn
 	}
+	defer func() {
+		tx.conn.currentTx = nil
+	}()
 	err = tx.transaction.Rollback(context.Background())
 	if err != nil {
 		return tx.conn.checkClosed(err)

sql_e2e_test.go

@@ -123,8 +123,13 @@ func TestDatabaseSql(t *testing.T) {
 		if err != nil {
 			return fmt.Errorf("cannot upsert views: %w", err)
 		}
-		row = db.QueryRowContext(
-			ydb.WithQueryMode(ctx, ydb.ScanQueryMode),
+		return nil
+	}, retry.Idempotent(true))
+	if err != nil {
+		t.Fatalf("begin tx failed: %v\n", err)
+	}
+	err = retry.DoTx(ctx, db, func(ctx context.Context, tx *sql.Tx) error {
+		row := tx.QueryRowContext(ctx,
 			render(
 				querySelect,
 				templateConfig{
@@ -135,6 +140,7 @@ func TestDatabaseSql(t *testing.T) {
 			sql.Named("seasonID", uint64(1)),
 			sql.Named("episodeID", uint64(1)),
 		)
+		var views sql.NullFloat64
 		if err = row.Scan(&views); err != nil {
 			return fmt.Errorf("cannot select current views: %w", err)
 		}