github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/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.With(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  	l.t.Logf(
    39  		"UnaryClientInterceptor: %s - err: %v\n\nreq:\n%v\n\nresp:\n%v",
    40  		method,
    41  		err,
    42  		req,
    43  		reply,
    44  	)
    45  
    46  	return err
    47  }
    48  
    49  func (l GrpcLogger) StreamClientInterceptor(
    50  	ctx context.Context,
    51  	desc *grpc.StreamDesc,
    52  	cc *grpc.ClientConn,
    53  	method string,
    54  	streamer grpc.Streamer,
    55  	opts ...grpc.CallOption,
    56  ) (grpc.ClientStream, error) {
    57  	stream, err := streamer(ctx, desc, cc, method, opts...)
    58  	streamWrapper := newGrpcLoggerStream(stream, l.t)
    59  	if stream != nil {
    60  		stream = streamWrapper
    61  	}
    62  	l.t.Logf(
    63  		"StreamStart: %v with err '%v' (streamID: %v)",
    64  		method,
    65  		err,
    66  		streamWrapper.streamID,
    67  	)
    68  
    69  	return stream, err
    70  }
    71  
    72  type grpcLoggerStream struct {
    73  	grpc.ClientStream
    74  	streamID int64
    75  	t        testing.TB
    76  }
    77  
    78  func newGrpcLoggerStream(stream grpc.ClientStream, t testing.TB) grpcLoggerStream {
    79  	return grpcLoggerStream{stream, atomic.AddInt64(&globalLastStreamID, 1), t}
    80  }
    81  
    82  func (g grpcLoggerStream) CloseSend() error {
    83  	err := g.ClientStream.CloseSend()
    84  	g.t.Logf("CloseSend: %v (streamID: %v)", err, g.streamID)
    85  
    86  	return err
    87  }
    88  
    89  func (g grpcLoggerStream) SendMsg(m interface{}) error {
    90  	err := g.ClientStream.SendMsg(m)
    91  	g.t.Logf("SendMsg (streamID: %v) with err '%v':\n%v ", g.streamID, err, m)
    92  
    93  	return err
    94  }
    95  
    96  func (g grpcLoggerStream) RecvMsg(m interface{}) error {
    97  	err := g.ClientStream.RecvMsg(m)
    98  	g.t.Logf("RecvMsg (streamID: %v) with err '%v':\n%v ", g.streamID, err, m)
    99  
   100  	return err
   101  }