correlation/context.go

package correlation

import (
	"context"
)

type ctxKey int

const (
	keyCorrelationID ctxKey = iota
	keyClientName
)

func extractFromContextByKey(ctx context.Context, key ctxKey) string {
	value := ctx.Value(key)

	str, ok := value.(string)
	if !ok {
		return ""
	}

	return str
}

// ExtractFromContext extracts the CollectionID from the provided context.
// Returns an empty string if it's unable to extract the CorrelationID for
// any reason.
func ExtractFromContext(ctx context.Context) string {
	return extractFromContextByKey(ctx, keyCorrelationID)
}

// ExtractFromContextOrGenerate extracts the CollectionID from the provided context or generates a random id if
// context does not contain one.
func ExtractFromContextOrGenerate(ctx context.Context) string {
	id := ExtractFromContext(ctx)
	if id == "" {
		id = SafeRandomID()
	}

	return id
}

// ContextWithCorrelation will create a new context containing the provided Correlation-ID value.
// This can be extracted using ExtractFromContext.
func ContextWithCorrelation(ctx context.Context, correlationID string) context.Context {
	return context.WithValue(ctx, keyCorrelationID, correlationID)
}

// ExtractClientNameFromContext extracts client name from incoming context.
// It will return an empty string if client name does not exist in the context.
func ExtractClientNameFromContext(ctx context.Context) string {
	return extractFromContextByKey(ctx, keyClientName)
}

// ContextWithClientName will create a new context containing client_name metadata.
func ContextWithClientName(ctx context.Context, clientName string) context.Context {
	return context.WithValue(ctx, keyClientName, clientName)
}