introduce contextual log level

context.go

@@ -14,8 +14,51 @@ func (key *ctxKey) String() string {
 	return key.name
 }
 
+// key is the context key for the attributes.
 var key = &ctxKey{"ctxslog"}
 
+type ctxLevelKey struct{ name string }
+
+func (key *ctxLevelKey) String() string {
+	return key.name
+}
+
+// levelKey is the context key for the minimum record level.
+var levelKey = &ctxLevelKey{"ctxslog.level"}
+
+var _ slog.Handler = (*wrapper)(nil)
+
+type wrapper struct {
+	parent slog.Handler
+}
+
+func (w *wrapper) Handle(ctx context.Context, record slog.Record) error {
+	attrs := contextAttrs(ctx)
+	newRecord := record.Clone()
+	attrs.addToRecord(&newRecord)
+	return w.parent.Handle(ctx, newRecord)
+}
+
+func (w *wrapper) Enabled(ctx context.Context, level slog.Level) bool {
+	minLevel, ok := contextLevel(ctx)
+	if !ok {
+		return w.parent.Enabled(ctx, level)
+	}
+	return level >= minLevel
+}
+
+func (w *wrapper) WithAttrs(attrs []slog.Attr) slog.Handler {
+	return &wrapper{
+		parent: w.parent.WithAttrs(attrs),
+	}
+}
+
+func (w *wrapper) WithGroup(name string) slog.Handler {
+	return &wrapper{
+		parent: w.parent.WithGroup(name),
+	}
+}
+
 type mergedAttrs struct {
 	parent *mergedAttrs
 	args   []any
@@ -105,14 +148,20 @@ func (attrs *mergedAttrs) addToRecord(record *slog.Record) {
 // New returns a new slog.Handler that injects the attributes from the context.
 func New(parent slog.Handler) slog.Handler {
 	return &wrapper{
-		handler: inject,
-		parent:  parent,
+		parent: parent,
 	}
 }
 
-func inject(ctx context.Context, parent func(ctx context.Context, record slog.Record) error, record slog.Record) error {
-	attrs := contextAttrs(ctx)
-	newRecord := record.Clone()
-	attrs.addToRecord(&newRecord)
-	return parent(ctx, newRecord)
+// WithLevel configures the minimum record level that will be logged in the context.
+func WithLevel(ctx context.Context, level slog.Level) context.Context {
+	return context.WithValue(ctx, levelKey, level)
+}
+
+// contextLevel returns the minimum record level that will be logged in the context.
+func contextLevel(ctx context.Context) (slog.Level, bool) {
+	level := ctx.Value(levelKey)
+	if level == nil {
+		return slog.LevelInfo, false
+	}
+	return level.(slog.Level), true
 }

context_test.go

@@ -7,10 +7,22 @@ import (
 	"io"
 	"log/slog"
 	"runtime"
-	"strings"
 	"testing"
 )
 
+var optsRemoveTime = &slog.HandlerOptions{
+	ReplaceAttr: removeTime,
+}
+
+// removeTime removes time attribute for stable testing.
+func removeTime(groups []string, a slog.Attr) slog.Attr {
+	// Remove time.
+	if a.Key == slog.TimeKey && len(groups) == 0 {
+		return slog.Attr{}
+	}
+	return a
+}
+
 func TestString(t *testing.T) {
 	ctx := context.Background()
 	ctx = With(ctx, "hello", 1)
@@ -23,7 +35,7 @@ func TestString(t *testing.T) {
 
 func TestWith(t *testing.T) {
 	buf := &bytes.Buffer{}
-	parent := slog.NewTextHandler(buf, nil)
+	parent := slog.NewTextHandler(buf, optsRemoveTime)
 	child := New(parent)
 	logger := slog.New(child)
 
@@ -32,54 +44,76 @@ func TestWith(t *testing.T) {
 	ctx = With(ctx, "world", 2)
 	ctx = With(ctx)
 	logger.InfoContext(ctx, "hello", "count", 42)
-	if !strings.HasSuffix(buf.String(), " level=INFO msg=hello count=42 hello=1 world=2\n") {
+	if buf.String() != "level=INFO msg=hello count=42 hello=1 world=2\n" {
 		t.Errorf("unexpected output: %q", buf.String())
 	}
 }
 
 func TestWithAttrs(t *testing.T) {
 	buf := &bytes.Buffer{}
-	parent := slog.NewTextHandler(buf, nil)
+	parent := slog.NewTextHandler(buf, optsRemoveTime)
 	child := New(parent)
 	logger := slog.New(child)
 
 	ctx := context.Background()
 	ctx = WithAttrs(ctx)
 	ctx = WithAttrs(ctx, slog.Int("hello", 1), slog.Int("world", 2))
 	logger.InfoContext(ctx, "hello", "count", 42)
-	if !strings.HasSuffix(buf.String(), " level=INFO msg=hello count=42 hello=1 world=2\n") {
+	if buf.String() != "level=INFO msg=hello count=42 hello=1 world=2\n" {
 		t.Errorf("unexpected output: %q", buf.String())
 	}
 }
 
 func TestHandlerWithAttrs(t *testing.T) {
 	buf := &bytes.Buffer{}
-	parent := slog.NewTextHandler(buf, nil)
+	parent := slog.NewTextHandler(buf, optsRemoveTime)
 	child := New(parent).WithAttrs([]slog.Attr{slog.Int("hello", 1)})
 	logger := slog.New(child)
 
 	ctx := context.Background()
 	ctx = WithAttrs(ctx, slog.Int("world", 2))
 	logger.InfoContext(ctx, "hello", "count", 42)
-	if !strings.HasSuffix(buf.String(), " level=INFO msg=hello hello=1 count=42 world=2\n") {
+	if buf.String() != "level=INFO msg=hello hello=1 count=42 world=2\n" {
 		t.Errorf("unexpected output: %q", buf.String())
 	}
 }
 
 func TestHandlerWithGroup(t *testing.T) {
 	buf := &bytes.Buffer{}
-	parent := slog.NewTextHandler(buf, nil)
+	parent := slog.NewTextHandler(buf, optsRemoveTime)
 	child := New(parent).WithGroup("my_group")
 	logger := slog.New(child)
 
 	ctx := context.Background()
 	ctx = WithAttrs(ctx, slog.Int("hello", 1), slog.Int("world", 2))
 	logger.InfoContext(ctx, "hello", "count", 42)
-	if !strings.HasSuffix(buf.String(), " level=INFO msg=hello my_group.count=42 my_group.hello=1 my_group.world=2\n") {
+	if buf.String() != "level=INFO msg=hello my_group.count=42 my_group.hello=1 my_group.world=2\n" {
 		t.Errorf("unexpected output: %q", buf.String())
 	}
 }
 
+func TestWithLevel(t *testing.T) {
+	buf := &bytes.Buffer{}
+	parent := slog.NewTextHandler(buf, optsRemoveTime)
+	child := New(parent)
+	logger := slog.New(child)
+
+	ctx := context.Background()
+	logger.InfoContext(ctx, "this log should be logged")
+	logger.DebugContext(ctx, "this log should not be logged")
+
+	ctx = WithLevel(ctx, slog.LevelDebug) // enable debug level
+	logger.DebugContext(ctx, "this log should be logged in this context")
+
+	got := buf.String()
+	want := `level=INFO msg="this log should be logged"
+level=DEBUG msg="this log should be logged in this context"
+`
+	if got != want {
+		t.Errorf("unexpected output: got %q, want %q", got, want)
+	}
+}
+
 func BenchmarkWith(b *testing.B) {
 	ctx := context.Background()
 	for i := 0; i < b.N; i++ {