context.go

package logging

import (
	"context"
	"sort"
)

// ContextExtractorFunc extracts a logging field from context.
// It returns the field Key and its value. A nil value means "no value in this context", the field is omitted.
// The function is also called with context.Background() once on registration to learn the Key,
// so it must not panic on an empty context.
type ContextExtractorFunc func(ctx context.Context) (Key, any)

// WithContextExtractors adds context value to logging field extraction functions to logging config.
// Predefined Key's could be added with Key.Extractor method.
// If several extractors report the same Key the last registered wins.
func WithContextExtractors(extractors ...ContextExtractorFunc) Option {
	return func(config *Config) {
		if config.contextExtractors == nil {
			config.contextExtractors = make(map[Key]ContextExtractorFunc, len(extractors))
		}

		ctx := context.Background()

		for _, extractor := range extractors {
			if extractor == nil {
				continue
			}

			k, _ := extractor(ctx)
			config.contextExtractors[k] = extractor
		}
	}
}

// contextKeys extracts logging keys from context using set with WithContextExtractors.
func (config *Config) contextKeys(ctx context.Context) Keys {
	res := make(Keys, len(config.contextExtractors))

	if ctx == nil {
		return res
	}

	for _, e := range config.contextExtractors {
		if k, v := e(ctx); v != nil {
			res[k] = v
		}
	}

	return res
}

// sortedExtractors returns registered extractors ordered by their keys to get a stable fields order.
func (config *Config) sortedExtractors() []ContextExtractorFunc {
	keys := make([]string, 0, len(config.contextExtractors))
	for k := range config.contextExtractors {
		keys = append(keys, k.String())
	}

	sort.Strings(keys)

	extractors := make([]ContextExtractorFunc, len(keys))
	for i, k := range keys {
		extractors[i] = config.contextExtractors[Key(k)]
	}

	return extractors
}

// loggableContextKey used internally to represent logging keys in context.
type loggableContextKey string

// loggingKey returns context key type for logging Key.
func (k Key) loggingKey() loggableContextKey {
	return loggableContextKey(k.String())
}

// SetToCtx sets Key value into context.
func (k Key) SetToCtx(ctx context.Context, value any) context.Context {
	return context.WithValue(ctx, k.loggingKey(), value)
}

// getFromCtx returns logging key value from context.
func (k Key) getFromCtx(ctx context.Context) any {
	return ctx.Value(k.loggingKey())
}

// Extractor returns ContextExtractorFunc for Key.
func (k Key) Extractor() ContextExtractorFunc {
	return func(ctx context.Context) (Key, any) {
		return k, k.getFromCtx(ctx)
	}
}

// KeysCtx returns logging keys from context using all extractors set with WithContextExtractors.
func KeysCtx(ctx context.Context) Keys {
	return CurrentConfig().contextKeys(ctx)
}