spanlogger: Take an interface implementation for extracting tenant IDs

CHANGELOG.md

@@ -7,9 +7,11 @@
 * [CHANGE] Added CHANGELOG.md and Pull Request template to reference the changelog
 * [CHANGE] Remove `cortex_` prefix for metrics registered in the ring. #46
 * [CHANGE] Rename `kv/kvtls` to `crypto/tls`. #39
+* [CHANGE] spanlogger: Take interface implementation for extracting tenant ID. #59
 * [ENHANCEMENT] Add middleware package. #38
 * [ENHANCEMENT] Add the ring package #45
 * [ENHANCEMENT] Add limiter package. #41
 * [ENHANCEMENT] Add grpcclient, grpcencoding and grpcutil packages. #39
 * [ENHANCEMENT] Replace go-kit/kit/log with go-kit/log. #52
-* [ENHANCEMENT] Add spanlogger package. #42
+* [ENHANCEMENT] Add spanlogger package. #42
+* [BUGFIX] spanlogger: Support multiple tenant IDs. #59

spanlogger/spanlogger.go

@@ -2,19 +2,25 @@ package spanlogger
 
 import (
 	"context"
-	"strings"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	opentracing "github.com/opentracing/opentracing-go"
 	"github.com/opentracing/opentracing-go/ext"
 	otlog "github.com/opentracing/opentracing-go/log"
 	"github.com/weaveworks/common/tracing"
-	"github.com/weaveworks/common/user"
 )
 
 type loggerCtxMarker struct{}
 
