Merge pull request #26 from shogo82148/inject-request-id

Inject request id

lib/AWS/Lambda.pm

@@ -6,6 +6,9 @@ use HTTP::Tiny;
 
 our $VERSION = "0.0.6";
 
+# the context of Lambda Function
+our $context;
+
 1;
 __END__
 

lib/AWS/Lambda/Bootstrap.pm

@@ -6,6 +6,7 @@ use warnings;
 use HTTP::Tiny;
 use JSON::XS qw/decode_json encode_json/;
 use Try::Tiny;
+use AWS::Lambda;
 use AWS::Lambda::Context;
 use Scalar::Util qw(blessed);
 
@@ -70,6 +71,7 @@ sub handle_event {
     $self->_init or return;
     my ($payload, $context) = $self->lambda_next;
     my $response = try {
+        local $AWS::Lambda::context = $context;
         $self->{function}->($payload, $context);
     } catch {
         $self->lambda_error($_, $context);

lib/AWS/Lambda/PSGI.pm

@@ -55,17 +55,21 @@ sub wrap {
 }
 
 sub call {
-    my($self, $env) = @_;
-    my $input = $self->format_input($env);
+    my($self, $env, $ctx) = @_;
+
+    # $ctx is added by #26
+    # fall back to $AWS::Lambda::context because of backward compatibility.
+    $ctx ||= $AWS::Lambda::context;
+
+    my $input = $self->format_input($env, $ctx);
     my $res = $self->app->($input);
     return $self->format_output($res);
 }
 
 # API Gateway https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html
 # Application Load Balancer https://docs.aws.amazon.com/elasticloadbalancing/latest/application/lambda-functions.html
 sub format_input {
-    my $self = shift;
-    my $payload = shift;
+    my ($self, $payload, $ctx) = @_;
     my $env = {};
 
     # merge queryStringParameters and multiValueQueryStringParameters
@@ -112,6 +116,12 @@ sub format_input {
     $env->{'psgix.harakiri'}    = Plack::Util::TRUE;
     $env->{'psgix.input.buffered'} = Plack::Util::TRUE;
 
+    # inject the request id that compatible with Plack::Middleware::RequestId
+    if ($ctx) {
+        $env->{'psgix.request_id'} = $ctx->aws_request_id;
+        $env->{'HTTP_X_REQUEST_ID'} = $ctx->aws_request_id;
+    }
+
     my $body = $payload->{body};
     if ($payload->{isBase64Encoded}) {
         $body = decode_base64 $body;
@@ -208,15 +218,22 @@ Add the following script into your Lambda code archive.
     my $func = AWS::Lambda::PSGI->wrap($app);
 
     sub handle {
-        my $payload = shift;
-        return $func->($payload);
+        return $func->(@_);
     }
 
     1;
 
 And then, L<Set up Lambda Proxy Integrations in API Gateway|https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html> or
 L<Lambda Functions as ALB Targets|https://docs.aws.amazon.com/elasticloadbalancing/latest/application/lambda-functions.html>
 
+=head1 DESCRIPTION
+
+=head2 Request ID
+
+L<AWS::Lambda::PSGI> injects the request id that compatible with L<Plack::Middleware::RequestId>.
+
+    env->{'psgix.request_id'} # It is same value with $context->aws_request_id
+
 =head1 LICENSE
 
 The MIT License (MIT)

t/20_psgi.t

@@ -11,6 +11,7 @@ use Test::Deep qw(cmp_deeply);
 use JSON::Types;
 use Encode;
 
+use AWS::Lambda::Context;
 use AWS::Lambda::PSGI;
 
 my $app = AWS::Lambda::PSGI->new;
@@ -211,4 +212,18 @@ subtest "query string enconding" => sub {
     is $req->query_string, 'gif+ref+=', 'query string';
 };
 
+# inject the request id that compatible with Plack::Middleware::RequestId
+subtest "request id" => sub {
+    my $context = AWS::Lambda::Context->new(
+        deadline_ms    => 1000,
+        aws_request_id => '8476a536-e9f4-11e8-9739-2dfe598c3fcd',
+        invoked_function_arn => 'arn:aws:lambda:us-east-2:123456789012:function:custom-runtime',
+    );
+    my $input = decode_json(slurp("$FindBin::Bin/testdata/apigateway-get-request.json"));
+    my $output = $app->format_input($input, $context);
+    my $req = Plack::Request->new($output);
+    is $req->env->{'psgix.request_id'}, '8476a536-e9f4-11e8-9739-2dfe598c3fcd', 'get request id from the env';
+    is $req->header('X-Request-Id'), '8476a536-e9f4-11e8-9739-2dfe598c3fcd', 'get request id from the header';
+};
+
 done_testing;

t/test_handlers/echo.pl

@@ -1,9 +1,13 @@
 use utf8;
 use strict;
 use warnings;
+use AWS::Lambda;
 
 sub handle {
-    my $payload = shift;
+    my ($payload, $context) = @_;
+    die "payload is empty" unless $payload;
+    die "context is empty" unless $context;
+    die "AWS::Lambda::context is invalid" if $AWS::Lambda::context != $context;
     return $payload;
 }