Skip to content

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

Open
reginaldalfret wants to merge 1 commit into
aws:masterfrom
reginaldalfret:fix-490-connection-context-manager
Open

reginaldalfret wants to merge 1 commit into
aws:masterfrom
reginaldalfret:fix-490-connection-context-manager

Conversation

@reginaldalfret

Copy link
Copy Markdown

Fixes #490

Problem

When using a database connection wrapped by \XRayTracedConn\ within a context manager (\with conn as c:), \wrapt.ObjectProxy\ forwards _enter_\ to the underlying database connection. Since DBAPI and driver \Connection.enter()\ returns \self\ (the underlying connection object), \c\ ends up referencing the unwrapped connection rather than \XRayTracedConn. Consequently, subsequent \c.cursor()\ calls return raw cursors, and SQL queries executed inside the context manager are not traced by X-Ray.

Root Cause

\XRayTracedCursor\ already implements _enter_\ to return \self\ when _wrapped_.enter()\ returns the underlying object, but \XRayTracedConn\ was missing the corresponding _enter_\ implementation.

Solution

Implement _enter_\ on \XRayTracedConn\ matching the existing pattern in \XRayTracedCursor:
\\python
def enter(self):
value = self.wrapped.enter()
if value is not self.wrapped:
return value
return self
\\

Validation

  • Added regression test \ est_connection_context_manager\ in \ ests/ext/sqlite3/test_sqlite3.py\ demonstrating that \with db as conn:\ preserves the traced connection wrapper and correctly records SQL subsegments.
  • Confirmed test fails without the fix (\IndexError: list index out of range\ due to missing subsegments) and passes cleanly with the fix.
  • Ran all existing SQLite3 and recorder tests (39 passed).

@reginaldalfret
reginaldalfret requested a review from a team as a code owner September 25, 2026 04:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

psycopg connection context manager drops X-Ray tracing wrapper, so cursor execute calls are not traced

1 participant