Merge pull request #1608 from ydb-platform/1607-bad-session-on-error

1607 bad session on error

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed drop session from pool unnecessary in query service 
+
 ## v3.96.0
 * Supported of list, set and struct for unmarshall using `sugar.Unmarshall...`
 

internal/query/client.go

@@ -212,8 +212,6 @@ func do(
 
 		err := op(ctx, s)
 		if err != nil {
-			s.SetStatus(session.StatusError)
-
 			return xerrors.WithStackTrace(err)
 		}
 
@@ -276,10 +274,6 @@ func doTx(
 
 		defer func() {
 			_ = tx.Rollback(ctx)
-
-			if opErr != nil {
-				s.SetStatus(session.StatusError)
-			}
 		}()
 
 		err = op(ctx, tx)
@@ -539,7 +533,7 @@ func CreateSession(ctx context.Context, c *Client) (*Session, error) {
 			return nil, xerrors.WithStackTrace(err)
 		}
 
-		s.laztTx = c.config.LazyTx()
+		s.lazyTx = c.config.LazyTx()
 
 		return s, nil
 	})
@@ -590,7 +584,7 @@ func New(ctx context.Context, cc grpc.ClientConnInterface, cfg *config.Config) *
 					return nil, xerrors.WithStackTrace(err)
 				}
 
-				s.laztTx = cfg.LazyTx()
+				s.lazyTx = cfg.LazyTx()
 
 				return s, nil
 			}),

internal/query/client_test.go

@@ -1580,7 +1580,7 @@ func newTestSessionWithClient(id string, client Ydb_Query_V1.QueryServiceClient,
 		Core:   &sessionControllerMock{id: id},
 		client: client,
 		trace:  &trace.Query{},
-		laztTx: lazyTx,
+		lazyTx: lazyTx,
 	}
 }
 

internal/query/session.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/result"
@@ -23,7 +24,7 @@ type (
 
 		client Ydb_Query_V1.QueryServiceClient
 		trace  *trace.Query
-		laztTx bool
+		lazyTx bool
 	}
 )
 
@@ -38,6 +39,8 @@ func (s *Session) QueryResultSet(
 
 	r, err := execute(ctx, s.ID(), s.client, q, options.ExecuteSettings(opts...), withTrace(s.trace))
 	if err != nil {
+		s.setStatusFromError(err)
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -54,6 +57,8 @@ func (s *Session) queryRow(
 ) (row query.Row, finalErr error) {
 	r, err := execute(ctx, s.ID(), s.client, q, settings, resultOpts...)
 	if err != nil {
+		s.setStatusFromError(err)
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -111,7 +116,7 @@ func (s *Session) Begin(
 		}
 	}()
 
-	if s.laztTx {
+	if s.lazyTx {
 		return &Transaction{
 			s:          s,
 			txSettings: txSettings,
@@ -120,6 +125,8 @@ func (s *Session) Begin(
 
 	txID, err := begin(ctx, s.client, s.ID(), txSettings)
 	if err != nil {
+		s.setStatusFromError(err)
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -140,6 +147,8 @@ func (s *Session) Exec(
 
 	r, err := execute(ctx, s.ID(), s.client, q, options.ExecuteSettings(opts...), withTrace(s.trace))
 	if err != nil {
+		s.setStatusFromError(err)
+
 		return xerrors.WithStackTrace(err)
 	}
 
@@ -162,8 +171,21 @@ func (s *Session) Query(
 
 	r, err := execute(ctx, s.ID(), s.client, q, options.ExecuteSettings(opts...), withTrace(s.trace))
 	if err != nil {
+		s.setStatusFromError(err)
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 
 	return r, nil
 }
+
+func (s *Session) setStatusFromError(err error) {
+	switch {
+	case xerrors.IsTransportError(err):
+		s.SetStatus(session.StatusError)
+	case 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

@@ -5,13 +5,11 @@ import (
 	"fmt"
 
 	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
-	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/result"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/session"
 	queryTx "github.com/ydb-platform/ydb-go-sdk/v3/internal/query/tx"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	baseTx "github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
@@ -116,6 +114,8 @@ func (tx *Transaction) QueryResultSet(
 	}
 	r, err := execute(ctx, tx.s.ID(), tx.s.client, q, settings, resultOpts...)
 	if err != nil {
+		tx.s.setStatusFromError(err)
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -165,6 +165,8 @@ func (tx *Transaction) QueryRow(
 	}
 	r, err := execute(ctx, tx.s.ID(), tx.s.client, q, settings, resultOpts...)
 	if err != nil {
+		tx.s.setStatusFromError(err)
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -231,6 +233,8 @@ func (tx *Transaction) Exec(ctx context.Context, q string, opts ...options.Execu
 
 	r, err := execute(ctx, tx.s.ID(), tx.s.client, q, settings, resultOpts...)
 	if err != nil {
+		tx.s.setStatusFromError(err)
+
 		return xerrors.WithStackTrace(err)
 	}
 
@@ -299,6 +303,8 @@ func (tx *Transaction) Query(ctx context.Context, q string, opts ...options.Exec
 	}
 	r, err := execute(ctx, tx.s.ID(), tx.s.client, q, settings, resultOpts...)
 	if err != nil {
+		tx.s.setStatusFromError(err)
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -338,9 +344,7 @@ func (tx *Transaction) CommitTx(ctx context.Context) (finalErr error) {
 
 	err = commitTx(ctx, tx.s.client, tx.s.ID(), tx.ID())
 	if err != nil {
-		if xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION) {
-			tx.s.SetStatus(session.StatusClosed)
-		}
+		tx.s.setStatusFromError(err)
 
 		return xerrors.WithStackTrace(err)
 	}
@@ -376,9 +380,7 @@ func (tx *Transaction) Rollback(ctx context.Context) (finalErr error) {
 
 	err := rollback(ctx, tx.s.client, tx.s.ID(), tx.ID())
 	if err != nil {
-		if xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION) {
-			tx.s.SetStatus(session.StatusClosed)
-		}
+		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)
+}