From 279decfbffe73fda561eb0c0753dabea0a8b93c4 Mon Sep 17 00:00:00 2001
From: Reginald Alfret <reginaldalfret@gmail.com>
Date: Fri, 25 Sep 2026 10:02:38 +0530
Subject: [PATCH] fix(dbapi2): preserve XRayTracedConn wrapper in connection
 context manager (fixes #490)

---
 aws_xray_sdk/ext/dbapi2.py        |  8 ++++++++
 tests/ext/sqlite3/test_sqlite3.py | 13 +++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/aws_xray_sdk/ext/dbapi2.py b/aws_xray_sdk/ext/dbapi2.py
index 3db4f44f..3e5b566c 100644
--- a/aws_xray_sdk/ext/dbapi2.py
+++ b/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
diff --git a/tests/ext/sqlite3/test_sqlite3.py b/tests/ext/sqlite3/test_sqlite3.py
index 3d0e555b..e68a9da4 100644
--- a/tests/ext/sqlite3/test_sqlite3.py
+++ b/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']