github.com/msales/pkg/v3@v3.24.0/grpcx/middleware/context.go (about)

     1  package middleware
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/grpc-ecosystem/go-grpc-middleware"
     8  	"github.com/msales/pkg/v3/log"
     9  	"github.com/msales/pkg/v3/stats"
    10  	"google.golang.org/grpc"
    11  )
    12  
    13  // WithUnaryServerLogger adds the logger instance to the unary request context.
    14  func WithUnaryServerLogger(l log.Logger) grpc.UnaryServerInterceptor {
    15  	return func(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    16  		return handler(log.WithLogger(ctx, l), req)
    17  	}
    18  }
    19  
    20  // WithStreamServerLogger adds the stats instance to the stream context.
    21  func WithStreamServerLogger(l log.Logger) grpc.StreamServerInterceptor {
    22  	return func(srv interface{}, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    23  		ws := grpc_middleware.WrapServerStream(ss)
    24  		ws.WrappedContext = log.WithLogger(ss.Context(), l)
    25  
    26  		return handler(srv, ws)
    27  	}
    28  }
    29  
    30  // WithUnaryServerStats adds the stats instance to the unary request context.
    31  func WithUnaryServerStats(s stats.Stats) grpc.UnaryServerInterceptor {
    32  	return func(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    33  		return handler(stats.WithStats(ctx, s), req)
    34  	}
    35  }
    36  
    37  // WithStreamServerStats adds the stats instance to the stream context.
    38  func WithStreamServerStats(s stats.Stats) grpc.StreamServerInterceptor {
    39  	return func(srv interface{}, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    40  		ws := grpc_middleware.WrapServerStream(ss)
    41  		ws.WrappedContext = stats.WithStats(ss.Context(), s)
    42  
    43  		return handler(srv, ws)
    44  	}
    45  }
    46  
    47  // WithUnaryClientContextTimeout adds timeout to unary client request context.
    48  func WithUnaryClientContextTimeout(d time.Duration) grpc.UnaryClientInterceptor {
    49  	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    50  		ctx, cancel := context.WithTimeout(ctx, d)
    51  		defer cancel()
    52  
    53  		return invoker(ctx, method, req, reply, cc, opts...)
    54  	}
    55  }
    56  
    57  // WithUnaryClientLogger adds the logger instance to the unary client request context.
    58  func WithUnaryClientLogger(l log.Logger) grpc.UnaryClientInterceptor {
    59  	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    60  		return invoker(log.WithLogger(ctx, l), method, req, reply, cc, opts...)
    61  	}
    62  }
    63  
    64  // WithUnaryClientStats adds the stats instance to the unary client request context.
    65  func WithUnaryClientStats(s stats.Stats) grpc.UnaryClientInterceptor {
    66  	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    67  		return invoker(stats.WithStats(ctx, s), method, req, reply, cc, opts...)
    68  	}
    69  }