Address PR review feedback

anzheyazzz · anzheyazzz · commit a4e53b235727 · 2026-01-14T23:19:20.000Z
diff --git a/lambdacontext/example_logger_test.go b/lambdacontext/example_logger_test.go
@@ -11,11 +11,11 @@ import (
 	"github.com/aws/aws-lambda-go/lambdacontext"
 )
 
-// ExampleLogHandler demonstrates basic usage of LogHandler for structured logging.
-// The handler automatically injects requestId from Lambda context into each log record.
-func ExampleLogHandler() {
-	// Set up the Lambda-aware slog handler
-	slog.SetDefault(slog.New(lambdacontext.LogHandler()))
+// ExampleNewLogger demonstrates the simplest usage of NewLogger for structured logging.
+// The logger automatically injects requestId from Lambda context into each log record.
+func ExampleNewLogger() {
+	// Set up the Lambda-aware slog logger
+	slog.SetDefault(lambdacontext.NewLogger())
 
 	lambda.Start(func(ctx context.Context) (string, error) {
 		// Use slog.InfoContext to include Lambda context in logs
@@ -24,12 +24,24 @@ func ExampleLogHandler() {
 	})
 }
 
-// ExampleLogHandler_withFields demonstrates LogHandler with additional fields.
-// Use WithFields with FieldFunctionARN() and FieldTenantID() to include extra context.
-func ExampleLogHandler_withFields() {
+// ExampleNewLogHandler demonstrates using NewLogHandler for more control.
+func ExampleNewLogHandler() {
+	// Set up the Lambda-aware slog handler
+	slog.SetDefault(slog.New(lambdacontext.NewLogHandler()))
+
+	lambda.Start(func(ctx context.Context) (string, error) {
+		slog.InfoContext(ctx, "processing request", "action", "example")
+		return "success", nil
+	})
+}
+
+// ExampleNewLogHandler_withOptions demonstrates NewLogHandler with additional fields.
+// Use WithFunctionARN() and WithTenantID() to include extra context.
+func ExampleNewLogHandler_withOptions() {
 	// Set up handler with function ARN and tenant ID fields
-	slog.SetDefault(slog.New(lambdacontext.LogHandler(
-		lambdacontext.WithFields(lambdacontext.FieldFunctionARN(), lambdacontext.FieldTenantID()),
+	slog.SetDefault(slog.New(lambdacontext.NewLogHandler(
+		lambdacontext.WithFunctionARN(),
+		lambdacontext.WithTenantID(),
 	)))
 
 	lambda.Start(func(ctx context.Context) (string, error) {
@@ -38,13 +50,12 @@ func ExampleLogHandler_withFields() {
 	})
 }
 
-// ExampleWithFields demonstrates using WithFields to include specific Lambda context fields.
-func ExampleWithFields() {
+// ExampleWithFunctionARN demonstrates using WithFunctionARN to include the function ARN.
+func ExampleWithFunctionARN() {
 	// Include only function ARN
-	handler := lambdacontext.LogHandler(
-		lambdacontext.WithFields(lambdacontext.FieldFunctionARN()),
-	)
-	slog.SetDefault(slog.New(handler))
+	slog.SetDefault(lambdacontext.NewLogger(
+		lambdacontext.WithFunctionARN(),
+	))
 
 	lambda.Start(func(ctx context.Context) (string, error) {
 		// Log output will include "functionArn" field
diff --git a/lambdacontext/logger.go b/lambdacontext/logger.go
@@ -11,44 +11,41 @@ import (
 	"os"
 )
 
-// Field represents a Lambda context field to include in log records.
-type Field struct {
+// field represents a Lambda context field to include in log records.
+type field struct {
 	key   string
 	value func(*LambdaContext) string
 }
 
-// FieldFunctionARN returns a Field that includes the invoked function ARN in log records.
-func FieldFunctionARN() Field {
-	return Field{"functionArn", func(lc *LambdaContext) string { return lc.InvokedFunctionArn }}
-}
-
-// FieldTenantID returns a Field that includes the tenant ID in log records (for multi-tenant functions).
-func FieldTenantID() Field {
-	return Field{"tenantId", func(lc *LambdaContext) string { return lc.TenantID }}
-}
-
 // logOptions holds configuration for the Lambda log handler.
 type logOptions struct {
-	fields []Field
+	fields []field
 }
 
 // LogOption is a functional option for configuring the Lambda log handler.
 type LogOption func(*logOptions)
 
-// WithFields includes the specified fields in log records.
-func WithFields(fields ...Field) LogOption {
+// WithFunctionARN includes the invoked function ARN in log records.
+func WithFunctionARN() LogOption {
+	return func(o *logOptions) {
+		o.fields = append(o.fields, field{"functionArn", func(lc *LambdaContext) string { return lc.InvokedFunctionArn }})
+	}
+}
+
+// WithTenantID includes the tenant ID in log records (for multi-tenant functions).
+func WithTenantID() LogOption {
 	return func(o *logOptions) {
-		o.fields = append(o.fields, fields...)
+		o.fields = append(o.fields, field{"tenantId", func(lc *LambdaContext) string { return lc.TenantID }})
 	}
 }
 
-// LogHandler returns a [slog.Handler] for AWS Lambda structured logging.
+// NewLogHandler returns a [slog.Handler] for AWS Lambda structured logging.
 // It reads AWS_LAMBDA_LOG_FORMAT and AWS_LAMBDA_LOG_LEVEL from environment,
 // and injects requestId from Lambda context into each log record.
 //
-// By default, only requestId is injected. Use WithFields to include more.
+// By default, only requestId is injected. Use WithFunctionARN or WithTenantID to include more.
 // See the package examples for usage.
-func LogHandler(opts ...LogOption) slog.Handler {
+func NewLogHandler(opts ...LogOption) slog.Handler {
 	options := &logOptions{}
 	for _, opt := range opts {
 		opt(options)
@@ -70,6 +67,12 @@ func LogHandler(opts ...LogOption) slog.Handler {
 	return &lambdaHandler{handler: h, fields: options.fields}
 }
 
+// NewLogger returns a [*slog.Logger] configured for AWS Lambda structured logging.
+// This is a convenience function equivalent to slog.New(NewLogHandler(opts...)).
+func NewLogger(opts ...LogOption) *slog.Logger {
+	return slog.New(NewLogHandler(opts...))
+}
+
 // ReplaceAttr maps slog's default keys to AWS Lambda's log format (time->timestamp, msg->message).
 func ReplaceAttr(groups []string, attr slog.Attr) slog.Attr {
 	if len(groups) > 0 {
@@ -85,16 +88,10 @@ func ReplaceAttr(groups []string, attr slog.Attr) slog.Attr {
 	return attr
 }
 
-// Attrs returns Lambda context fields as slog-compatible key-value pairs.
-// For most use cases, using [LogHandler] with slog.InfoContext is preferred.
-func (lc *LambdaContext) Attrs() []any {
-	return []any{"requestId", lc.AwsRequestID}
-}
-
 // lambdaHandler wraps a slog.Handler to inject Lambda context fields.
 type lambdaHandler struct {
 	handler slog.Handler
-	fields  []Field
+	fields  []field
 }
 
 // Enabled implements slog.Handler.
diff --git a/lambdacontext/logger_test.go b/lambdacontext/logger_test.go
@@ -16,14 +16,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestAttrs(t *testing.T) {
-	lc := &LambdaContext{
-		AwsRequestID: "test-request-id",
-	}
-	attrs := lc.Attrs()
-	assert.Equal(t, []any{"requestId", "test-request-id"}, attrs)
-}
-
 func TestReplaceAttr(t *testing.T) {
 	tests := []struct {
 		name     string
@@ -268,9 +260,15 @@ func TestLogHandler_WithFields(t *testing.T) {
 		ReplaceAttr: ReplaceAttr,
 	}
 	baseHandler := slog.NewJSONHandler(&buf, opts)
+
+	// Create options with fields
+	options := &logOptions{}
+	WithFunctionARN()(options)
+	WithTenantID()(options)
+
 	handler := &lambdaHandler{
 		handler: baseHandler,
-		fields:  []Field{FieldFunctionARN(), FieldTenantID()},
+		fields:  options.fields,
 	}
 
 	lc := &LambdaContext{
@@ -300,9 +298,13 @@ func TestLogHandler_WithFieldFunctionARNOnly(t *testing.T) {
 		ReplaceAttr: ReplaceAttr,
 	}
 	baseHandler := slog.NewJSONHandler(&buf, opts)
+
+	options := &logOptions{}
+	WithFunctionARN()(options)
+
 	handler := &lambdaHandler{
 		handler: baseHandler,
-		fields:  []Field{FieldFunctionARN()},
+		fields:  options.fields,
 	}
 
 	lc := &LambdaContext{
@@ -332,9 +334,14 @@ func TestLogHandler_FieldsEmpty(t *testing.T) {
 		ReplaceAttr: ReplaceAttr,
 	}
 	baseHandler := slog.NewJSONHandler(&buf, opts)
+
+	options := &logOptions{}
+	WithFunctionARN()(options)
+	WithTenantID()(options)
+
 	handler := &lambdaHandler{
 		handler: baseHandler,
-		fields:  []Field{FieldFunctionARN(), FieldTenantID()},
+		fields:  options.fields,
 	}
 
 	lc := &LambdaContext{
@@ -356,26 +363,43 @@ func TestLogHandler_FieldsEmpty(t *testing.T) {
 	assert.NotContains(t, logOutput, "tenantId")
 }
 
-func TestWithFields(t *testing.T) {
+func TestWithFunctionARN(t *testing.T) {
 	options := &logOptions{}
-	WithFields(FieldFunctionARN(), FieldTenantID())(options)
+	WithFunctionARN()(options)
 
-	assert.Len(t, options.fields, 2)
+	assert.Len(t, options.fields, 1)
 	assert.Equal(t, "functionArn", options.fields[0].key)
-	assert.Equal(t, "tenantId", options.fields[1].key)
+
+	lc := &LambdaContext{InvokedFunctionArn: "arn:aws:lambda:us-east-1:123456789:function:test"}
+	assert.Equal(t, "arn:aws:lambda:us-east-1:123456789:function:test", options.fields[0].value(lc))
 }
 
-func TestFieldFunctions(t *testing.T) {
-	lc := &LambdaContext{
-		InvokedFunctionArn: "arn:aws:lambda:us-east-1:123456789:function:test",
-		TenantID:           "tenant-abc",
-	}
+func TestWithTenantID(t *testing.T) {
+	options := &logOptions{}
+	WithTenantID()(options)
+
+	assert.Len(t, options.fields, 1)
+	assert.Equal(t, "tenantId", options.fields[0].key)
+
+	lc := &LambdaContext{TenantID: "tenant-abc"}
+	assert.Equal(t, "tenant-abc", options.fields[0].value(lc))
+}
+
+func TestNewLogger(t *testing.T) {
+	LogFormat = "JSON"
+	LogLevel = "INFO"
+
+	logger := NewLogger()
+	assert.NotNil(t, logger)
+}
+
+func TestNewLogHandler(t *testing.T) {
+	LogFormat = "JSON"
+	LogLevel = "INFO"
 
-	arnField := FieldFunctionARN()
-	assert.Equal(t, "functionArn", arnField.key)
-	assert.Equal(t, "arn:aws:lambda:us-east-1:123456789:function:test", arnField.value(lc))
+	handler := NewLogHandler()
+	assert.NotNil(t, handler)
 
-	tenantField := FieldTenantID()
-	assert.Equal(t, "tenantId", tenantField.key)
-	assert.Equal(t, "tenant-abc", tenantField.value(lc))
+	handlerWithOpts := NewLogHandler(WithFunctionARN(), WithTenantID())
+	assert.NotNil(t, handlerWithOpts)
 }
