Close segment even if error was raised

Author: FerumFlex

aws_xray_sdk/ext/sqlalchemy/util/decorators.py

@@ -58,11 +58,14 @@ def wrapper(*args, **kw):
                 subsegment = xray_recorder.begin_subsegment(sub_name, namespace='remote')
             else:
                 subsegment = None
-        res = func(*args, **kw)
-        if subsegment is not None:
-            subsegment.set_sql(sql)
-            subsegment.put_annotation("sqlalchemy", class_name+'.'+func.__name__)
-            xray_recorder.end_subsegment()
+
+        try:
+            res = func(*args, **kw)
+        finally:
+            if subsegment is not None:
+                subsegment.set_sql(sql)
+                subsegment.put_annotation("sqlalchemy", class_name+'.'+func.__name__)
+                xray_recorder.end_subsegment()
         return res
     return wrapper
 # URL Parse output