From f719972e9419d5fb2f23eb16c944e5909f7dc80c Mon Sep 17 00:00:00 2001 From: Ivana Kellyer Date: Wed, 23 Oct 2024 13:11:39 +0200 Subject: [PATCH 1/2] Fix asyncpg breadcrumbs --- sentry_sdk/integrations/asyncpg.py | 82 ++++++++++++---------- tests/integrations/asyncpg/test_asyncpg.py | 20 ++---- 2 files changed, 49 insertions(+), 53 deletions(-) diff --git a/sentry_sdk/integrations/asyncpg.py b/sentry_sdk/integrations/asyncpg.py index c1f2557a20..21458f8b61 100644 --- a/sentry_sdk/integrations/asyncpg.py +++ b/sentry_sdk/integrations/asyncpg.py @@ -1,6 +1,6 @@ from __future__ import annotations import contextlib -from typing import Any, TypeVar, Callable, Awaitable, Iterator +from typing import Any, TypeVar, Callable, Awaitable, Iterator, Optional import sentry_sdk from sentry_sdk.consts import OP, SPANDATA @@ -21,6 +21,7 @@ except ImportError: raise DidNotEnable("asyncpg not installed.") + # asyncpg.__version__ is a string containing the semantic version in the form of ".." asyncpg_version = parse_version(asyncpg.__version__) @@ -123,10 +124,13 @@ def _wrap_connection_method( async def _inner(*args: Any, **kwargs: Any) -> T: if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None: return await f(*args, **kwargs) + query = args[1] params_list = args[2] if len(args) > 2 else None + with _record(None, query, params_list, executemany=executemany) as span: - _set_db_data(span, args[0]) + data = _get_db_data(conn=args[0]) + _set_on_span(span, data) res = await f(*args, **kwargs) return res @@ -146,7 +150,8 @@ def _inner(*args: Any, **kwargs: Any) -> T: # noqa: N807 params_list, executemany=False, ) as span: - _set_db_data(span, args[0]) + data = _get_db_data(conn=args[0]) + _set_on_span(span, data) res = f(*args, **kwargs) span.set_attribute("db.cursor", _serialize_span_attribute(res)) @@ -160,40 +165,23 @@ async def _inner(*args: Any, **kwargs: Any) -> T: if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None: return await f(*args, **kwargs) - user = kwargs["params"].user - database = kwargs["params"].database - with sentry_sdk.start_span( op=OP.DB, name="connect", origin=AsyncPGIntegration.origin, ) as span: - span.set_attribute(SPANDATA.DB_SYSTEM, "postgresql") - addr = kwargs.get("addr") - if addr: - try: - span.set_attribute(SPANDATA.SERVER_ADDRESS, addr[0]) - span.set_attribute(SPANDATA.SERVER_PORT, addr[1]) - except IndexError: - pass - - span.set_attribute(SPANDATA.DB_NAME, database) - span.set_attribute(SPANDATA.DB_USER, user) + data = _get_db_data( + addr=kwargs.get("addr"), + database=kwargs["params"].database, + user=kwargs["params"].user, + ) + _set_on_span(span, data) with capture_internal_exceptions(): - data = {} - for attr in ( - "db.cursor", - "db.params", - "db.paramstyle", - SPANDATA.DB_NAME, - SPANDATA.DB_SYSTEM, - SPANDATA.DB_USER, - SPANDATA.SERVER_ADDRESS, - SPANDATA.SERVER_PORT, - ): - if span.get_attribute(attr): - data[attr] = span.get_attribute(attr) + # XXX + # "db.cursor", + # "db.params", + # "db.paramstyle", sentry_sdk.add_breadcrumb( message="connect", category="query", data=data @@ -206,21 +194,37 @@ async def _inner(*args: Any, **kwargs: Any) -> T: return _inner -def _set_db_data(span: Span, conn: Any) -> None: - span.set_attribute(SPANDATA.DB_SYSTEM, "postgresql") +def _get_db_data( + conn: Any = None, + addr: Optional[tuple[str]] = None, + database: Optional[str] = None, + user: Optional[str] = None, +) -> dict[str, str]: + if conn is not None: + addr = conn._addr + database = conn._params.database + user = conn._params.user + + data = { + SPANDATA.DB_SYSTEM: "postgresql", + } - addr = conn._addr if addr: try: - span.set_attribute(SPANDATA.SERVER_ADDRESS, addr[0]) - span.set_attribute(SPANDATA.SERVER_PORT, addr[1]) + data[SPANDATA.SERVER_ADDRESS] = addr[0] + data[SPANDATA.SERVER_PORT] = addr[1] except IndexError: pass - database = conn._params.database if database: - span.set_attribute(SPANDATA.DB_NAME, database) + data[SPANDATA.DB_NAME] = database - user = conn._params.user if user: - span.set_attribute(SPANDATA.DB_USER, user) + data[SPANDATA.DB_USER] = user + + return data + + +def _set_on_span(span: Span, data: dict[str, Any]): + for key, value in data.items(): + span.set_attribute(key, value) diff --git a/tests/integrations/asyncpg/test_asyncpg.py b/tests/integrations/asyncpg/test_asyncpg.py index adeef37d38..aae80770bf 100644 --- a/tests/integrations/asyncpg/test_asyncpg.py +++ b/tests/integrations/asyncpg/test_asyncpg.py @@ -74,7 +74,6 @@ async def _clean_pg(): async def test_connect(sentry_init, capture_events) -> None: sentry_init( integrations=[AsyncPGIntegration()], - traces_sample_rate=1.0, _experiments={"record_sql_params": True}, ) events = capture_events() @@ -97,7 +96,6 @@ async def test_connect(sentry_init, capture_events) -> None: async def test_execute(sentry_init, capture_events) -> None: sentry_init( integrations=[AsyncPGIntegration()], - traces_sample_rate=1.0, _experiments={"record_sql_params": True}, ) events = capture_events() @@ -163,7 +161,6 @@ async def test_execute(sentry_init, capture_events) -> None: async def test_execute_many(sentry_init, capture_events) -> None: sentry_init( integrations=[AsyncPGIntegration()], - traces_sample_rate=1.0, _experiments={"record_sql_params": True}, ) events = capture_events() @@ -202,7 +199,6 @@ async def test_execute_many(sentry_init, capture_events) -> None: async def test_record_params(sentry_init, capture_events) -> None: sentry_init( integrations=[AsyncPGIntegration(record_params=True)], - traces_sample_rate=1.0, _experiments={"record_sql_params": True}, ) events = capture_events() @@ -243,7 +239,6 @@ async def test_record_params(sentry_init, capture_events) -> None: async def test_cursor(sentry_init, capture_events) -> None: sentry_init( integrations=[AsyncPGIntegration()], - traces_sample_rate=1.0, _experiments={"record_sql_params": True}, ) events = capture_events() @@ -308,7 +303,6 @@ async def test_cursor(sentry_init, capture_events) -> None: async def test_cursor_manual(sentry_init, capture_events) -> None: sentry_init( integrations=[AsyncPGIntegration()], - traces_sample_rate=1.0, _experiments={"record_sql_params": True}, ) events = capture_events() @@ -375,7 +369,6 @@ async def test_cursor_manual(sentry_init, capture_events) -> None: async def test_prepared_stmt(sentry_init, capture_events) -> None: sentry_init( integrations=[AsyncPGIntegration()], - traces_sample_rate=1.0, _experiments={"record_sql_params": True}, ) events = capture_events() @@ -425,7 +418,6 @@ async def test_prepared_stmt(sentry_init, capture_events) -> None: async def test_connection_pool(sentry_init, capture_events) -> None: sentry_init( integrations=[AsyncPGIntegration()], - traces_sample_rate=1.0, _experiments={"record_sql_params": True}, ) events = capture_events() @@ -497,7 +489,7 @@ async def test_connection_pool(sentry_init, capture_events) -> None: async def test_query_source_disabled(sentry_init, capture_events): sentry_options = { "integrations": [AsyncPGIntegration()], - "enable_tracing": True, + "traces_sample_rate": 1.0, "enable_db_query_source": False, "db_query_source_threshold_ms": 0, } @@ -535,7 +527,7 @@ async def test_query_source_enabled( ): sentry_options = { "integrations": [AsyncPGIntegration()], - "enable_tracing": True, + "traces_sample_rate": 1.0, "db_query_source_threshold_ms": 0, } if enable_db_query_source is not None: @@ -571,7 +563,7 @@ async def test_query_source_enabled( async def test_query_source(sentry_init, capture_events): sentry_init( integrations=[AsyncPGIntegration()], - enable_tracing=True, + traces_sample_rate=1.0, enable_db_query_source=True, db_query_source_threshold_ms=0, ) @@ -621,7 +613,7 @@ async def test_query_source_with_module_in_search_path(sentry_init, capture_even """ sentry_init( integrations=[AsyncPGIntegration()], - enable_tracing=True, + traces_sample_rate=1.0, enable_db_query_source=True, db_query_source_threshold_ms=0, ) @@ -667,7 +659,7 @@ async def test_query_source_with_module_in_search_path(sentry_init, capture_even async def test_no_query_source_if_duration_too_short(sentry_init, capture_events): sentry_init( integrations=[AsyncPGIntegration()], - enable_tracing=True, + traces_sample_rate=1.0, enable_db_query_source=True, db_query_source_threshold_ms=100, ) @@ -714,7 +706,7 @@ def fake_record_sql_queries(*args, **kwargs): async def test_query_source_if_duration_over_threshold(sentry_init, capture_events): sentry_init( integrations=[AsyncPGIntegration()], - enable_tracing=True, + traces_sample_rate=1.0, enable_db_query_source=True, db_query_source_threshold_ms=100, ) From 29be6db855ee85fb5d91e7bda8fca592094a5325 Mon Sep 17 00:00:00 2001 From: Ivana Kellyer Date: Wed, 23 Oct 2024 13:13:38 +0200 Subject: [PATCH 2/2] . --- sentry_sdk/integrations/asyncpg.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/sentry_sdk/integrations/asyncpg.py b/sentry_sdk/integrations/asyncpg.py index 21458f8b61..1113fd199c 100644 --- a/sentry_sdk/integrations/asyncpg.py +++ b/sentry_sdk/integrations/asyncpg.py @@ -178,11 +178,6 @@ async def _inner(*args: Any, **kwargs: Any) -> T: _set_on_span(span, data) with capture_internal_exceptions(): - # XXX - # "db.cursor", - # "db.params", - # "db.paramstyle", - sentry_sdk.add_breadcrumb( message="connect", category="query", data=data )