Support for logging event (#11)

pankajagrawal16 · web-flow · commit 1478899ed719 · 2020-07-31T14:48:42.000+02:00
diff --git a/example/HelloWorldFunction/src/main/java/helloworld/App.java b/example/HelloWorldFunction/src/main/java/helloworld/App.java
@@ -23,7 +23,7 @@ public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatew
 
     Logger log = LogManager.getLogger();
 
-    @PowerToolsLogging
+    @PowerToolsLogging(logEvent = true)
     public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
         Map<String, String> headers = new HashMap<>();
         headers.put("Content-Type", "application/json");
diff --git a/example/HelloWorldFunction/src/main/java/helloworld/AppStream.java b/example/HelloWorldFunction/src/main/java/helloworld/AppStream.java
@@ -0,0 +1,23 @@
+package helloworld;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.Map;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import software.aws.lambda.logging.PowerToolsLogging;
+
+public class AppStream implements RequestStreamHandler {
+    private static final ObjectMapper mapper = new ObjectMapper();
+
+    @Override
+    @PowerToolsLogging(logEvent = true)
+    public void handleRequest(InputStream input, OutputStream output, Context context) throws IOException {
+        Map map = mapper.readValue(input, Map.class);
+
+        System.out.println(map.size());
+    }
+}
diff --git a/example/template.yaml b/example/template.yaml
@@ -28,6 +28,23 @@ Resources:
             Path: /hello
             Method: get
 
+  HelloWorldStreamFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      CodeUri: HelloWorldFunction
+      Handler: helloworld.AppStream::handleRequest
+      Runtime: java8
+      MemorySize: 512
+      Environment:
+        Variables:
+          PARAM1: VALUE
+      Events:
+        HelloWorld:
+          Type: Api
+          Properties:
+            Path: /hellostream
+            Method: get
+
 Outputs:
   # ServerlessRestApi is an implicit API created out of Events key under Serverless::Function
   # Find out more about other implicit resources you can reference within SAM
diff --git a/src/main/java/software/aws/lambda/logging/LambdaAspect.java b/src/main/java/software/aws/lambda/logging/LambdaAspect.java
@@ -1,11 +1,22 @@
 package software.aws.lambda.logging;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.OutputStreamWriter;
+import java.util.Map;
 import java.util.Optional;
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
 import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.ThreadContext;
+import org.apache.logging.log4j.core.util.IOUtils;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
@@ -17,14 +28,16 @@
 @Aspect
 public final class LambdaAspect {
     static Boolean IS_COLD_START = null;
+    private static final ObjectMapper mapper = new ObjectMapper();
 
     @Pointcut("@annotation(powerToolsLogging)")
     public void callAt(PowerToolsLogging powerToolsLogging) {
     }
 
-    @Around(value = "callAt(powerToolsLogging)")
+    @Around(value = "callAt(powerToolsLogging)", argNames = "pjp,powerToolsLogging")
     public Object around(ProceedingJoinPoint pjp,
                          PowerToolsLogging powerToolsLogging) throws Throwable {
+        Object[] proceedArgs = pjp.getArgs();
 
         extractContext(pjp)
                 .ifPresent(context -> {
@@ -34,24 +47,84 @@ public Object around(ProceedingJoinPoint pjp,
 
         IS_COLD_START = false;
 
+        if (powerToolsLogging.logEvent()) {
+            proceedArgs = logEvent(pjp);
+        }
 
-        return pjp.proceed();
+        return pjp.proceed(proceedArgs);
     }
 
     private Optional<Context> extractContext(ProceedingJoinPoint pjp) {
 
-        if ("handleRequest".equals(pjp.getSignature().getName())) {
-            if (RequestHandler.class.isAssignableFrom(pjp.getSignature().getDeclaringType())
-                    && pjp.getArgs().length == 2 && pjp.getArgs()[1] instanceof Context) {
+        if (isHandlerMethod(pjp)) {
+            if (placedOnRequestHandler(pjp)) {
                 return of((Context) pjp.getArgs()[1]);
             }
 
-            if (RequestStreamHandler.class.isAssignableFrom(pjp.getSignature().getDeclaringType())
-                    && pjp.getArgs().length == 3 && pjp.getArgs()[2] instanceof Context) {
+            if (placedOnStreamHandler(pjp)) {
                 return of((Context) pjp.getArgs()[2]);
             }
         }
 
         return empty();
     }
+
+    private Object[] logEvent(ProceedingJoinPoint pjp) {
+        Object[] args = pjp.getArgs();
+
+        if (isHandlerMethod(pjp)) {
+            if (placedOnRequestHandler(pjp)) {
+                Logger log = logger(pjp);
+                log.info(pjp.getArgs()[0]);
+            }
+
+            if (placedOnStreamHandler(pjp)) {
+                args = logFromInputStream(pjp);
+            }
+        }
+
+        return args;
+    }
+
+    private Object[] logFromInputStream(ProceedingJoinPoint pjp) {
+        Object[] args = pjp.getArgs();
+
+        try (ByteArrayOutputStream out = new ByteArrayOutputStream();
+             OutputStreamWriter writer = new OutputStreamWriter(out);
+             InputStreamReader reader = new InputStreamReader((InputStream) pjp.getArgs()[0])) {
+
+            IOUtils.copy(reader, writer);
+            writer.flush();
+            byte[] bytes = out.toByteArray();
+            args[0] = new ByteArrayInputStream(bytes);
+
+            Logger log = logger(pjp);
+            log.info(mapper.readValue(bytes, Map.class));
+
+        } catch (IOException e) {
+            Logger log = logger(pjp);
+            log.debug("Failed to log event from supplied input stream.", e);
+        }
+
+        return args;
+    }
+
+    private Logger logger(ProceedingJoinPoint pjp) {
+        return LogManager.getLogger(pjp.getSignature().getDeclaringType());
+    }
+
+    private boolean isHandlerMethod(ProceedingJoinPoint pjp) {
+        return "handleRequest".equals(pjp.getSignature().getName());
+    }
+
+    private boolean placedOnRequestHandler(ProceedingJoinPoint pjp) {
+        return RequestHandler.class.isAssignableFrom(pjp.getSignature().getDeclaringType())
+                && pjp.getArgs().length == 2 && pjp.getArgs()[1] instanceof Context;
+    }
+
+    private boolean placedOnStreamHandler(ProceedingJoinPoint pjp) {
+        return RequestStreamHandler.class.isAssignableFrom(pjp.getSignature().getDeclaringType())
+                && pjp.getArgs().length == 3 && pjp.getArgs()[0] instanceof InputStream
+                && pjp.getArgs()[2] instanceof Context;
+    }
 }
diff --git a/src/main/java/software/aws/lambda/logging/PowerToolsLogging.java b/src/main/java/software/aws/lambda/logging/PowerToolsLogging.java
@@ -8,4 +8,6 @@
 @Retention(RetentionPolicy.RUNTIME)
 @Target(ElementType.METHOD)
 public @interface PowerToolsLogging {
+
+    boolean logEvent() default false;
 }
diff --git a/src/test/java/software/aws/lambda/logging/LambdaAspectTest.java b/src/test/java/software/aws/lambda/logging/LambdaAspectTest.java
@@ -1,10 +1,16 @@
 package software.aws.lambda.logging;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
 import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.logging.log4j.ThreadContext;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -49,7 +55,7 @@ void shouldSetLambdaContextWhenEnabled() {
     void shouldSetLambdaContextForStreamHandlerWhenEnabled() throws IOException {
         requestStreamHandler = new PowerLogToolEnabledForStream();
 
-        requestStreamHandler.handleRequest(null, null, context);
+        requestStreamHandler.handleRequest(new ByteArrayInputStream(new byte[]{}), null, context);
 
         assertThat(ThreadContext.getImmutableContext())
                 .hasSize(5)
@@ -62,13 +68,13 @@ void shouldSetLambdaContextForStreamHandlerWhenEnabled() throws IOException {
 
     @Test
     void shouldSetColdStartFlag() throws IOException {
-        requestStreamHandler.handleRequest(null, null, context);
+        requestStreamHandler.handleRequest(new ByteArrayInputStream(new byte[]{}), null, context);
 
         assertThat(ThreadContext.getImmutableContext())
                 .hasSize(5)
                 .containsEntry("coldStart", "true");
 
-        requestStreamHandler.handleRequest(null, null, context);
+        requestStreamHandler.handleRequest(new ByteArrayInputStream(new byte[]{}), null, context);
 
         assertThat(ThreadContext.getImmutableContext())
                 .hasSize(5)
@@ -105,6 +111,33 @@ void shouldHaveNoEffectIfNotUsedOnLambdaHandler() {
                 .isEmpty();
     }
 
+    @Test
+    void shouldLogEventForHandler() {
+        requestHandler = new PowerToolLogEventEnabled();
+
+        requestHandler.handleRequest(new Object(), context);
+
+        assertThat(ThreadContext.getImmutableContext())
+                .hasSize(5);
+    }
+
+    @Test
+    void shouldLogEventForStreamAndLambdaStreamIsValid() throws IOException {
+        requestStreamHandler = new PowerToolLogEventEnabledForStream();
+        ByteArrayOutputStream output = new ByteArrayOutputStream();
+
+        Map<String, String> testPayload = new HashMap<>();
+        testPayload.put("test", "payload");
+
+        requestStreamHandler.handleRequest(new ByteArrayInputStream(new ObjectMapper().writeValueAsBytes(testPayload)), output, context);
+
+        assertThat(new String(output.toByteArray(), StandardCharsets.UTF_8))
+                .isEqualTo("{\"test\":\"payload\"}");
+
+        assertThat(ThreadContext.getImmutableContext())
+                .hasSize(5);
+    }
+
     private void setupContext() {
         when(context.getFunctionName()).thenReturn("testFunction");
         when(context.getInvokedFunctionArn()).thenReturn("testArn");
diff --git a/src/test/java/software/aws/lambda/logging/PowerToolLogEventEnabled.java b/src/test/java/software/aws/lambda/logging/PowerToolLogEventEnabled.java
@@ -0,0 +1,13 @@
+package software.aws.lambda.logging;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+
+public class PowerToolLogEventEnabled implements RequestHandler<Object, Object> {
+
+    @PowerToolsLogging(logEvent = true)
+    @Override
+    public Object handleRequest(Object input, Context context) {
+        return null;
+    }
+}
diff --git a/src/test/java/software/aws/lambda/logging/PowerToolLogEventEnabledForStream.java b/src/test/java/software/aws/lambda/logging/PowerToolLogEventEnabledForStream.java
@@ -0,0 +1,20 @@
+package software.aws.lambda.logging;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.Map;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+public class PowerToolLogEventEnabledForStream implements RequestStreamHandler {
+
+    @PowerToolsLogging(logEvent = true)
+    @Override
+    public void handleRequest(InputStream input, OutputStream output, Context context) throws IOException {
+        ObjectMapper mapper = new ObjectMapper();
+        mapper.writeValue(output, mapper.readValue(input, Map.class));
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -8,4 +8,6 @@`
`8`	`8`	`@Retention(RetentionPolicy.RUNTIME)`
`9`	`9`	`@Target(ElementType.METHOD)`
`10`	`10`	`public @interface PowerToolsLogging {`
	`11`	`+`
	`12`	`+ boolean logEvent() default false;`
`11`	`13`	`}`