Pass context.Context from Lambda runtime to http.Request.

Fixes awslabs#27
snabb · Feb 7, 2019 · 0057c9d · 0057c9d
1 parent 5e02399
commit 0057c9d
Show file tree

Hide file tree

Showing 13 changed files with 35 additions and 22 deletions.
diff --git a/chi/adapter.go b/chi/adapter.go
@@ -4,6 +4,7 @@
 package chiadapter
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -30,15 +31,15 @@ func New(chi *chi.Mux) *ChiLambda {
 // Proxy receives an API Gateway proxy event, transforms it into an http.Request
 // object, and sends it to the chi.Mux for routing.
 // It returns a proxy response object gneerated from the http.ResponseWriter.
-func (g *ChiLambda) Proxy(req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (g *ChiLambda) Proxy(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	chiRequest, err := g.ProxyEventToHTTPRequest(req)
 
 	if err != nil {
 		return core.GatewayTimeout(), core.NewLoggedError("Could not convert proxy event to request: %v", err)
 	}
 
 	respWriter := core.NewProxyResponseWriter()
-	g.chiMux.ServeHTTP(http.ResponseWriter(respWriter), chiRequest)
+	g.chiMux.ServeHTTP(http.ResponseWriter(respWriter), chiRequest.WithContext(ctx))
 
 	proxyResponse, err := respWriter.GetProxyResponse()
 	if err != nil {

diff --git a/chi/chilambda_test.go b/chi/chilambda_test.go
@@ -1,6 +1,7 @@
 package chiadapter_test
 
 import (
+	"context"
 	"log"
 	"net/http"
 
@@ -29,7 +30,7 @@ var _ = Describe("ChiLambda tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			resp, err := adapter.Proxy(req)
+			resp, err := adapter.Proxy(context.Background(), req)
 
 			Expect(err).To(BeNil())
 			Expect(resp.StatusCode).To(Equal(200))

diff --git a/gin/adapter.go b/gin/adapter.go
@@ -4,6 +4,7 @@
 package ginadapter
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -30,15 +31,15 @@ func New(gin *gin.Engine) *GinLambda {
 // Proxy receives an API Gateway proxy event, transforms it into an http.Request
 // object, and sends it to the gin.Engine for routing.
 // It returns a proxy response object gneerated from the http.ResponseWriter.
-func (g *GinLambda) Proxy(req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (g *GinLambda) Proxy(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	ginRequest, err := g.ProxyEventToHTTPRequest(req)
 
 	if err != nil {
 		return core.GatewayTimeout(), core.NewLoggedError("Could not convert proxy event to request: %v", err)
 	}
 
 	respWriter := core.NewProxyResponseWriter()
-	g.ginEngine.ServeHTTP(http.ResponseWriter(respWriter), ginRequest)
+	g.ginEngine.ServeHTTP(http.ResponseWriter(respWriter), ginRequest.WithContext(ctx))
 
 	proxyResponse, err := respWriter.GetProxyResponse()
 	if err != nil {

diff --git a/gin/ginlambda_test.go b/gin/ginlambda_test.go
@@ -1,6 +1,7 @@
 package ginadapter_test
 
 import (
+	"context"
 	"log"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -30,7 +31,7 @@ var _ = Describe("GinLambda tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			resp, err := adapter.Proxy(req)
+			resp, err := adapter.Proxy(context.Background(), req)
 
 			Expect(err).To(BeNil())
 			Expect(resp.StatusCode).To(Equal(200))

diff --git a/gorillamux/adapter.go b/gorillamux/adapter.go
@@ -1,6 +1,7 @@
 package gorillamux
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -19,14 +20,14 @@ func New(router *mux.Router) *GorillaMuxAdapter {
 	}
 }
 
-func (h *GorillaMuxAdapter) Proxy(event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (h *GorillaMuxAdapter) Proxy(ctx context.Context, event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	req, err := h.ProxyEventToHTTPRequest(event)
 	if err != nil {
 		return core.GatewayTimeout(), core.NewLoggedError("Could not convert proxy event to request: %v", err)
 	}
 
 	w := core.NewProxyResponseWriter()
-	h.router.ServeHTTP(http.ResponseWriter(w), req)
+	h.router.ServeHTTP(http.ResponseWriter(w), req.WithContext(ctx))
 
 	resp, err := w.GetProxyResponse()
 	if err != nil {

diff --git a/gorillamux/adapter_test.go b/gorillamux/adapter_test.go
@@ -1,6 +1,7 @@
 package gorillamux_test
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 
@@ -36,7 +37,7 @@ var _ = Describe("GorillaMuxAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			homePageResp, homePageReqErr := adapter.Proxy(homePageReq)
+			homePageResp, homePageReqErr := adapter.Proxy(context.Background(), homePageReq)
 
 			Expect(homePageReqErr).To(BeNil())
 			Expect(homePageResp.StatusCode).To(Equal(200))
@@ -47,7 +48,7 @@ var _ = Describe("GorillaMuxAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			productsPageResp, productsPageReqErr := adapter.Proxy(productsPageReq)
+			productsPageResp, productsPageReqErr := adapter.Proxy(context.Background(), productsPageReq)
 
 			Expect(productsPageReqErr).To(BeNil())
 			Expect(productsPageResp.StatusCode).To(Equal(200))

diff --git a/handlerfunc/adapter.go b/handlerfunc/adapter.go
@@ -1,6 +1,7 @@
 package handlerfunc
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -18,14 +19,14 @@ func New(handlerFunc http.HandlerFunc) *HandlerFuncAdapter {
 	}
 }
 
-func (h *HandlerFuncAdapter) Proxy(event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (h *HandlerFuncAdapter) Proxy(ctx context.Context, event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	req, err := h.ProxyEventToHTTPRequest(event)
 	if err != nil {
 		return core.GatewayTimeout(), core.NewLoggedError("Could not convert proxy event to request: %v", err)
 	}
 
 	w := core.NewProxyResponseWriter()
-	h.handlerFunc.ServeHTTP(http.ResponseWriter(w), req)
+	h.handlerFunc.ServeHTTP(http.ResponseWriter(w), req.WithContext(ctx))
 
 	resp, err := w.GetProxyResponse()
 	if err != nil {

diff --git a/handlerfunc/adapter_test.go b/handlerfunc/adapter_test.go
@@ -1,6 +1,7 @@
 package handlerfunc_test
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -29,7 +30,7 @@ var _ = Describe("HandlerFuncAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			resp, err := adapter.Proxy(req)
+			resp, err := adapter.Proxy(context.Background(), req)
 
 			Expect(err).To(BeNil())
 			Expect(resp.StatusCode).To(Equal(200))

diff --git a/httpadapter/adapter.go b/httpadapter/adapter.go
@@ -1,6 +1,7 @@
 package httpadapter
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -18,14 +19,14 @@ func New(handler http.Handler) *HandlerAdapter {
 	}
 }
 
-func (h *HandlerAdapter) Proxy(event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (h *HandlerAdapter) Proxy(ctx context.Context, event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	req, err := h.ProxyEventToHTTPRequest(event)
 	if err != nil {
 		return core.GatewayTimeout(), core.NewLoggedError("Could not convert proxy event to request: %v", err)
 	}
 
 	w := core.NewProxyResponseWriter()
-	h.handler.ServeHTTP(http.ResponseWriter(w), req)
+	h.handler.ServeHTTP(http.ResponseWriter(w), req.WithContext(ctx))
 
 	resp, err := w.GetProxyResponse()
 	if err != nil {

diff --git a/httpadapter/adapter_test.go b/httpadapter/adapter_test.go
@@ -1,6 +1,7 @@
 package httpadapter_test
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -36,7 +37,7 @@ var _ = Describe("HTTPAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			resp, err := adapter.Proxy(req)
+			resp, err := adapter.Proxy(context.Background(), req)
 
 			Expect(err).To(BeNil())
 			Expect(resp.StatusCode).To(Equal(200))

diff --git a/negroni/adapter.go b/negroni/adapter.go
@@ -1,6 +1,7 @@
 package negroniadapter
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -19,14 +20,14 @@ func New(n *negroni.Negroni) *NegroniAdapter {
 	}
 }
 
-func (h *NegroniAdapter) Proxy(event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (h *NegroniAdapter) Proxy(ctx context.Context, event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	req, err := h.ProxyEventToHTTPRequest(event)
 	if err != nil {
 		return core.GatewayTimeout(), core.NewLoggedError("Could not convert proxy event to request: %v", err)
 	}
 
 	w := core.NewProxyResponseWriter()
-	h.n.ServeHTTP(http.ResponseWriter(w), req)
+	h.n.ServeHTTP(http.ResponseWriter(w), req.WithContext(ctx))
 
 	resp, err := w.GetProxyResponse()
 	if err != nil {

diff --git a/negroni/adapter_test.go b/negroni/adapter_test.go
@@ -1,6 +1,7 @@
 package negroniadapter_test
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -42,7 +43,7 @@ var _ = Describe("NegroniAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			homePageResp, homePageReqErr := adapter.Proxy(homePageReq)
+			homePageResp, homePageReqErr := adapter.Proxy(context.Background(), homePageReq)
 
 			Expect(homePageReqErr).To(BeNil())
 			Expect(homePageResp.StatusCode).To(Equal(200))
@@ -53,7 +54,7 @@ var _ = Describe("NegroniAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			productsPageResp, productsPageReqErr := adapter.Proxy(productsPageReq)
+			productsPageResp, productsPageReqErr := adapter.Proxy(context.Background(), productsPageReq)
 
 			Expect(productsPageReqErr).To(BeNil())
 			Expect(productsPageResp.StatusCode).To(Equal(200))

diff --git a/sample/main.go b/sample/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"log"
 	"net/http"
 	"strconv"
@@ -15,7 +16,7 @@ var ginLambda *ginadapter.GinLambda
 
 // Handler is the main entry point for Lambda. Receives a proxy request and
 // returns a proxy response
-func Handler(req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func Handler(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	if ginLambda == nil {
 		// stdout and stderr are sent to AWS CloudWatch Logs
 		log.Printf("Gin cold start")
@@ -27,7 +28,7 @@ func Handler(req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse,
 		ginLambda = ginadapter.New(r)
 	}
 
-	return ginLambda.Proxy(req)
+	return ginLambda.Proxy(ctx, req)
 }
 
 func main() {