From 33954ca03b007e99e7e07cfc805d0fb8c16f5780 Mon Sep 17 00:00:00 2001 From: Daniel <61800298+ffe4@users.noreply.github.com> Date: Wed, 2 Sep 2020 20:26:16 +0200 Subject: [PATCH] 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 --- .../CHANGELOG.md | 5 ++++- .../instrumentation/dbapi/__init__.py | 14 ++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/instrumentation/opentelemetry-instrumentation-dbapi/CHANGELOG.md b/instrumentation/opentelemetry-instrumentation-dbapi/CHANGELOG.md index bdb9236ac..99e1e09b5 100644 --- a/instrumentation/opentelemetry-instrumentation-dbapi/CHANGELOG.md +++ b/instrumentation/opentelemetry-instrumentation-dbapi/CHANGELOG.md @@ -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 \ No newline at end of file +- Initial release diff --git a/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py b/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py index 035c823bc..551f71555 100644 --- a/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py @@ -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)