feat(tracing): ability to override object mapper (#698)

* feat(tracing): ability to override object mapper used for serialization

* chore(tracing): fix static code warnings

* docs(tracing): ability to override object mapper used for serialization

Author: pankajagrawal16

docs/core/tracing.md

@@ -164,6 +164,38 @@ context for an operation using any native object.
     }
     ```
 
+## Override default object mapper
+
+You can optionally choose to override default object mapper which is used to serialize method response and exceptions when enabled. You might
+want to supply custom object mapper in order to control how serialisation is done, for example, when you want to log only
+specific fields from received event due to security.
+
+=== "App.java"
+
+    ```java hl_lines="10-14"
+    import software.amazon.lambda.powertools.tracing.Tracing;
+    import software.amazon.lambda.powertools.tracing.TracingUtils;
+    import static software.amazon.lambda.powertools.tracing.CaptureMode.RESPONSE;
+
+    /**
+     * Handler for requests to Lambda function.
+     */
+    public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> { 
+        static {
+            ObjectMapper objectMapper = new ObjectMapper();
+            SimpleModule simpleModule = new SimpleModule();
+            objectMapper.registerModule(simpleModule);
+
+            TracingUtils.defaultObjectMapper(objectMapper);
+        }
+    
+        @Tracing(captureMode = RESPONSE)
+        public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
+            ...
+        }
+    }
+    ```
+
 ## Utilities
 
 Tracing modules comes with certain utility method when you don't want to use annotation for capturing a code block

powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/TracingUtils.java

@@ -13,11 +13,11 @@
  */
 package software.amazon.lambda.powertools.tracing;
 
+import java.util.function.Consumer;
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.entities.Entity;
 import com.amazonaws.xray.entities.Subsegment;
-
-import java.util.function.Consumer;
+import com.fasterxml.jackson.databind.ObjectMapper;
 
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.serviceName;
 
@@ -27,6 +27,7 @@
  *
  */
 public final class TracingUtils {
+    private static ObjectMapper objectMapper;
 
     /**
      * Put an annotation to the current subsegment with a String value.
@@ -155,4 +156,18 @@ public static void withSubsegment(String namespace, String name, Consumer<Subseg
             AWSXRay.endSubsegment();
         }
     }
+
+    /**
+     * Sets the instance of ObjectMapper object which is used for serialising object response when
+     * {@code @Tracing(captureMode=CaptureMode.RESPONSE)}.
+     *
+     * @param objectMapper Custom implementation of object mapper to be used for serializing response
+     */
+    public static void defaultObjectMapper(ObjectMapper objectMapper) {
+        TracingUtils.objectMapper = objectMapper;
+    }
+
+    public static ObjectMapper objectMapper() {
+        return objectMapper;
+    }
 }

powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/internal/LambdaTracingAspect.java

@@ -14,7 +14,6 @@
 package software.amazon.lambda.powertools.tracing.internal;
 
 import java.util.function.Supplier;
-
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.entities.Subsegment;
 import org.aspectj.lang.ProceedingJoinPoint;
@@ -30,6 +29,7 @@
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnRequestHandler;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnStreamHandler;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.serviceName;
+import static software.amazon.lambda.powertools.tracing.TracingUtils.objectMapper;
 
 @Aspect
 public final class LambdaTracingAspect {
@@ -59,7 +59,7 @@ public Object around(ProceedingJoinPoint pjp,
         try {
             Object methodReturn = pjp.proceed(proceedArgs);
             if (captureResponse) {
-                segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " response", methodReturn);
+                segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " response", null != objectMapper() ? objectMapper().writeValueAsString(methodReturn): methodReturn);
             }
 
             if (placedOnHandlerMethod(pjp)) {
@@ -69,7 +69,7 @@ public Object around(ProceedingJoinPoint pjp,
             return methodReturn;
         } catch (Exception e) {
             if (captureError) {
-                segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " error", e);
+                segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " error", null != objectMapper() ? objectMapper().writeValueAsString(e) : e);
             }
             throw e;
         } finally {

powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/handlers/PowerTracerToolEnabledForResponseWithCustomMapper.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates.
+ * Licensed under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package software.amazon.lambda.powertools.tracing.handlers;
+
+import java.io.IOException;
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.module.SimpleModule;
+import com.fasterxml.jackson.databind.ser.std.StdSerializer;
+import software.amazon.lambda.powertools.tracing.Tracing;
+import software.amazon.lambda.powertools.tracing.TracingUtils;
+
+import static software.amazon.lambda.powertools.tracing.CaptureMode.RESPONSE;
+
+public class PowerTracerToolEnabledForResponseWithCustomMapper implements RequestHandler<Object, Object> {
+    static {
+        ObjectMapper objectMapper = new ObjectMapper();
+        SimpleModule simpleModule = new SimpleModule();
+        simpleModule.addSerializer(ChildClass.class, new ChildSerializer());
+        objectMapper.registerModule(simpleModule);
+
+        TracingUtils.defaultObjectMapper(objectMapper);
+    }
+    @Override
+    @Tracing(namespace = "lambdaHandler", captureMode = RESPONSE)
+    public Object handleRequest(Object input, Context context) {
+        ParentClass parentClass = new ParentClass("parent");
+        ChildClass childClass = new ChildClass("child", parentClass);
+        parentClass.setC(childClass);
+        return parentClass;
+    }
+
+    public class ParentClass {
+        public String name;
+        public ChildClass c;
+
+        public ParentClass(String name) {
+            this.name = name;
+        }
+
+        public void setC(ChildClass c) {
+            this.c = c;
+        }
+    }
+
+    public class ChildClass {
+        public String name;
+        public ParentClass p;
+
+        public ChildClass(String name, ParentClass p) {
+            this.name = name;
+            this.p = p;
+        }
+    }
+
+    public static class ChildSerializer extends StdSerializer<ChildClass> {
+
+        public ChildSerializer() {
+            this(null);
+        }
+
+        public ChildSerializer(Class<ChildClass> t) {
+            super(t);
+        }
+
+        @Override
+        public void serialize(ChildClass value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
+            jgen.writeStartObject();
+            jgen.writeStringField("name", value.name);
+            jgen.writeStringField("p", value.p.name);
+            jgen.writeEndObject();
+        }
+    }
+}