+// TenantResolver provides methods for extracting tenant IDs from a context.
+type TenantResolver interface {
+	// TenantID tries to extract a tenant ID from a context.
+	TenantID(context.Context) (string, error)
+	// TenantIDs tries to extract tenant IDs from a context.
+	TenantIDs(context.Context) ([]string, error)
+}
+
 const (
 	// TenantIDsTagName is the tenant IDs tag name.
 	TenantIDsTagName = "tenant_ids"
@@ -32,13 +38,13 @@ type SpanLogger struct {
 
 // New makes a new SpanLogger with a log.Logger to send logs to. The provided context will have the logger attached
 // to it and can be retrieved with FromContext.
-func New(ctx context.Context, logger log.Logger, method string, kvps ...interface{}) (*SpanLogger, context.Context) {
+func New(ctx context.Context, logger log.Logger, method string, resolver TenantResolver, kvps ...interface{}) (*SpanLogger, context.Context) {
 	span, ctx := opentracing.StartSpanFromContext(ctx, method)
-	if id := tenantID(ctx); id != "" {
-		span.SetTag(TenantIDsTagName, []string{id})
+	if ids, err := resolver.TenantIDs(ctx); err == nil && len(ids) > 0 {
+		span.SetTag(TenantIDsTagName, ids)
 	}
 	l := &SpanLogger{
-		Logger: log.With(withContext(ctx, logger), "method", method),
+		Logger: log.With(withContext(ctx, logger, resolver), "method", method),
 		Span:   span,
 	}
 	if len(kvps) > 0 {
@@ -53,7 +59,7 @@ func New(ctx context.Context, logger log.Logger, method string, kvps ...interfac
 // If there is no parent span, the SpanLogger will only log to the logger
 // within the context. If the context doesn't have a logger, the fallback
 // logger is used.
-func FromContext(ctx context.Context, fallback log.Logger) *SpanLogger {
+func FromContext(ctx context.Context, fallback log.Logger, resolver TenantResolver) *SpanLogger {
 	logger, ok := ctx.Value(loggerCtxKey).(log.Logger)
 	if !ok {
 		logger = fallback
@@ -63,7 +69,7 @@ func FromContext(ctx context.Context, fallback log.Logger) *SpanLogger {
 		sp = defaultNoopSpan
 	}
 	return &SpanLogger{
-		Logger: withContext(ctx, logger),
+		Logger: withContext(ctx, logger, resolver),
 		Span:   sp,
 	}
 }
@@ -90,34 +96,18 @@ func (s *SpanLogger) Error(err error) error {
 	return err
 }
 
-// tenantID tries to extract the tenant ID from ctx.
-func tenantID(ctx context.Context) string {
-	//lint:ignore faillint wrapper around upstream method
-	id, err := user.ExtractOrgID(ctx)
-	if err != nil {
-		return ""
-	}
-
-	// handle the relative reference to current and parent path.
-	if id == "." || id == ".." || strings.ContainsAny(id, `\/`) {
-		return ""
-	}
-
-	return id
-}
-
-func withContext(ctx context.Context, l log.Logger) log.Logger {
+func withContext(ctx context.Context, logger log.Logger, resolver TenantResolver) log.Logger {
 	// Weaveworks uses "orgs" and "orgID" to represent Cortex users,
 	// even though the code-base generally uses `userID` to refer to the same thing.
-	userID := tenantID(ctx)
-	if userID != "" {
-		l = log.With(l, "org_id", userID)
+	userID, err := resolver.TenantID(ctx)
+	if err == nil && userID != "" {
+		logger = log.With(logger, "org_id", userID)
 	}
 
 	traceID, ok := tracing.ExtractSampledTraceID(ctx)
 	if !ok {
-		return l
+		return logger
 	}
 
-	return log.With(l, "traceID", traceID)
+	return log.With(logger, "traceID", traceID)
 }

spanlogger/spanlogger_test.go

@@ -2,6 +2,7 @@ package spanlogger
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	"github.com/go-kit/log"
@@ -14,12 +15,13 @@ import (
 
 func TestSpanLogger_Log(t *testing.T) {
 	logger := log.NewNopLogger()
-	span, ctx := New(context.Background(), logger, "test", "bar")
+	resolver := fakeResolver{}
+	span, ctx := New(context.Background(), logger, "test", resolver, "bar")
 	_ = span.Log("foo")
-	newSpan := FromContext(ctx, logger)
+	newSpan := FromContext(ctx, logger, resolver)
 	require.Equal(t, span.Span, newSpan.Span)
 	_ = newSpan.Log("bar")
-	noSpan := FromContext(context.Background(), logger)
+	noSpan := FromContext(context.Background(), logger, resolver)
 	_ = noSpan.Log("foo")
 	require.Error(t, noSpan.Error(errors.New("err")))
 	require.NoError(t, noSpan.Error(nil))
@@ -31,13 +33,14 @@ func TestSpanLogger_CustomLogger(t *testing.T) {
 		logged = append(logged, keyvals)
 		return nil
 	}
-	span, ctx := New(context.Background(), logger, "test")
+	resolver := fakeResolver{}
+	span, ctx := New(context.Background(), logger, "test", resolver)
 	_ = span.Log("msg", "original spanlogger")
 
-	span = FromContext(ctx, log.NewNopLogger())
+	span = FromContext(ctx, log.NewNopLogger(), resolver)
 	_ = span.Log("msg", "restored spanlogger")
 
-	span = FromContext(context.Background(), logger)
+	span = FromContext(context.Background(), logger, resolver)
 	_ = span.Log("msg", "fallback spanlogger")
 
 	expect := [][]interface{}{
@@ -57,15 +60,15 @@ func TestSpanCreatedWithTenantTag(t *testing.T) {
 func TestSpanCreatedWithoutTenantTag(t *testing.T) {
 	mockSpan := createSpan(context.Background())
 
-	_, exist := mockSpan.Tags()[TenantIDsTagName]
-	require.False(t, exist)
+	_, exists := mockSpan.Tags()[TenantIDsTagName]
+	require.False(t, exists)
 }
 
 func createSpan(ctx context.Context) *mocktracer.MockSpan {
 	mockTracer := mocktracer.New()
 	opentracing.SetGlobalTracer(mockTracer)
 
-	logger, _ := New(ctx, log.NewNopLogger(), "name")
+	logger, _ := New(ctx, log.NewNopLogger(), "name", fakeResolver{})
 	return logger.Span.(*mocktracer.MockSpan)
 }
 
@@ -74,3 +77,32 @@ type funcLogger func(keyvals ...interface{}) error
 func (f funcLogger) Log(keyvals ...interface{}) error {
 	return f(keyvals...)
 }
+
+type fakeResolver struct {
+}
+
+func (fakeResolver) TenantID(ctx context.Context) (string, error) {
+	id, err := user.ExtractOrgID(ctx)
+	if err != nil {
+		return "", err
+	}
+
+	// handle the relative reference to current and parent path.
+	if id == "." || id == ".." || strings.ContainsAny(id, `\/`) {
+		return "", nil
+	}
+
+	return id, nil
+}
+
+func (r fakeResolver) TenantIDs(ctx context.Context) ([]string, error) {
+	id, err := r.TenantID(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if id == "" {
+		return nil, nil
+	}
+
+	return []string{id}, nil
+}