fix(dbapi2): preserve XRayTracedConn wrapper in connection context manager

aws_xray_sdk/ext/dbapi2.py

@@ -13,12 +13,20 @@ def __init__(self, conn, meta={}):
         super().__init__(conn)
         self._xray_meta = meta
 
+    def __enter__(self):
+
+        value = self.__wrapped__.__enter__()
+        if value is not self.__wrapped__:
+            return value
+        return self
+
     def cursor(self, *args, **kwargs):
 
         cursor = self.__wrapped__.cursor(*args, **kwargs)
         return XRayTracedCursor(cursor, self._xray_meta)
 
 
+
 class XRayTracedCursor(wrapt.ObjectProxy):
 
     _xray_meta = None

tests/ext/sqlite3/test_sqlite3.py

@@ -54,3 +54,16 @@ def test_invalid_syntax(db):
 
     exception = subsegment.cause['exceptions'][0]
     assert exception.type == 'OperationalError'
+
+
+def test_connection_context_manager(db):
+    with db as conn:
+        q = 'SELECT name FROM sqlite_master'
+        cur = conn.cursor()
+        cur.execute(q)
+
+    subsegment = xray_recorder.current_segment().subsegments[0]
+    assert subsegment.name == ':memory:'
+    sql = subsegment.sql
+    assert sql['database_type'] == 'sqlite3'
+    assert sql['database_version']