From 8ee1b6a6515bc02b5621cca241aecb65bba328a4 Mon Sep 17 00:00:00 2001
From: Aleksey Myasnikov <asmyasnikov@ydb.tech>
Date: Mon, 17 Jun 2024 14:46:40 +0300
Subject: [PATCH] wrapped error BAD_SESSION in query session calls

---
 internal/query/config/config.go |  1 +
 internal/query/session.go       | 19 +++++++++++++++++++
 internal/query/transaction.go   | 31 +++++++++++++++++++++++++++----
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/internal/query/config/config.go b/internal/query/config/config.go
index 8bd1262f6..ab1db6490 100644
--- a/internal/query/config/config.go
+++ b/internal/query/config/config.go
@@ -43,6 +43,7 @@ func defaults() *Config {
 		sessionCreateTimeout: DefaultSessionCreateTimeout,
 		sessionDeleteTimeout: DefaultSessionDeleteTimeout,
 		trace:                &trace.Query{},
+		useSessionPool:       true,
 	}
 }
 
diff --git a/internal/query/session.go b/internal/query/session.go
index 3ee888eda..9e6155d81 100644
--- a/internal/query/session.go
+++ b/internal/query/session.go
@@ -6,6 +6,7 @@ import (
 	"sync/atomic"
 
 	"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"
@@ -78,6 +79,16 @@ func createSession(ctx context.Context, client Ydb_Query_V1.QueryServiceClient,
 		cfg:        cfg,
 		grpcClient: client,
 		statusCode: statusUnknown,
+		checks: []func(s *Session) bool{
+			func(s *Session) bool {
+				switch s.status() {
+				case statusClosed, statusClosing, statusError:
+					return false
+				default:
+					return true
+				}
+			},
+		},
 	}
 	defer func() {
 		if finalErr != nil && s != nil {
@@ -263,6 +274,10 @@ func (s *Session) Begin(
 
 	tx, err = begin(ctx, s.grpcClient, s, txSettings)
 	if err != nil {
+		if xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION) {
+			s.setStatus(statusClosed)
+		}
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 	tx.s = s
@@ -301,6 +316,10 @@ func (s *Session) Execute(
 
 	tx, r, err := execute(ctx, s, s.grpcClient, q, options.ExecuteSettings(opts...))
 	if err != nil {
+		if xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION) {
+			s.setStatus(statusClosed)
+		}
+
 		return nil, nil, xerrors.WithStackTrace(err)
 	}
 
diff --git a/internal/query/transaction.go b/internal/query/transaction.go
index 4e316a840..94df6e23a 100644
--- a/internal/query/transaction.go
+++ b/internal/query/transaction.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-genproto/protos/Ydb_Query"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
@@ -80,6 +81,10 @@ func (tx transaction) Execute(ctx context.Context, q string, opts ...options.TxE
 
 	_, res, err := execute(ctx, tx.s, tx.s.grpcClient, q, options.TxExecuteSettings(tx.ID(), opts...).ExecuteSettings)
 	if err != nil {
+		if xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION) {
+			tx.s.setStatus(statusClosed)
+		}
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -98,8 +103,17 @@ func commitTx(ctx context.Context, client Ydb_Query_V1.QueryServiceClient, sessi
 	return nil
 }
 
-func (tx transaction) CommitTx(ctx context.Context) (err error) {
-	return commitTx(ctx, tx.s.grpcClient, tx.s.id, tx.ID())
+func (tx transaction) CommitTx(ctx context.Context) error {
+	err := commitTx(ctx, tx.s.grpcClient, tx.s.id, tx.ID())
+	if err != nil {
+		if xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION) {
+			tx.s.setStatus(statusClosed)
+		}
+
+		return xerrors.WithStackTrace(err)
+	}
+
+	return nil
 }
 
 func rollback(ctx context.Context, client Ydb_Query_V1.QueryServiceClient, sessionID, txID string) error {
@@ -114,6 +128,15 @@ func rollback(ctx context.Context, client Ydb_Query_V1.QueryServiceClient, sessi
 	return nil
 }
 
-func (tx transaction) Rollback(ctx context.Context) (err error) {
-	return rollback(ctx, tx.s.grpcClient, tx.s.id, tx.ID())
+func (tx transaction) Rollback(ctx context.Context) error {
+	err := rollback(ctx, tx.s.grpcClient, tx.s.id, tx.ID())
+	if err != nil {
+		if xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION) {
+			tx.s.setStatus(statusClosed)
+		}
+
+		return xerrors.WithStackTrace(err)
+	}
+
+	return nil
 }