@@ -11,6 +11,7 @@ import (
1111 "google.golang.org/grpc"
1212 "google.golang.org/grpc/codes"
1313 "google.golang.org/grpc/connectivity"
14+ grpcStatus "google.golang.org/grpc/status"
1415
1516 "github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
1617
@@ -186,8 +187,8 @@ func (c *conn) take(ctx context.Context) (cc *grpc.ClientConn, err error) {
186187
187188 if c .isClosed () {
188189 return nil , errors .NewGrpcError (
189- codes .Unavailable ,
190- errors .WithMsg ( "ydb driver conn closed early" ),
190+ errors . WithStatus ( grpcStatus . New ( codes .Unavailable , "ydb driver conn closed early" )) ,
191+ errors .WithErr ( err ),
191192 )
192193 }
193194
@@ -312,6 +313,7 @@ func (c *conn) Close(ctx context.Context) (err error) {
312313 return err
313314}
314315
316+ // invoke have behavior like grpc call
315317func (c * conn ) invoke (
316318 ctx context.Context ,
317319 method string ,
@@ -323,25 +325,34 @@ func (c *conn) invoke(
323325 cc , err = c .take (ctx )
324326 if err != nil {
325327 return errors .NewGrpcError (
326- codes .Unavailable ,
327- errors .WithMsg ("ydb driver conn take failed" ),
328+ errors .WithStatus (grpcStatus .New (codes .Unavailable , "ydb driver conn take failed" )),
328329 errors .WithErr (err ),
329330 )
330331 }
331332
332333 ctx , err = c .config .Meta ().Meta (ctx )
333334 if err != nil {
334335 return errors .NewGrpcError (
335- codes .Unavailable ,
336- errors .WithMsg ("ydb driver conn apply meta failed" ),
336+ errors .WithStatus (grpcStatus .New (codes .Unavailable , "ydb driver conn apply meta failed" )),
337337 errors .WithErr (err ),
338338 )
339339 }
340340
341341 c .changeUsages (1 )
342342 defer c .changeUsages (- 1 )
343343
344- return cc .Invoke (ctx , method , req , res , opts ... )
344+ err = cc .Invoke (ctx , method , req , res , opts ... )
345+
346+ if err != nil {
347+ if s , ok := grpcStatus .FromError (err ); ok {
348+ return errors .NewGrpcError (
349+ errors .WithStatus (s ),
350+ )
351+ }
352+ return errors .WithStackTrace (err )
353+ }
354+
355+ return nil
345356}
346357
347358func (c * conn ) Invoke (
@@ -400,6 +411,7 @@ func (c *conn) Invoke(
400411 return err
401412}
402413
414+ // newStream have behavior like grpc call
403415func (c * conn ) newStream (
404416 ctx context.Context ,
405417 desc * grpc.StreamDesc ,
@@ -410,25 +422,27 @@ func (c *conn) newStream(
410422 cc , err = c .take (ctx )
411423 if err != nil {
412424 return nil , errors .NewGrpcError (
413- codes .Unavailable ,
414- errors .WithMsg ("ydb driver conn take failed" ),
415- errors .WithErr (err ),
416- )
417- }
418-
419- ctx , err = c .config .Meta ().Meta (ctx )
420- if err != nil {
421- return nil , errors .NewGrpcError (
422- codes .Unavailable ,
423- errors .WithMsg ("ydb driver conn apply meta failed" ),
425+ errors .WithStatus (grpcStatus .New (codes .Unavailable , "ydb driver conn take failed" )),
424426 errors .WithErr (err ),
425427 )
426428 }
427429
428430 c .changeStreamUsages (1 )
429431 defer c .changeStreamUsages (- 1 )
430432
431- return cc .NewStream (ctx , desc , method , opts ... )
433+ var client grpc.ClientStream
434+ client , err = cc .NewStream (ctx , desc , method , opts ... )
435+
436+ if err != nil {
437+ if s , ok := grpcStatus .FromError (err ); ok {
438+ return nil , errors .NewGrpcError (
439+ errors .WithStatus (s ),
440+ )
441+ }
442+ return nil , errors .WithStackTrace (err )
443+ }
444+
445+ return client , nil
432446}
433447
434448func (c * conn ) NewStream (
0 commit comments