wip

Author: yacchi

aws/adaptor.go

@@ -4,12 +4,15 @@ import (
 	"context"
 	"fmt"
 	"github.com/aws/aws-lambda-go/lambda"
+	"github.com/aws/aws-lambda-go/lambda/handlertrace"
 	"github.com/yacchi/lambda-http-adaptor/registry"
 	"github.com/yacchi/lambda-http-adaptor/types"
 	"net/http"
 	"os"
 )
 
+var DEBUGDumpPayload = os.Getenv("DEBUG_DUMP_PAYLOAD")
+
 type LambdaIntegrationType int
 
 const (
@@ -79,11 +82,22 @@ func LambdaDetector() bool {
 }
 
 type LambdaAdaptor struct {
-	h lambda.Handler
+	h *LambdaHandler
 }
 
 func (l LambdaAdaptor) ListenAndServe() error {
-	lambda.Start(l.h)
+	ctx := context.Background()
+	if DEBUGDumpPayload != "" && (DEBUGDumpPayload == "1" || DEBUGDumpPayload == "true") {
+		ctx = handlertrace.NewContext(ctx, handlertrace.HandlerTrace{
+			RequestEvent: func(ctx context.Context, payload interface{}) {
+				fmt.Printf("Request payload: %s\n", payload)
+			},
+			ResponseEvent: func(ctx context.Context, payload interface{}) {
+				fmt.Printf("Response payload: %s\n", payload)
+			},
+		})
+	}
+	lambda.StartHandlerFunc(l.h.Invoke, lambda.WithContext(ctx))
 	return nil
 }
 

aws/lambda.go

@@ -5,18 +5,13 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/aws/aws-lambda-go/events"
-	"github.com/aws/aws-lambda-go/lambda"
-	"github.com/aws/aws-lambda-go/lambda/handlertrace"
 	"github.com/aws/aws-lambda-go/lambdaurl"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"net/http"
-	"os"
 )
 
-var DEBUGDumpPayload = os.Getenv("DEBUG_DUMP_PAYLOAD")
-
 const DefaultNonHTTPEventPath = "/events"
 
 type LambdaHandlerOption func(handler *LambdaHandler)
@@ -60,7 +55,9 @@ type LambdaHandler struct {
 	nonHTTPEventPath string
 }
 
-func NewLambdaHandlerWithOption(h http.Handler, options []interface{}) lambda.Handler {
+type HandlerFunc func(ctx context.Context, payload json.RawMessage) (res any, err error)
+
+func NewLambdaHandlerWithOption(h http.Handler, options []interface{}) *LambdaHandler {
 	handler := &LambdaHandler{
 		httpHandler: h,
 		confProv: func(ctx context.Context) (aws.Config, error) {
@@ -79,7 +76,7 @@ func NewLambdaHandlerWithOption(h http.Handler, options []interface{}) lambda.Ha
 	return handler
 }
 
-func NewLambdaHandler(h http.Handler) lambda.Handler {
+func NewLambdaHandler(h http.Handler) *LambdaHandler {
 	return NewLambdaHandlerWithOption(h, nil)
 }
 
@@ -182,30 +179,11 @@ func (l *LambdaHandler) InvokeWebsocketAPI(ctx context.Context, request *events.
 	}
 }
 
-func (l *LambdaHandler) Invoke(ctx context.Context, payload []byte) ([]byte, error) {
-	if DEBUGDumpPayload != "" && (DEBUGDumpPayload == "1" || DEBUGDumpPayload == "true") {
-		ctx = handlertrace.NewContext(ctx, handlertrace.HandlerTrace{
-			RequestEvent: func(ctx context.Context, payload interface{}) {
-				fmt.Printf("Request payload: %s\n", payload)
-			},
-			ResponseEvent: func(ctx context.Context, payload interface{}) {
-				fmt.Printf("Response payload: %s\n", payload)
-			},
-		})
-	}
-
-	trace := handlertrace.FromContext(ctx)
-
+func (l *LambdaHandler) Invoke(ctx context.Context, payload json.RawMessage) (res any, err error) {
 	var (
 		checker integrationTypeChecker
-		res     interface{}
-		err     error
 	)
 
-	if trace.RequestEvent != nil {
-		trace.RequestEvent(ctx, payload)
-	}
-
 	if err = json.Unmarshal(payload, &checker); err != nil {
 		res, err = l.HandleNonHTTPEvent(ctx, payload, http.DetectContentType(payload))
 	} else {
@@ -245,23 +223,5 @@ func (l *LambdaHandler) Invoke(ctx context.Context, payload []byte) ([]byte, err
 		}
 	}
 
-	if err != nil {
-		return nil, err
-	}
-
-	if b, ok := res.([]byte); ok {
-		if trace.ResponseEvent != nil {
-			trace.ResponseEvent(ctx, b)
-		}
-		return b, nil
-	} else {
-		if responseBytes, err := json.Marshal(res); err != nil {
-			return nil, err
-		} else {
-			if trace.ResponseEvent != nil {
-				trace.ResponseEvent(ctx, responseBytes)
-			}
-			return responseBytes, nil
-		}
-	}
+	return res, err
 }