diff --git a/aws_xray_sdk/ext/dbapi2.py b/aws_xray_sdk/ext/dbapi2.py index 4fda3164..7a015750 100644 --- a/aws_xray_sdk/ext/dbapi2.py +++ b/aws_xray_sdk/ext/dbapi2.py @@ -33,6 +33,13 @@ def __init__(self, cursor, meta={}): db_type = cursor.__class__.__module__.split('.')[0] self._xray_meta['database_type'] = db_type + def __enter__(self): + + value = self.__wrapped__.__enter__() + if value is not self.__wrapped__: + return value + return self + @xray_recorder.capture() def execute(self, query, *args, **kwargs):