Fix lambda url handling

Author: yacchi

aws/adaptor.go

@@ -12,6 +12,7 @@ import (
 )
 
 var DEBUGDumpPayload = os.Getenv("DEBUG_DUMP_PAYLOAD")
+var LambdaInvokeMode = os.Getenv("LAMBDA_INVOKE_MODE")
 
 type LambdaIntegrationType int
 

aws/lambda.go

@@ -45,14 +45,15 @@ func WithoutNonHTTPEventPassThrough() LambdaHandlerOption {
 }
 
 type LambdaHandler struct {
-	httpHandler      http.Handler
-	sessProv         SDKSessionProvider
-	sess             *session.Session
-	confProv         SDKConfigProvider
-	conf             *aws.Config
-	apiGW            APIGatewayManagementAPI
-	wsPathPrefix     string
-	nonHTTPEventPath string
+	httpHandler            http.Handler
+	sessProv               SDKSessionProvider
+	sess                   *session.Session
+	confProv               SDKConfigProvider
+	conf                   *aws.Config
+	apiGW                  APIGatewayManagementAPI
+	wsPathPrefix           string
+	nonHTTPEventPath       string
+	invokeLambdaWithStream func(ctx context.Context, request *events.LambdaFunctionURLRequest) (*events.LambdaFunctionURLStreamingResponse, error)
 }
 
 type HandlerFunc func(ctx context.Context, payload json.RawMessage) (res any, err error)
@@ -63,8 +64,9 @@ func NewLambdaHandlerWithOption(h http.Handler, options []interface{}) *LambdaHa
 		confProv: func(ctx context.Context) (aws.Config, error) {
 			return config.LoadDefaultConfig(ctx)
 		},
-		wsPathPrefix:     DefaultWebsocketPathPrefix,
-		nonHTTPEventPath: DefaultNonHTTPEventPath,
+		wsPathPrefix:           DefaultWebsocketPathPrefix,
+		nonHTTPEventPath:       DefaultNonHTTPEventPath,
+		invokeLambdaWithStream: lambdaurl.Wrap(h),
 	}
 
 	for _, opt := range options {
@@ -213,11 +215,19 @@ func (l *LambdaHandler) Invoke(ctx context.Context, payload json.RawMessage) (re
 			}
 			res, err = l.InvokeWebsocketAPI(ctx, event)
 		case LambdaFunctionURLIntegration:
-			event := &events.LambdaFunctionURLRequest{}
-			if err := json.Unmarshal(payload, event); err != nil {
-				return nil, err
+			if LambdaInvokeMode == "response_stream" {
+				event := &events.LambdaFunctionURLRequest{}
+				if err := json.Unmarshal(payload, event); err != nil {
+					return nil, err
+				}
+				res, err = l.invokeLambdaWithStream(ctx, event)
+			} else {
+				event := &events.APIGatewayV2HTTPRequest{}
+				if err := json.Unmarshal(payload, event); err != nil {
+					return nil, err
+				}
+				res, err = l.InvokeHTTPAPI(ctx, event)
 			}
-			res, err = lambdaurl.Wrap(l.httpHandler)(ctx, event)
 		default:
 			res, err = l.HandleNonHTTPEvent(ctx, payload, "application/json")
 		}