gitlab.com/gitlab-org/labkit@v1.21.0/correlation/context.go (about)

     1  package correlation
     2  
     3  import (
     4  	"context"
     5  )
     6  
     7  type ctxKey int
     8  
     9  const (
    10  	keyCorrelationID ctxKey = iota
    11  	keyClientName
    12  )
    13  
    14  func extractFromContextByKey(ctx context.Context, key ctxKey) string {
    15  	value := ctx.Value(key)
    16  
    17  	str, ok := value.(string)
    18  	if !ok {
    19  		return ""
    20  	}
    21  
    22  	return str
    23  }
    24  
    25  // ExtractFromContext extracts the CollectionID from the provided context.
    26  // Returns an empty string if it's unable to extract the CorrelationID for
    27  // any reason.
    28  func ExtractFromContext(ctx context.Context) string {
    29  	return extractFromContextByKey(ctx, keyCorrelationID)
    30  }
    31  
    32  // ExtractFromContextOrGenerate extracts the CollectionID from the provided context or generates a random id if
    33  // context does not contain one.
    34  func ExtractFromContextOrGenerate(ctx context.Context) string {
    35  	id := ExtractFromContext(ctx)
    36  	if id == "" {
    37  		id = SafeRandomID()
    38  	}
    39  	return id
    40  }
    41  
    42  // ContextWithCorrelation will create a new context containing the provided Correlation-ID value.
    43  // This can be extracted using ExtractFromContext.
    44  func ContextWithCorrelation(ctx context.Context, correlationID string) context.Context {
    45  	return context.WithValue(ctx, keyCorrelationID, correlationID)
    46  }
    47  
    48  // ExtractClientNameFromContext extracts client name from incoming context.
    49  // It will return an empty string if client name does not exist in the context.
    50  func ExtractClientNameFromContext(ctx context.Context) string {
    51  	return extractFromContextByKey(ctx, keyClientName)
    52  }
    53  
    54  // ContextWithClientName will create a new context containing client_name metadata.
    55  func ContextWithClientName(ctx context.Context, clientName string) context.Context {
    56  	return context.WithValue(ctx, keyClientName, clientName)
    57  }