chore: Addition of Warn Message If Invalid Annotation Key While Tracing

#1511 (#1512)
aws-powertools · Nov 14, 2023 · 30c8d33 · 30c8d33
1 parent de547d0
commit 30c8d33
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 2 deletions.
diff --git a/powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/TracingUtils.java b/powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/TracingUtils.java
@@ -21,12 +21,15 @@
 import com.amazonaws.xray.entities.Subsegment;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import java.util.function.Consumer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * A class of helper functions to add additional functionality and ease
  * of use.
  */
 public final class TracingUtils {
+    private static final Logger LOG = LoggerFactory.getLogger(TracingUtils.class);
     private static ObjectMapper objectMapper;
 
     /**
@@ -36,6 +39,10 @@ public final class TracingUtils {
      * @param value the value of the annotation
      */
     public static void putAnnotation(String key, String value) {
+        if (!isValidAnnotationKey(key)) {
+            LOG.warn("Ignoring annotation with unsupported characters in key: {}", key);
+            return;
+        }
         AWSXRay.getCurrentSubsegmentOptional()
                 .ifPresent(segment -> segment.putAnnotation(key, value));
     }
@@ -47,10 +54,24 @@ public static void putAnnotation(String key, String value) {
      * @param value the value of the annotation
      */
     public static void putAnnotation(String key, Number value) {
+        if (!isValidAnnotationKey(key)) {
+            LOG.warn("Ignoring annotation with unsupported characters in key: {}", key);
+            return;
+        }
         AWSXRay.getCurrentSubsegmentOptional()
                 .ifPresent(segment -> segment.putAnnotation(key, value));
     }
 
+    /**
+     Make sure that the annotation key is valid according to
+     <a href='https://docs.aws.amazon.com/xray/latest/devguide/xray-api-segmentdocuments.html#api-segmentdocuments-annotations'>the documentation</a>.
+
+     Annotation keys that are added that are invalid are ignored by x-ray.
+     **/
+    private static boolean isValidAnnotationKey(String key) {
+        return key.matches("^[a-zA-Z0-9_]+$");
+    }
+
     /**
      * Put an annotation to the current subsegment with a Boolean value.
      *

diff --git a/...ols-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java b/...ols-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java
@@ -28,7 +28,6 @@
 import org.junit.jupiter.api.Test;
 
 class TracingUtilsTest {
-
     @BeforeEach
     void setUp() {
         AWSXRay.beginSegment("test");
@@ -123,6 +122,24 @@ void shouldInvokeCodeBlockWrappedWithinSubsegment() {
                 });
     }
 
+    @Test
+    void shouldNotAddAnnotationIfInvalidCharacterInKey() {
+        AWSXRay.beginSubsegment("subSegment");
+        String inputKey = "stringKey with spaces";
+        TracingUtils.putAnnotation(inputKey, "val");
+        AWSXRay.getCurrentSubsegmentOptional()
+                .ifPresent(segment -> assertThat(segment.getAnnotations()).size().isEqualTo(0));
+    }
+
+    @Test
+    void shouldAddAnnotationIfValidCharactersInKey() {
+        AWSXRay.beginSubsegment("subSegment");
+        String inputKey = "validKey";
+        TracingUtils.putAnnotation(inputKey, "val");
+        AWSXRay.getCurrentSubsegmentOptional()
+                .ifPresent(segment -> assertThat(segment.getAnnotations()).size().isEqualTo(1));
+    }
+
     @Test
     void shouldInvokeCodeBlockWrappedWithinNamespacedSubsegment() {
         Context test = mock(Context.class);
@@ -221,4 +238,4 @@ void shouldInvokeCodeBlockWrappedWithinNamespacedEntitySubsegment() throws Inter
                             .containsEntry("key", "val");
                 });
     }
-}
+}