github.com/openfga/openfga@v1.5.4-rc1/pkg/middleware/requestid/requestid.go (about)

     1  package requestid
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/google/uuid"
     7  	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors"
     8  	"go.opentelemetry.io/otel/attribute"
     9  	"go.opentelemetry.io/otel/trace"
    10  	"google.golang.org/grpc"
    11  	"google.golang.org/grpc/metadata"
    12  )
    13  
    14  const (
    15  	requestIDCtxKey   = "request-id-context-key"
    16  	requestIDTraceKey = "request_id"
    17  
    18  	// RequestIDHeader defines the HTTP header that is set in each HTTP response
    19  	// for a given request. The value of the header is unique per request.
    20  	RequestIDHeader = "X-Request-Id"
    21  )
    22  
    23  // FromContext extracts the request-id from the context, if it exists.
    24  func FromContext(ctx context.Context) (string, bool) {
    25  	if md, ok := metadata.FromOutgoingContext(ctx); ok {
    26  		if vals := md.Get(requestIDCtxKey); len(vals) > 0 {
    27  			return vals[0], true
    28  		}
    29  	}
    30  
    31  	return "", false
    32  }
    33  
    34  // NewUnaryInterceptor creates a grpc.UnaryServerInterceptor which must
    35  // come after the trace interceptor and before the logging interceptor.
    36  func NewUnaryInterceptor() grpc.UnaryServerInterceptor {
    37  	return interceptors.UnaryServerInterceptor(reportable())
    38  }
    39  
    40  // NewStreamingInterceptor creates a grpc.StreamServerInterceptor which must
    41  // come after the trace interceptor and before the logging interceptor.
    42  func NewStreamingInterceptor() grpc.StreamServerInterceptor {
    43  	return interceptors.StreamServerInterceptor(reportable())
    44  }
    45  
    46  func reportable() interceptors.CommonReportableFunc {
    47  	return func(ctx context.Context, c interceptors.CallMeta) (interceptors.Reporter, context.Context) {
    48  		// TODO use ulid library?
    49  		id, _ := uuid.NewRandom()
    50  		requestID := id.String()
    51  
    52  		ctx = metadata.AppendToOutgoingContext(ctx, requestIDCtxKey, requestID)
    53  
    54  		trace.SpanFromContext(ctx).SetAttributes(attribute.String(requestIDTraceKey, requestID))
    55  
    56  		_ = grpc.SetHeader(ctx, metadata.Pairs(RequestIDHeader, requestID))
    57  
    58  		return interceptors.NoopReporter{}, ctx
    59  	}
    60  }