Support for custom keys in the logs

example/HelloWorldFunction/src/main/java/helloworld/App.java

@@ -14,6 +14,7 @@
 import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import software.aws.lambda.logging.PowerLogger;
 import software.aws.lambda.logging.PowerToolsLogging;
 
 /**
@@ -29,6 +30,8 @@ public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEv
         headers.put("Content-Type", "application/json");
         headers.put("X-Custom-Header", "application/json");
 
+        PowerLogger.customKey("test", "willBeLogged");
+
         APIGatewayProxyResponseEvent response = new APIGatewayProxyResponseEvent()
                 .withHeaders(headers);
         try {

src/main/java/org/apache/logging/log4j/core/layout/LambdaJsonLayout.java

@@ -1,5 +1,15 @@
 package org.apache.logging.log4j.core.layout;
 
+import java.io.IOException;
+import java.io.Writer;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
+import java.time.ZoneId;
+import java.time.ZonedDateTime;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
 import com.fasterxml.jackson.annotation.JsonAnyGetter;
 import com.fasterxml.jackson.annotation.JsonGetter;
 import com.fasterxml.jackson.annotation.JsonRootName;
@@ -18,22 +28,11 @@
 import org.apache.logging.log4j.core.util.KeyValuePair;
 import org.apache.logging.log4j.util.Strings;
 
-import java.io.IOException;
-import java.io.Writer;
-import java.nio.charset.Charset;
-import java.nio.charset.StandardCharsets;
-import java.time.ZoneId;
-import java.time.ZonedDateTime;
-import java.util.HashMap;
-import java.util.LinkedHashMap;
-import java.util.Map;
-
 import static java.time.Instant.ofEpochMilli;
 import static java.time.format.DateTimeFormatter.ISO_ZONED_DATE_TIME;
 
 @Plugin(name = "LambdaJsonLayout", category = Node.CATEGORY, elementType = Layout.ELEMENT_TYPE, printObject = true)
 public class LambdaJsonLayout extends AbstractJacksonLayout {
-
     private static final String DEFAULT_FOOTER = "]";
 
     private static final String DEFAULT_HEADER = "[";

src/main/java/software/aws/lambda/logging/PowerLogger.java

@@ -0,0 +1,10 @@
+package software.aws.lambda.logging;
+
+import org.apache.logging.log4j.ThreadContext;
+
+public class PowerLogger {
+
+    public static void customKey(String key, String value) {
+        ThreadContext.put(key, value);
+    }
+}

src/main/java/software/aws/lambda/logging/DefaultLambdaFields.java → src/main/java/software/aws/lambda/logging/internal/DefaultLambdaFields.java

@@ -1,4 +1,4 @@
-package software.aws.lambda.logging;
+package software.aws.lambda.logging.internal;
 
 import com.amazonaws.services.lambda.runtime.Context;
 

src/main/java/software/aws/lambda/logging/LambdaAspect.java → src/main/java/software/aws/lambda/logging/internal/LambdaAspect.java

@@ -1,4 +1,4 @@
-package software.aws.lambda.logging;
+package software.aws.lambda.logging.internal;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -21,6 +21,7 @@
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.annotation.Pointcut;
+import software.aws.lambda.logging.PowerToolsLogging;
 
 import static java.util.Optional.empty;
 import static java.util.Optional.of;

src/test/java/software/aws/lambda/logging/PowerLoggerTest.java

@@ -0,0 +1,26 @@
+package software.aws.lambda.logging;
+
+import org.apache.logging.log4j.ThreadContext;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import software.aws.lambda.logging.PowerLogger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+
+class PowerLoggerTest {
+
+    @BeforeEach
+    void setUp() {
+        ThreadContext.clearAll();
+    }
+
+    @Test
+    void shouldSetCustomKeyOnThreadContext() {
+        PowerLogger.customKey("test", "value");
+
+        assertThat(ThreadContext.getImmutableContext())
+                .hasSize(1)
+                .containsEntry("test", "value");
+    }
+}

src/test/java/software/aws/lambda/logging/LambdaAspectTest.java → src/test/java/software/aws/lambda/logging/internal/LambdaAspectTest.java

@@ -1,4 +1,4 @@
-package software.aws.lambda.logging;
+package software.aws.lambda.logging.internal;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;

src/test/java/software/aws/lambda/logging/PowerLogToolDisabled.java → src/test/java/software/aws/lambda/logging/internal/PowerLogToolDisabled.java

@@ -1,4 +1,4 @@
-package software.aws.lambda.logging;
+package software.aws.lambda.logging.internal;
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;

src/test/java/software/aws/lambda/logging/PowerLogToolDisabledForStream.java → src/test/java/software/aws/lambda/logging/internal/PowerLogToolDisabledForStream.java

@@ -1,4 +1,4 @@
-package software.aws.lambda.logging;
+package software.aws.lambda.logging.internal;
 
 import java.io.InputStream;
 import java.io.OutputStream;

src/test/java/software/aws/lambda/logging/PowerLogToolEnabled.java → src/test/java/software/aws/lambda/logging/internal/PowerLogToolEnabled.java

@@ -1,7 +1,8 @@
-package software.aws.lambda.logging;
+package software.aws.lambda.logging.internal;
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
+import software.aws.lambda.logging.PowerToolsLogging;
 
 public class PowerLogToolEnabled implements RequestHandler<Object, Object> {
 

src/test/java/software/aws/lambda/logging/PowerLogToolEnabledForStream.java → src/test/java/software/aws/lambda/logging/internal/PowerLogToolEnabledForStream.java

@@ -1,10 +1,11 @@
-package software.aws.lambda.logging;
+package software.aws.lambda.logging.internal;
 
 import java.io.InputStream;
 import java.io.OutputStream;
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import software.aws.lambda.logging.PowerToolsLogging;
 
 public class PowerLogToolEnabledForStream implements RequestStreamHandler {
 

src/test/java/software/aws/lambda/logging/PowerToolLogEventEnabled.java → src/test/java/software/aws/lambda/logging/internal/PowerToolLogEventEnabled.java

@@ -1,7 +1,8 @@
-package software.aws.lambda.logging;
+package software.aws.lambda.logging.internal;
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
+import software.aws.lambda.logging.PowerToolsLogging;
 
 public class PowerToolLogEventEnabled implements RequestHandler<Object, Object> {
 

src/test/java/software/aws/lambda/logging/PowerToolLogEventEnabledForStream.java → src/test/java/software/aws/lambda/logging/internal/PowerToolLogEventEnabledForStream.java

@@ -1,4 +1,4 @@
-package software.aws.lambda.logging;
+package software.aws.lambda.logging.internal;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -8,6 +8,7 @@
 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 PowerToolLogEventEnabledForStream implements RequestStreamHandler {