Add support for db cursors and connections in context managers (#1028)

Here is an example snippet that will not report tracing without this patch:

with psycopg2.connect(...) as conn, conn.cursor() as cursor:
    cursor.execute("select 1;")

Co-authored-by: Carl Bordum Hansen <carl@bordum.dk>
This commit is contained in:
Daniel
2020-09-02 20:26:16 +02:00
committed by GitHub
parent 1107b97209
commit 33954ca03b
2 changed files with 18 additions and 1 deletions

View File

@ -2,6 +2,9 @@
## Unreleased
- bugfix: cursors and connections now produce spans when used with context managers
([#1028](https://github.com/open-telemetry/opentelemetry-python/pull/1028))
## Version 0.12b0
Released 2020-08-14
@ -19,4 +22,4 @@ Released 2020-05-12
Released 2020-02-21
- Initial release
- Initial release

View File

@ -294,6 +294,13 @@ def get_traced_connection_proxy(
self.__wrapped__.cursor(*args, **kwargs), db_api_integration
)
def __enter__(self):
self.__wrapped__.__enter__()
return self
def __exit__(self, *args, **kwargs):
self.__wrapped__.__exit__(*args, **kwargs)
return TracedConnectionProxy(connection, *args, **kwargs)
@ -366,4 +373,11 @@ def get_traced_cursor_proxy(cursor, db_api_integration, *args, **kwargs):
self.__wrapped__.callproc, *args, **kwargs
)
def __enter__(self):
self.__wrapped__.__enter__()
return self
def __exit__(self, *args, **kwargs):
self.__wrapped__.__exit__(*args, **kwargs)
return TracedCursorProxy(cursor, *args, **kwargs)