feat: custom segment names (#221)

* add segmentName attribute to Trace annotation
* if segmentName is provided, use it as subsegment name intead of traced
method's name

Author: volomas

docs/content/core/tracing.mdx

@@ -62,6 +62,8 @@ public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatew
 ```
 
 By default this annotation will automatically record method responses and exceptions.
+If you want to customize segment name that appears in traces, use:
+`@Tracing(segmentName="yourCustomName")`
 
 <Note type="warning">
   <strong>Returning sensitive information from your Lambda handler or functions, where Tracer is used?</strong>

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

@@ -50,4 +50,5 @@
     String namespace() default "";
     boolean captureResponse() default true;
     boolean captureError() default true;
+    String segmentName() default "";
 }

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

@@ -15,6 +15,7 @@
 
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.entities.Subsegment;
+import java.util.function.Supplier;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
@@ -42,7 +43,9 @@ public Object around(ProceedingJoinPoint pjp,
                          Tracing tracing) throws Throwable {
         Object[] proceedArgs = pjp.getArgs();
 
-        Subsegment segment = AWSXRay.beginSubsegment("## " + pjp.getSignature().getName());
+        Subsegment segment = AWSXRay.beginSubsegment(
+            customSegmentNameOrDefault(tracing,
+                () -> "## " + pjp.getSignature().getName()));
         segment.setNamespace(namespace(tracing));
 
         if (placedOnHandlerMethod(pjp)) {
@@ -69,6 +72,11 @@ public Object around(ProceedingJoinPoint pjp,
         }
     }
 
+    private String customSegmentNameOrDefault(Tracing powerToolsTracing, Supplier<String> defaultSegmentName) {
+        String segmentName = powerToolsTracing.segmentName();
+        return segmentName.isEmpty() ? defaultSegmentName.get() : segmentName;
+    }
+
     private String namespace(Tracing powerToolsTracing) {
         return powerToolsTracing.namespace().isEmpty() ? serviceName() : powerToolsTracing.namespace();
     }

powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/internal/LambdaTracingAspectTest.java

@@ -26,6 +26,7 @@
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
 import software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor;
+import software.amazon.lambda.powertools.tracing.nonhandler.PowerToolNonHandler;
 import software.amazon.lambda.powertools.tracing.handlers.PowerToolDisabled;
 import software.amazon.lambda.powertools.tracing.handlers.PowerToolDisabledForStream;
 import software.amazon.lambda.powertools.tracing.handlers.PowerTracerToolEnabled;
@@ -43,6 +44,7 @@
 class LambdaTracingAspectTest {
     private RequestHandler<Object, Object> requestHandler;
     private RequestStreamHandler streamHandler;
+    private PowerToolNonHandler nonHandlerMethod;
 
     @Mock
     private Context context;
@@ -54,6 +56,7 @@ void setUp() throws IllegalAccessException {
         setupContext();
         requestHandler = new PowerTracerToolEnabled();
         streamHandler = new PowerTracerToolEnabledForStream();
+        nonHandlerMethod = new PowerToolNonHandler();
         AWSXRay.beginSegment(LambdaTracingAspectTest.class.getName());
     }
 
@@ -63,6 +66,24 @@ void tearDown() {
     }
 
     @Test
+    void shouldCaptureNonHandlerMethod() {
+        nonHandlerMethod.doSomething();
+        assertThat(AWSXRay.getTraceEntity().getSubsegments())
+            .hasSize(1)
+            .anySatisfy(segment ->
+                assertThat(segment.getName()).isEqualTo("## doSomething"));
+    }
+
+    @Test
+    void shouldCaptureNonHandlerMethodWithCustomSegmentName() {
+        nonHandlerMethod.doSomethingCustomName();
+        assertThat(AWSXRay.getTraceEntity().getSubsegments())
+            .hasSize(1)
+            .anySatisfy(segment ->
+                assertThat(segment.getName()).isEqualTo("custom"));
+    }
+
+  @Test
     void shouldCaptureTraces() {
         requestHandler.handleRequest(new Object(), context);
 
@@ -177,4 +198,4 @@ private void setupContext() {
         when(context.getFunctionVersion()).thenReturn("1");
         when(context.getMemoryLimitInMB()).thenReturn(10);
     }
-}
+}

powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/nonhandler/PowerToolNonHandler.java

@@ -0,0 +1,14 @@
+package software.amazon.lambda.powertools.tracing.nonhandler;
+
+import software.amazon.lambda.powertools.tracing.Tracing;
+
+public class PowerToolNonHandler {
+
+    @Tracing
+    public void doSomething() {
+    }
+
+    @Tracing(segmentName = "custom")
+    public void doSomethingCustomName() {
+    }
+}