github.com/grailbio/base@v0.0.11/common/log/context.go (about)

     1  package log
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/gin-gonic/gin"
     7  	"github.com/google/uuid"
     8  )
     9  
    10  const (
    11  	// Best practices for avoiding key collisions in context say this shouldn't be a string type.
    12  	// Gin forces us to use a string for their context, so choose a name that is unlikely to collide with anything else.
    13  	RequestIDContextKey = "grail_logger_request_id"
    14  )
    15  
    16  // WithRequestID sets the uuid value for the RequestIDContextKey key in the context.
    17  func WithRequestID(ctx context.Context, requestID uuid.UUID) context.Context {
    18  	return context.WithValue(ctx, RequestIDContextKey, requestID)
    19  }
    20  
    21  // WithGinRequestID creates a uuid that is set as a string on the gin Context and as
    22  // a uuid on the regular-flavor Request context that it wraps. The context should
    23  // be passed to the methods in this package to prefix logs with the identifier.
    24  func WithGinRequestID(ctx *gin.Context) {
    25  	requuid := uuid.New()
    26  	uuidStr := requuid.String()
    27  	if _, ok := ctx.Get(RequestIDContextKey); ok {
    28  		return // Avoid overwriting the original value in case this middleware is invoked twice
    29  	}
    30  	ctx.Set(RequestIDContextKey, uuidStr)
    31  	// TODO: ideally we'd pass the  X-Amzn-Trace-Id header from our ALB, but we're not using ALBs yet.
    32  	ctx.Request = ctx.Request.WithContext(WithRequestID(ctx.Request.Context(), requuid))
    33  }