add regression test

Author: rekby

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed drop session from pool unnecessary in query service 
+
 ## v3.95.6
 * Fixed panic on span reporting in `xsql/Tx`
 

internal/query/session.go

@@ -39,7 +39,7 @@ func (s *Session) QueryResultSet(
 
 	r, err := execute(ctx, s.ID(), s.client, q, options.ExecuteSettings(opts...), withTrace(s.trace))
 	if err != nil {
-		s.handleSessionErrorStatus(err)
+		s.setStatusFromError(err)
 
 		return nil, xerrors.WithStackTrace(err)
 	}
@@ -57,7 +57,7 @@ func (s *Session) queryRow(
 ) (row query.Row, finalErr error) {
 	r, err := execute(ctx, s.ID(), s.client, q, settings, resultOpts...)
 	if err != nil {
-		s.handleSessionErrorStatus(err)
+		s.setStatusFromError(err)
 
 		return nil, xerrors.WithStackTrace(err)
 	}
@@ -125,7 +125,7 @@ func (s *Session) Begin(
 
 	txID, err := begin(ctx, s.client, s.ID(), txSettings)
 	if err != nil {
-		s.handleSessionErrorStatus(err)
+		s.setStatusFromError(err)
 
 		return nil, xerrors.WithStackTrace(err)
 	}
@@ -147,7 +147,7 @@ func (s *Session) Exec(
 
 	r, err := execute(ctx, s.ID(), s.client, q, options.ExecuteSettings(opts...), withTrace(s.trace))
 	if err != nil {
-		s.handleSessionErrorStatus(err)
+		s.setStatusFromError(err)
 
 		return xerrors.WithStackTrace(err)
 	}
@@ -171,17 +171,18 @@ func (s *Session) Query(
 
 	r, err := execute(ctx, s.ID(), s.client, q, options.ExecuteSettings(opts...), withTrace(s.trace))
 	if err != nil {
-		s.handleSessionErrorStatus(err)
+		s.setStatusFromError(err)
 
 		return nil, xerrors.WithStackTrace(err)
 	}
 
 	return r, nil
 }
 
-func (s *Session) handleSessionErrorStatus(err error) {
+func (s *Session) setStatusFromError(err error) {
 	switch {
-	case xerrors.IsTransportError(err) || xerrors.IsOperationError(err, Ydb.StatusIds_SESSION_BUSY, Ydb.StatusIds_BAD_SESSION):
+	case xerrors.IsTransportError(err) ||
+		xerrors.IsOperationError(err, Ydb.StatusIds_SESSION_BUSY, Ydb.StatusIds_BAD_SESSION):
 		s.SetStatus(session.StatusError)
 	case xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION):
 		s.SetStatus(session.StatusClosed)

internal/query/transaction.go

@@ -336,7 +336,8 @@ func (tx *Transaction) CommitTx(ctx context.Context) (finalErr error) {
 
 	err = commitTx(ctx, tx.s.client, tx.s.ID(), tx.ID())
 	if err != nil {
-		tx.s.handleSessionErrorStatus(err)
+		tx.s.setStatusFromError(err)
+
 		return xerrors.WithStackTrace(err)
 	}
 
@@ -371,7 +372,7 @@ func (tx *Transaction) Rollback(ctx context.Context) (finalErr error) {
 
 	err := rollback(ctx, tx.s.client, tx.s.ID(), tx.ID())
 	if err != nil {
-		tx.s.handleSessionErrorStatus(err)
+		tx.s.setStatusFromError(err)
 
 		return xerrors.WithStackTrace(err)
 	}

tests/integration/query_regression_test.go

@@ -376,3 +376,22 @@ func TestReadTwoPartsIntoMemoryIssue1559(t *testing.T) {
 	require.Equal(t, targetCount, len(rows))
 	require.Greater(t, partReaded, 1)
 }
+
+// https://github.com/ydb-platform/ydb-go-sdk/issues/1607
+func TestCloseSessionOnCustomerErrorsIssue1607(t *testing.T) {
+	scope := newScope(t)
+
+	sessionID1 := ""
+	_ = scope.Driver().Query().Do(scope.Ctx, func(ctx context.Context, s query.Session) error {
+		sessionID1 = s.ID()
+		return errors.New("test")
+	})
+
+	sessionID2 := ""
+	_ = scope.Driver().Query().Do(scope.Ctx, func(ctx context.Context, s query.Session) error {
+		sessionID2 = s.ID()
+		return nil
+	})
+
+	scope.Require.Equal(sessionID1, sessionID2)
+}