@@ -2,6 +2,7 @@ package query
22
33import (
44 "context"
5+ "io"
56 "sync/atomic"
67
78 "github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
@@ -11,7 +12,6 @@ import (
1112 "github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
1213 "github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
1314 "github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
14- "github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
1515 "github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
1616 "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
1717 "github.com/ydb-platform/ydb-go-sdk/v3/query"
@@ -170,28 +170,14 @@ func (s *Session) attach(ctx context.Context) (finalErr error) {
170170 onDone (finalErr )
171171 }()
172172
173- attachCtx , cancelAttach := xcontext .WithCancel (context .Background ())
174- defer func () {
175- if finalErr != nil {
176- cancelAttach ()
177- }
178- }()
179-
180- attach , err := s .grpcClient .AttachSession (attachCtx , & Ydb_Query.AttachSessionRequest {
173+ attach , err := s .grpcClient .AttachSession (context .Background (), & Ydb_Query.AttachSessionRequest {
181174 SessionId : s .id ,
182175 })
183176 if err != nil {
184177 return xerrors .WithStackTrace (
185178 xerrors .Transport (err ),
186179 )
187180 }
188-
189- defer func () {
190- if finalErr != nil {
191- _ = attach .CloseSend ()
192- }
193- }()
194-
195181 state , err := attach .Recv ()
196182 if err != nil {
197183 return xerrors .WithStackTrace (xerrors .Transport (err ))
@@ -202,13 +188,21 @@ func (s *Session) attach(ctx context.Context) (finalErr error) {
202188 }
203189
204190 go func () {
205- defer func () {
206- _ = s .closeOnce (ctx )
207- }()
208-
209191 for {
192+ if ! s .IsAlive () {
193+ return
194+ }
210195 recv , recvErr := attach .Recv ()
211- if recvErr != nil || recv .GetStatus () != Ydb .StatusIds_SUCCESS {
196+ if recvErr != nil {
197+ if xerrors .Is (recvErr , io .EOF ) {
198+ s .setStatus (statusClosed )
199+ } else {
200+ s .setStatus (statusError )
201+ }
202+
203+ return
204+ }
205+ if recv .GetStatus () != Ydb .StatusIds_SUCCESS {
212206 s .setStatus (statusError )
213207
214208 return
@@ -227,16 +221,32 @@ func (s *Session) attach(ctx context.Context) (finalErr error) {
227221 }()
228222 }
229223
230- err = attach . CloseSend ()
231-
232- cancelAttach ()
224+ if err = deleteSession ( ctx , s . grpcClient , s . id ); err != nil {
225+ return xerrors . WithStackTrace ( err )
226+ }
233227
234- return err
228+ return nil
235229 })
236230
237231 return nil
238232}
239233
234+ func deleteSession (ctx context.Context , client Ydb_Query_V1.QueryServiceClient , sessionID string ) error {
235+ response , err := client .DeleteSession (ctx ,
236+ & Ydb_Query.DeleteSessionRequest {
237+ SessionId : sessionID ,
238+ },
239+ )
240+ if err != nil {
241+ return xerrors .WithStackTrace (xerrors .Transport (err ))
242+ }
243+ if response .GetStatus () != Ydb .StatusIds_SUCCESS {
244+ return xerrors .WithStackTrace (xerrors .FromOperation (response ))
245+ }
246+
247+ return nil
248+ }
249+
240250func (s * Session ) IsAlive () bool {
241251 switch s .status () {
242252 case statusIdle , statusInUse :
0 commit comments