Fixed implementation of TracingStatement when no segment is present.

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

@@ -97,8 +97,24 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
 
             Subsegment subsegment = createSubsegment();
             if (subsegment == null) {
-                // don't trace if failed to create subsegment
-                return method.invoke(delegate, args);
+                try {
+                    // don't trace if failed to create subsegment
+                    return method.invoke(delegate, args);
+                } catch (Throwable 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) {
+                            throw ite.getTargetException();
+                        }
+                        if (ite.getCause() != null) {
+                            throw ite.getCause();
+                        }
+                    }
+
+                    throw t;
+                }
             }
 
             logger.debug("Invoking statement execution with X-Ray tracing.");

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();