handle session error on server response instead of callback error

Author: rekby

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)

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"
@@ -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.handleSessionErrorStatus(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.handleSessionErrorStatus(err)
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -120,6 +125,8 @@ func (s *Session) Begin(
 
 	txID, err := begin(ctx, s.client, s.ID(), txSettings)
 	if err != nil {
+		s.handleSessionErrorStatus(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.handleSessionErrorStatus(err)
+
 		return xerrors.WithStackTrace(err)
 	}
 
@@ -162,8 +171,19 @@ 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)
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 
 	return r, nil
 }
+
+func (s *Session) handleSessionErrorStatus(err error) {
+	switch {
+	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

@@ -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"
@@ -338,10 +336,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.handleSessionErrorStatus(err)
 		return xerrors.WithStackTrace(err)
 	}
 
@@ -376,9 +371,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.handleSessionErrorStatus(err)
 
 		return xerrors.WithStackTrace(err)
 	}