From 6538a323f59cec0bdc45291b16517adc8fbf5e86 Mon Sep 17 00:00:00 2001
From: Aleksey Myasnikov <asmyasnikov@ydb.tech>
Date: Tue, 5 Nov 2024 17:21:17 +0300
Subject: [PATCH] fix begin tx

---
 internal/table/conn/conn.go  | 18 +++++++++++-------
 internal/xerrors/tx.go       |  1 +
 tests/integration/tx_test.go |  1 +
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/internal/table/conn/conn.go b/internal/table/conn/conn.go
index f3490e701..bc6d5265e 100644
--- a/internal/table/conn/conn.go
+++ b/internal/table/conn/conn.go
@@ -448,7 +448,14 @@ func (c *Conn) ID() string {
 	return c.session.ID()
 }
 
-func (c *Conn) beginTx(ctx context.Context, txOptions driver.TxOptions) (currentTx, error) {
+func (c *Conn) beginTx(ctx context.Context, txOptions driver.TxOptions) (tx currentTx, finalErr error) {
+	onDone := trace.DatabaseSQLOnConnBegin(c.parent.Trace(), &ctx,
+		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/table/conn.(*Conn).beginTx"),
+	)
+	defer func() {
+		onDone(tx, finalErr)
+	}()
+
 	if c.currentTx != nil {
 		return nil, badconn.Map(
 			xerrors.WithStackTrace(xerrors.AlreadyHasTx(c.currentTx.ID())),
@@ -470,13 +477,10 @@ func (c *Conn) beginTx(ctx context.Context, txOptions driver.TxOptions) (current
 }
 
 func (c *Conn) BeginTx(ctx context.Context, txOptions driver.TxOptions) (driver.Tx, error) {
-	onDone := trace.DatabaseSQLOnConnBegin(c.parent.Trace(), &ctx,
-		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/table/conn.(*Conn).BeginTx"),
-	)
 	tx, err := c.beginTx(ctx, txOptions)
-	defer func() {
-		onDone(tx, err)
-	}()
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
+	}
 
 	c.currentTx = tx
 
diff --git a/internal/xerrors/tx.go b/internal/xerrors/tx.go
index 795b0a583..b4eb75fd6 100644
--- a/internal/xerrors/tx.go
+++ b/internal/xerrors/tx.go
@@ -25,5 +25,6 @@ func AlreadyHasTx(txID string) error {
 
 func IsAlreadyHasTx(err error) bool {
 	var txErr *alreadyHasTxError
+
 	return As(err, &txErr)
 }
diff --git a/tests/integration/tx_test.go b/tests/integration/tx_test.go
index bcf68f0c7..bcb16610a 100644
--- a/tests/integration/tx_test.go
+++ b/tests/integration/tx_test.go
@@ -143,6 +143,7 @@ func TestNoEffectsIfForgetCommitTx(t *testing.T) {
 
 		// second tx on existing conn === session
 		_, err = cc.BeginTx(ctx, &sql.TxOptions{})
+		require.Error(t, err)
 		require.True(t, xerrors.IsAlreadyHasTx(err))
 	})
 }