diff --git a/internal/query/config/config.go b/internal/query/config/config.go index 8bd1262f6..9cd96303f 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: false, } } 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..fce975f9b 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) { + 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 }