github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/xtest/grpclogger.go (about)

     1  package xtest
     2  
     3  import (
     4  	"context"
     5  	"sync/atomic"
     6  	"testing"
     7  
     8  	"google.golang.org/grpc"
     9  )
    10  
    11  var globalLastStreamID = int64(0)
    12  
    13  // GrpcLogger use for log raw grpc messages
    14  //
    15  // Usage:
    16  //
    17  //		db, err := ydb.Open(context.Background(), connectionString,
    18  //	     ...
    19  //			ydb.Change(config.WithGrpcOptions(grpc.WithChainUnaryInterceptor(xtest.NewGrpcLogger(t).UnaryClientInterceptor))),
    20  //		)
    21  type GrpcLogger struct {
    22  	t testing.TB
    23  }
    24  
    25  func NewGrpcLogger(t testing.TB) GrpcLogger {
    26  	return GrpcLogger{t: t}
    27  }
    28  
    29  func (l GrpcLogger) UnaryClientInterceptor(
    30  	ctx context.Context,
    31  	method string,
    32  	req, reply interface{},
    33  	cc *grpc.ClientConn,
    34  	invoker grpc.UnaryInvoker,
    35  	opts ...grpc.CallOption,
    36  ) error {
    37  	err := invoker(ctx, method, req, reply, cc, opts...)
    38  	if err != nil {
    39  		l.t.Logf("UnaryClientInterceptor: %s - err: %v\n\nreq:\n%v\n\nresp:\n%v", method, err, req, reply)
    40  	} else {
    41  		l.t.Logf("UnaryClientInterceptor: %s:\n\nreq:\n%v\n\nresp:\n%v", method, req, reply)
    42  	}
    43  
    44  	return err
    45  }
    46  
    47  func (l GrpcLogger) StreamClientInterceptor(
    48  	ctx context.Context,
    49  	desc *grpc.StreamDesc,
    50  	cc *grpc.ClientConn,
    51  	method string,
    52  	streamer grpc.Streamer,
    53  	opts ...grpc.CallOption,
    54  ) (grpc.ClientStream, error) {
    55  	stream, err := streamer(ctx, desc, cc, method, opts...)
    56  	streamWrapper := newGrpcLoggerStream(stream, l.t)
    57  	if stream != nil {
    58  		stream = streamWrapper
    59  	}
    60  	if err != nil {
    61  		l.t.Logf("StreamStart: %v with err '%v' (streamID: %v)", method, err, streamWrapper.streamID)
    62  	} else {
    63  		l.t.Logf("StreamStart: %v (streamID: %v)", method, streamWrapper.streamID)
    64  	}
    65  
    66  	return stream, err
    67  }
    68  
    69  type grpcLoggerStream struct {
    70  	grpc.ClientStream
    71  	streamID int64
    72  	t        testing.TB
    73  }
    74  
    75  func newGrpcLoggerStream(stream grpc.ClientStream, t testing.TB) grpcLoggerStream {
    76  	return grpcLoggerStream{stream, atomic.AddInt64(&globalLastStreamID, 1), t}
    77  }
    78  
    79  func (g grpcLoggerStream) CloseSend() error {
    80  	err := g.ClientStream.CloseSend()
    81  	if err != nil {
    82  		g.t.Logf("CloseSend: %v (streamID: %v)", err, g.streamID)
    83  	} else {
    84  		g.t.Logf("CloseSend (streamID: %v)", g.streamID)
    85  	}
    86  
    87  	return err
    88  }
    89  
    90  func (g grpcLoggerStream) SendMsg(m interface{}) error {
    91  	err := g.ClientStream.SendMsg(m)
    92  	if err != nil {
    93  		g.t.Logf("SendMsg (streamID: %v) with err '%v':\n%v ", g.streamID, err, m)
    94  	} else {
    95  		g.t.Logf("SendMsg (streamID: %v):\n%v ", g.streamID, m)
    96  	}
    97  
    98  	return err
    99  }
   100  
   101  func (g grpcLoggerStream) RecvMsg(m interface{}) error {
   102  	err := g.ClientStream.RecvMsg(m)
   103  	if err != nil {
   104  		g.t.Logf("RecvMsg (streamID: %v) with err '%v':\n%v ", g.streamID, err, m)
   105  	} else {
   106  		g.t.Logf("RecvMsg (streamID: %v):\n%v ", g.streamID, m)
   107  	}
   108  
   109  	return err
   110  }