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 index 51bcddce1..ec1416ad6 100644 --- 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 @@ -29,7 +29,7 @@ public final class TracingUtils { /** - * Put an annotation to the current subsegment. + * Put an annotation to the current subsegment with a String value. * * @param key the key of the annotation * @param value the value of the annotation @@ -39,6 +39,28 @@ public static void putAnnotation(String key, String value) { .ifPresent(segment -> segment.putAnnotation(key, value)); } + /** + * Put an annotation to the current subsegment with a Number value. + * + * @param key the key of the annotation + * @param value the value of the annotation + */ + public static void putAnnotation(String key, Number value) { + AWSXRay.getCurrentSubsegmentOptional() + .ifPresent(segment -> segment.putAnnotation(key, value)); + } + + /** + * Put an annotation to the current subsegment with a Boolean value. + * + * @param key the key of the annotation + * @param value the value of the annotation + */ + public static void putAnnotation(String key, Boolean value) { + AWSXRay.getCurrentSubsegmentOptional() + .ifPresent(segment -> segment.putAnnotation(key, value)); + } + /** * Put additional metadata for the current subsegment. * diff --git a/powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java b/powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java index 2c65c90f1..d2a96ec65 100644 --- a/powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java +++ b/powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java @@ -21,6 +21,7 @@ import org.junit.jupiter.api.Test; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.entry; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static software.amazon.lambda.powertools.tracing.TracingUtils.withEntitySubsegment; @@ -45,11 +46,17 @@ void tearDown() { void shouldSetAnnotationOnCurrentSubSegment() { AWSXRay.beginSubsegment("subSegment"); - TracingUtils.putAnnotation("key", "val"); + TracingUtils.putAnnotation("stringKey", "val"); + TracingUtils.putAnnotation("numberKey", 10); + TracingUtils.putAnnotation("booleanKey", false); assertThat(AWSXRay.getTraceEntity().getAnnotations()) - .hasSize(1) - .containsEntry("key", "val"); + .hasSize(3) + .contains( + entry("stringKey", "val"), + entry("numberKey", 10), + entry("booleanKey", false) + ); } @Test