github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/conn/middleware.go (about)

     1  package conn
     2  
     3  import (
     4  	"context"
     5  
     6  	"google.golang.org/grpc"
     7  )
     8  
     9  var _ grpc.ClientConnInterface = (*middleware)(nil)
    10  
    11  type (
    12  	invoker  func(context.Context, string, interface{}, interface{}, ...grpc.CallOption) error
    13  	streamer func(context.Context, *grpc.StreamDesc, string, ...grpc.CallOption) (grpc.ClientStream, error)
    14  )
    15  
    16  type middleware struct {
    17  	invoke    invoker
    18  	newStream streamer
    19  }
    20  
    21  func (m *middleware) Invoke(
    22  	ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption,
    23  ) error {
    24  	return m.invoke(ctx, method, args, reply, opts...)
    25  }
    26  
    27  func (m *middleware) NewStream(
    28  	ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption,
    29  ) (grpc.ClientStream, error) {
    30  	return m.newStream(ctx, desc, method, opts...)
    31  }
    32  
    33  func WithContextModifier(
    34  	cc grpc.ClientConnInterface,
    35  	modifyCtx func(ctx context.Context) context.Context,
    36  ) grpc.ClientConnInterface {
    37  	return &middleware{
    38  		invoke: func(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error {
    39  			ctx = modifyCtx(ctx)
    40  
    41  			return cc.Invoke(ctx, method, args, reply, opts...)
    42  		},
    43  		newStream: func(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (
    44  			grpc.ClientStream, error,
    45  		) {
    46  			ctx = modifyCtx(ctx)
    47  
    48  			return cc.NewStream(ctx, desc, method, opts...)
    49  		},
    50  	}
    51  }
    52  
    53  func WithAppendOptions(cc grpc.ClientConnInterface, appendOpts ...grpc.CallOption) grpc.ClientConnInterface {
    54  	return &middleware{
    55  		invoke: func(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error {
    56  			opts = append(opts, appendOpts...)
    57  
    58  			return cc.Invoke(ctx, method, args, reply, opts...)
    59  		},
    60  		newStream: func(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (
    61  			grpc.ClientStream, error,
    62  		) {
    63  			opts = append(opts, appendOpts...)
    64  
    65  			return cc.NewStream(ctx, desc, method, opts...)
    66  		},
    67  	}
    68  }
    69  
    70  func WithBeforeFunc(
    71  	cc grpc.ClientConnInterface,
    72  	before func(),
    73  ) grpc.ClientConnInterface {
    74  	return &middleware{
    75  		invoke: func(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error {
    76  			before()
    77  
    78  			return cc.Invoke(ctx, method, args, reply, opts...)
    79  		},
    80  		newStream: func(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (
    81  			grpc.ClientStream, error,
    82  		) {
    83  			before()
    84  
    85  			return cc.NewStream(ctx, desc, method, opts...)
    86  		},
    87  	}
    88  }