some fixes && test

asmyasnikov · asmyasnikov · commit 26c8bfe080bf · 2025-02-28T18:48:14.000+03:00
diff --git a/internal/xsql/badconn/badconn.go b/internal/xsql/badconn/badconn.go
@@ -2,6 +2,7 @@ package badconn
 
 import (
 	"database/sql/driver"
+	"fmt"
 	"io"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
@@ -11,6 +12,10 @@ type Error struct {
 	err error
 }
 
+func Errorf(format string, args ...interface{}) error {
+	return &Error{err: fmt.Errorf(format, args...)}
+}
+
 func (e Error) Origin() error {
 	return e.err
 }
diff --git a/internal/xsql/badconn/badconn_test.go b/internal/xsql/badconn/badconn_test.go
diff --git a/internal/xsql/conn.go b/internal/xsql/conn.go
@@ -23,6 +23,14 @@ type Conn struct {
 	lastUsage xsync.LastUsage
 }
 
+func (c *Conn) ID() string {
+	return c.cc.ID()
+}
+
+func (c *Conn) NodeID() uint32 {
+	return c.cc.NodeID()
+}
+
 func (c *Conn) Ping(ctx context.Context) (finalErr error) {
 	onDone := trace.DatabaseSQLOnConnPing(c.connector.trace, &c.ctx,
 		stack.FunctionID("database/sql.(*Conn).Ping", stack.Package("database/sql")),
diff --git a/internal/xsql/tx.go b/internal/xsql/tx.go
@@ -6,8 +6,8 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/common"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
diff --git a/internal/xsql/unwrap.go b/internal/xsql/unwrap.go
@@ -5,7 +5,7 @@ import (
 	"fmt"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 )
 
 func Unwrap[T *sql.DB | *sql.Conn](v T) (connector *Connector, _ error) {
diff --git a/internal/xsql/xquery/tx_fake.go b/internal/xsql/xquery/tx_fake.go
@@ -7,8 +7,8 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/common"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"
 )
 
 type txFake struct {
diff --git a/internal/xsql/xtable/conn.go b/internal/xsql/xtable/conn.go
@@ -12,8 +12,8 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/common"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/scripting"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/options"
diff --git a/internal/xsql/xtable/rows.go b/internal/xsql/xtable/rows.go
@@ -10,7 +10,7 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/scanner"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/options"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/result"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/result/indexed"
diff --git a/internal/xsql/xtable/tx.go b/internal/xsql/xtable/tx.go
@@ -7,8 +7,8 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/common"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
 )
 
diff --git a/internal/xsql/xtable/tx_fake.go b/internal/xsql/xtable/tx_fake.go
@@ -7,8 +7,8 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/common"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"
 )
 
 type txFake struct {
diff --git a/retry/errors.go b/retry/errors.go
@@ -2,7 +2,7 @@ package retry
 
 import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 )
 
 func unwrapErrBadConn(err error) error {
diff --git a/retry/sql.go b/retry/sql.go
@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"database/sql/driver"
 	"fmt"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
@@ -93,6 +94,12 @@ func DoWithResult[T any](ctx context.Context, db *sql.DB,
 			return zeroValue, unwrapErrBadConn(xerrors.WithStackTrace(err))
 		}
 		defer func() {
+			if finalErr != nil && mustDeleteConn(finalErr, cc) {
+				_ = cc.Raw(func(driverConn any) error {
+					return xerrors.WithStackTrace(badconn.Errorf("close connection because: %w", finalErr))
+				})
+			}
+
 			_ = cc.Close()
 		}()
 		v, err := op(xcontext.MarkRetryCall(ctx), cc)
@@ -233,10 +240,12 @@ func DoTxWithResult[T any](ctx context.Context, db *sql.DB,
 	return v, nil
 }
 
-func mustDeleteConn(err error, conn *sql.Conn) bool {
+func mustDeleteConn[T interface {
+	*sql.Conn
+}](err error, conn T) bool {
 	if xerrors.Is(err, driver.ErrBadConn) {
 		return true
 	}
 
-	return xerrors.IsValid(err, conn)
+	return !xerrors.IsValid(err, conn)
 }
diff --git a/retry/sql_test.go b/retry/sql_test.go
@@ -14,7 +14,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/backoff"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 )
 
 type mockConnector struct {
diff --git a/tests/integration/database_sql_regression_test.go b/tests/integration/database_sql_regression_test.go
@@ -9,6 +9,7 @@ import (
 	"database/sql/driver"
 	"errors"
 	"fmt"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 	"math/rand"
 	"strings"
 	"testing"
@@ -21,7 +22,6 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
diff --git a/tests/integration/database_sql_retryable_error_test.go b/tests/integration/database_sql_retryable_error_test.go
@@ -0,0 +1,83 @@
+//go:build integration
+// +build integration
+
+package integration
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
+	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
+)
+
+func TestDatabaseSqlRetryableError(t *testing.T) {
+	scope := newScope(t)
+
+	var (
+		newConnCounter   = 0
+		closeConnCounter = 0
+	)
+
+	db := scope.SQLDriverWithFolder(
+		ydb.WithTablePathPrefix(scope.Folder()),
+		ydb.WithAutoDeclare(),
+		ydb.WithNumericArgs(),
+		ydb.WithDatabaseSQLTrace(trace.DatabaseSQL{
+			OnConnectorConnect: func(info trace.DatabaseSQLConnectorConnectStartInfo) func(trace.DatabaseSQLConnectorConnectDoneInfo) {
+				return func(info trace.DatabaseSQLConnectorConnectDoneInfo) {
+					if info.Error == nil {
+						newConnCounter++
+					}
+				}
+			},
+			OnConnClose: func(info trace.DatabaseSQLConnCloseStartInfo) func(trace.DatabaseSQLConnCloseDoneInfo) {
+				closeConnCounter++
+
+				return nil
+			},
+		}),
+	)
+
+	var id *string
+
+	err := retry.Do(scope.Ctx, db, func(ctx context.Context, cc *sql.Conn) error {
+		var ccID *string
+		cc.Raw(func(driverConn any) error {
+			if ider, has := driverConn.(interface{ ID() string }); has {
+				v := ider.ID()
+				ccID = &v
+			}
+
+			return nil
+		})
+
+		require.NotNil(t, ccID)
+
+		if id != nil {
+			require.NotEqual(t, *ccID, *id)
+		}
+
+		id = ccID
+
+		if closeConnCounter > 0 {
+			return nil
+		}
+
+		return retry.RetryableError(errors.New("test error"), xerrors.Invalid(cc))
+	})
+
+	require.NoError(t, err)
+	require.Equal(t, 2, newConnCounter)
+	require.Equal(t, 1, closeConnCounter)
+
+	require.NoError(t, db.Close())
+	require.Equal(t, 2, newConnCounter)
+	require.Equal(t, 2, closeConnCounter)
+}

Original file line number	Diff line number	Diff line change
`@@ -6,8 +6,8 @@ import (`
`6`	`6`
`7`	`7`	`"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"`
`8`	`8`	`"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"`
	`9`	`+ "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"`
`9`	`10`	`"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/common"`
`10`		`- "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"`
`11`	`11`	`"github.com/ydb-platform/ydb-go-sdk/v3/trace"`
`12`	`12`	`)`
`13`	`13`
Original file line number	Diff line number	Diff line change
`@@ -5,7 +5,7 @@ import (`
`5`	`5`	`"fmt"`
`6`	`6`
`7`	`7`	`"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"`
`8`		`- "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"`
	`8`	`+ "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"`
`9`	`9`	`)`
`10`	`10`
`11`	`11`	`func Unwrap[T sql.DB \| sql.Conn](v T) (connector *Connector, _ error) {`
Original file line number	Diff line number	Diff line change
`@@ -7,8 +7,8 @@ import (`
`7`	`7`	`"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"`
`8`	`8`	`"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"`
`9`	`9`	`"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"`
	`10`	`+ "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"`
`10`	`11`	`"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/common"`
`11`		`- "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable/badconn"`
`12`	`12`	`)`
`13`	`13`
`14`	`14`	`type txFake struct {`