Fixed implementation of TracingStatement when no segment is present.

aws-xray-recorder-sdk-sql/src/main/java/com/amazonaws/xray/sql/TracingStatement.java

@@ -90,42 +90,39 @@ private static class TracingStatementHandler implements InvocationHandler {
         }
 
         public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
-            if (!isExecution(method)) {
-                // don't trace non execution methods
-                return method.invoke(delegate, args);
-            }
+            Subsegment subsegment = null;
 
-            Subsegment subsegment = createSubsegment();
-            if (subsegment == null) {
-                // don't trace if failed to create subsegment
-                return method.invoke(delegate, args);
+            if (isExecution(method)) {
+                // only trace on execution methods
+                subsegment = createSubsegment();
             }
 
-            logger.debug("Invoking statement execution with X-Ray tracing.");
+            logger.debug(String.format("Invoking statement execution with X-Ray tracing. Tracing active: %s", subsegment != null));
             try {
                 // execute the query "wrapped" in a XRay Subsegment
                 return method.invoke(delegate, args);
             } catch (Throwable t) {
+                Throwable rootThrowable = t;
                 if (t instanceof InvocationTargetException) {
                     // the reflection may wrap the actual error with an InvocationTargetException.
                     // we want to use the root cause to make the instrumentation seamless
                     InvocationTargetException ite = (InvocationTargetException) t;
                     if (ite.getTargetException() != null) {
-                        subsegment.addException(ite.getTargetException());
-                        throw ite.getTargetException();
+                        rootThrowable = ite.getTargetException();
                     }
-                    if (ite.getCause() != null) {
-                        subsegment.addException(ite.getCause());
-                        throw ite.getCause();
+                    else if (ite.getCause() != null) {
+                        rootThrowable = ite.getCause();
                     }
-                    subsegment.addException(ite);
-                    throw ite;
                 }
 
-                subsegment.addException(t);
-                throw t;
+                if (subsegment != null) {
+                    subsegment.addException(rootThrowable);
+                }
+                throw rootThrowable;
             } finally {
-                AWSXRay.endSubsegment();
+                if (subsegment != null && isExecution(method)) {
+                    AWSXRay.endSubsegment();
+                }
             }
         }
 

aws-xray-recorder-sdk-sql/src/test/java/com/amazonaws/xray/sql/TracingStatementTest.java

@@ -10,6 +10,7 @@
 import java.sql.DatabaseMetaData;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
+import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.HashMap;
 import java.util.Map;
@@ -25,6 +26,8 @@
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.entities.Namespace;
 import com.amazonaws.xray.entities.Subsegment;
+import com.amazonaws.xray.strategy.ContextMissingStrategy;
+import com.amazonaws.xray.strategy.IgnoreErrorContextMissingStrategy;
 
 @RunWith(MockitoJUnitRunner.class)
 public class TracingStatementTest {
@@ -169,7 +172,7 @@ public void testCalledStatementExecuteUpdate() throws Exception {
     }
 
     @Test
-    public void testCaptureException() throws Exception {
+    public void testCaptureRuntimeException() throws Exception {
         Throwable exception = new RuntimeException("foo");
         when(delegate.execute(SQL)).thenThrow(exception);
         try {
@@ -183,6 +186,59 @@ public void testCaptureException() throws Exception {
         }
     }
 
+    @Test
+    public void testCaptureSqlException() throws Exception {
+        Throwable exception = new SQLException("foo");
+        when(delegate.execute(SQL)).thenThrow(exception);
+        try {
+            statement.execute(SQL);
+            fail("Expected exception is not thrown");
+        } catch (Throwable th) {
+            assertEquals(exception, th);
+        } finally {
+            assertEquals(exception, AWSXRay.getCurrentSegment().getSubsegments().get(0).getCause().getExceptions().get(0).getThrowable());
+            assertSubsegment();
+        }
+    }
+
+
+
+    @Test
+    public void testCaptureRuntimeExceptionWithoutSegment() throws Exception {
+        ContextMissingStrategy oldStrategy = AWSXRay.getGlobalRecorder().getContextMissingStrategy();
+        AWSXRay.getGlobalRecorder().setContextMissingStrategy(new IgnoreErrorContextMissingStrategy());
+        try {
+            Throwable exception = new RuntimeException("foo");
+            when(delegate.execute(SQL)).thenThrow(exception);
+            try {
+                statement.execute(SQL);
+                fail("Expected exception is not thrown");
+            } catch (Throwable th) {
+                assertEquals(exception, th);
+            }
+        } finally {
+            AWSXRay.getGlobalRecorder().setContextMissingStrategy(oldStrategy);
+        }
+    }
+
+    @Test
+    public void testCaptureSqlExceptionWithoutSegment() throws Exception {
+        ContextMissingStrategy oldStrategy = AWSXRay.getGlobalRecorder().getContextMissingStrategy();
+        AWSXRay.getGlobalRecorder().setContextMissingStrategy(new IgnoreErrorContextMissingStrategy());
+        try {
+            Throwable exception = new SQLException("foo");
+            when(delegate.execute(SQL)).thenThrow(exception);
+            try {
+                statement.execute(SQL);
+                fail("Expected exception is not thrown");
+            } catch (Throwable th) {
+                assertEquals(exception, th);
+            }
+        } finally {
+            AWSXRay.getGlobalRecorder().setContextMissingStrategy(oldStrategy);
+        }
+    }
+
     @Test
     public void testNonTracedMethod() throws Exception {
         statement.close();