Skip to content

Commit 0dae564

Browse files
authored
Add capability to define optional handler behavior (#444)
* Add capability to define optional handler behavior * add Option WithSetEscapeHTML to set this same named option on the underlying json encoder * refactor implementations of ***WithContext functions to use the Option WithContext * mark a bunch of not useful related stuff as deprecated so that they are hidden in the godocs * private variable renames * add some usage examples * variable rename
1 parent ec8f96b commit 0dae564

8 files changed

+227
-94
lines changed

lambda/entry.go

+22-6
Original file line numberDiff line numberDiff line change
@@ -37,12 +37,14 @@ import (
3737
// Where "TIn" and "TOut" are types compatible with the "encoding/json" standard library.
3838
// See https://golang.org/pkg/encoding/json/#Unmarshal for how deserialization behaves
3939
func Start(handler interface{}) {
40-
StartWithContext(context.Background(), handler)
40+
StartWithOptions(handler)
4141
}
4242

4343
// StartWithContext is the same as Start except sets the base context for the function.
44+
//
45+
// Deprecated: use lambda.StartWithOptions(handler, lambda.WithContext(ctx)) instead
4446
func StartWithContext(ctx context.Context, handler interface{}) {
45-
StartHandlerWithContext(ctx, NewHandler(handler))
47+
StartWithOptions(handler, WithContext(ctx))
4648
}
4749

4850
// StartHandler takes in a Handler wrapper interface which can be implemented either by a
@@ -51,13 +53,20 @@ func StartWithContext(ctx context.Context, handler interface{}) {
5153
// Handler implementation requires a single "Invoke()" function:
5254
//
5355
// func Invoke(context.Context, []byte) ([]byte, error)
56+
//
57+
// Deprecated: use lambda.Start(handler) instead
5458
func StartHandler(handler Handler) {
55-
StartHandlerWithContext(context.Background(), handler)
59+
StartWithOptions(handler)
60+
}
61+
62+
// StartWithOptions is the same as Start after the application of any handler options specified
63+
func StartWithOptions(handler interface{}, options ...Option) {
64+
start(newHandler(handler, options...))
5665
}
5766

5867
type startFunction struct {
5968
env string
60-
f func(ctx context.Context, envValue string, handler Handler) error
69+
f func(envValue string, handler Handler) error
6170
}
6271

6372
var (
@@ -66,7 +75,7 @@ var (
6675
// To drop the rpc dependencies, compile with `-tags lambda.norpc`
6776
rpcStartFunction = &startFunction{
6877
env: "_LAMBDA_SERVER_PORT",
69-
f: func(c context.Context, p string, h Handler) error {
78+
f: func(_ string, _ Handler) error {
7079
return errors.New("_LAMBDA_SERVER_PORT was present but the function was compiled without RPC support")
7180
},
7281
}
@@ -85,17 +94,24 @@ var (
8594
// Handler implementation requires a single "Invoke()" function:
8695
//
8796
// func Invoke(context.Context, []byte) ([]byte, error)
97+
//
98+
// Deprecated: use lambda.StartWithOptions(handler, lambda.WithContext(ctx)) instead
8899
func StartHandlerWithContext(ctx context.Context, handler Handler) {
100+
StartWithOptions(handler, WithContext(ctx))
101+
}
102+
103+
func start(handler *handlerOptions) {
89104
var keys []string
90105
for _, start := range startFunctions {
91106
config := os.Getenv(start.env)
92107
if config != "" {
93108
// in normal operation, the start function never returns
94109
// if it does, exit!, this triggers a restart of the lambda function
95-
err := start.f(ctx, config, handler)
110+
err := start.f(config, handler)
96111
logFatalf("%v", err)
97112
}
98113
keys = append(keys, start.env)
99114
}
100115
logFatalf("expected AWS Lambda environment variables %s are not defined", keys)
116+
101117
}

lambda/function.go

+11-25
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,17 @@ import (
1313
)
1414

1515
// Function struct which wrap the Handler
16+
//
17+
// Deprecated: The Function type is public for the go1.x runtime internal use of the net/rpc package
1618
type Function struct {
17-
handler Handler
18-
ctx context.Context
19+
handler *handlerOptions
1920
}
2021

2122
// NewFunction which creates a Function with a given Handler
23+
//
24+
// Deprecated: The Function type is public for the go1.x runtime internal use of the net/rpc package
2225
func NewFunction(handler Handler) *Function {
23-
return &Function{handler: handler}
26+
return &Function{newHandler(handler)}
2427
}
2528

2629
// Ping method which given a PingRequest and a PingResponse parses the PingResponse
@@ -38,7 +41,7 @@ func (fn *Function) Invoke(req *messages.InvokeRequest, response *messages.Invok
3841
}()
3942

4043
deadline := time.Unix(req.Deadline.Seconds, req.Deadline.Nanos).UTC()
41-
invokeContext, cancel := context.WithDeadline(fn.context(), deadline)
44+
invokeContext, cancel := context.WithDeadline(fn.baseContext(), deadline)
4245
defer cancel()
4346

4447
lc := &lambdacontext.LambdaContext{
@@ -70,26 +73,9 @@ func (fn *Function) Invoke(req *messages.InvokeRequest, response *messages.Invok
7073
return nil
7174
}
7275

73-
// context returns the base context used for the fn.
74-
func (fn *Function) context() context.Context {
75-
if fn.ctx == nil {
76-
return context.Background()
76+
func (fn *Function) baseContext() context.Context {
77+
if fn.handler.baseContext != nil {
78+
return fn.handler.baseContext
7779
}
78-
79-
return fn.ctx
80-
}
81-
82-
// withContext returns a shallow copy of Function with its context changed
83-
// to the provided ctx. If the provided ctx is non-nil a Background context is set.
84-
func (fn *Function) withContext(ctx context.Context) *Function {
85-
if ctx == nil {
86-
ctx = context.Background()
87-
}
88-
89-
fn2 := new(Function)
90-
*fn2 = *fn
91-
92-
fn2.ctx = ctx
93-
94-
return fn2
80+
return context.Background()
9581
}

lambda/function_test.go

+21-20
Original file line numberDiff line numberDiff line change
@@ -36,14 +36,14 @@ func (h testWrapperHandler) Invoke(ctx context.Context, payload []byte) ([]byte,
3636
var _ Handler = (testWrapperHandler)(nil)
3737

3838
func TestInvoke(t *testing.T) {
39-
srv := &Function{handler: testWrapperHandler(
39+
srv := NewFunction(testWrapperHandler(
4040
func(ctx context.Context, input []byte) (interface{}, error) {
4141
if deadline, ok := ctx.Deadline(); ok {
4242
return deadline.UnixNano(), nil
4343
}
4444
return nil, errors.New("!?!?!?!?!")
4545
},
46-
)}
46+
))
4747
deadline := time.Now()
4848
var response messages.InvokeResponse
4949
err := srv.Invoke(&messages.InvokeRequest{
@@ -59,15 +59,17 @@ func TestInvoke(t *testing.T) {
5959

6060
func TestInvokeWithContext(t *testing.T) {
6161
key := struct{}{}
62-
srv := NewFunction(testWrapperHandler(
63-
func(ctx context.Context, input []byte) (interface{}, error) {
64-
assert.Equal(t, "dummy", ctx.Value(key))
65-
if deadline, ok := ctx.Deadline(); ok {
66-
return deadline.UnixNano(), nil
67-
}
68-
return nil, errors.New("!?!?!?!?!")
69-
}))
70-
srv = srv.withContext(context.WithValue(context.Background(), key, "dummy"))
62+
srv := NewFunction(&handlerOptions{
63+
Handler: testWrapperHandler(
64+
func(ctx context.Context, input []byte) (interface{}, error) {
65+
assert.Equal(t, "dummy", ctx.Value(key))
66+
if deadline, ok := ctx.Deadline(); ok {
67+
return deadline.UnixNano(), nil
68+
}
69+
return nil, errors.New("!?!?!?!?!")
70+
}),
71+
baseContext: context.WithValue(context.Background(), key, "dummy"),
72+
})
7173
deadline := time.Now()
7274
var response messages.InvokeResponse
7375
err := srv.Invoke(&messages.InvokeRequest{
@@ -86,12 +88,11 @@ type CustomError struct{}
8688
func (e CustomError) Error() string { return "Something bad happened!" }
8789

8890
func TestCustomError(t *testing.T) {
89-
90-
srv := &Function{handler: testWrapperHandler(
91+
srv := NewFunction(testWrapperHandler(
9192
func(ctx context.Context, input []byte) (interface{}, error) {
9293
return nil, CustomError{}
9394
},
94-
)}
95+
))
9596
var response messages.InvokeResponse
9697
err := srv.Invoke(&messages.InvokeRequest{}, &response)
9798
assert.NoError(t, err)
@@ -106,11 +107,11 @@ func (e *CustomError2) Error() string { return "Something bad happened!" }
106107

107108
func TestCustomErrorRef(t *testing.T) {
108109

109-
srv := &Function{handler: testWrapperHandler(
110+
srv := NewFunction(testWrapperHandler(
110111
func(ctx context.Context, input []byte) (interface{}, error) {
111112
return nil, &CustomError2{}
112113
},
113-
)}
114+
))
114115
var response messages.InvokeResponse
115116
err := srv.Invoke(&messages.InvokeRequest{}, &response)
116117
assert.NoError(t, err)
@@ -120,12 +121,12 @@ func TestCustomErrorRef(t *testing.T) {
120121
}
121122

122123
func TestContextPlumbing(t *testing.T) {
123-
srv := &Function{handler: testWrapperHandler(
124+
srv := NewFunction(testWrapperHandler(
124125
func(ctx context.Context, input []byte) (interface{}, error) {
125126
lc, _ := lambdacontext.FromContext(ctx)
126127
return lc, nil
127128
},
128-
)}
129+
))
129130
var response messages.InvokeResponse
130131
err := srv.Invoke(&messages.InvokeRequest{
131132
CognitoIdentityId: "dummyident",
@@ -172,14 +173,14 @@ func TestXAmznTraceID(t *testing.T) {
172173
Ctx string
173174
}
174175

175-
srv := &Function{handler: testWrapperHandler(
176+
srv := NewFunction(testWrapperHandler(
176177
func(ctx context.Context, input []byte) (interface{}, error) {
177178
return &XRayResponse{
178179
Env: os.Getenv("_X_AMZN_TRACE_ID"),
179180
Ctx: ctx.Value("x-amzn-trace-id").(string),
180181
}, nil
181182
},
182-
)}
183+
))
183184

184185
sequence := []struct {
185186
Input string

0 commit comments

Comments
 (0)