Skip to content

Commit a40f128

Browse files
antonpirkerantonpirker
authored
Use new scopes API in Django, SQLAlchemy, and asyncpg integration. (getsentry#2845)
Use new scopes API in Django, SQLAlchemy, and asyncpg integration.
1 parent 27d5ee1 commit a40f128

11 files changed

Lines changed: 157 additions & 196 deletions

File tree

‎sentry_sdk/integrations/asyncpg.py‎

Lines changed: 20 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,17 @@
22
import contextlib
33
from typing import Any, TypeVar, Callable, Awaitable, Iterator
44

5-
from sentry_sdk import Hub
5+
import sentry_sdk
66
from sentry_sdk.consts import OP, SPANDATA
77
from sentry_sdk.integrations import Integration, DidNotEnable
88
from sentry_sdk.tracing import Span
99
from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
10-
from sentry_sdk.utils import parse_version, capture_internal_exceptions
10+
from sentry_sdk.utils import (
11+
ensure_integration_enabled,
12+
ensure_integration_enabled_async,
13+
parse_version,
14+
capture_internal_exceptions,
15+
)
1116

1217
try:
1318
import asyncpg # type: ignore[import-not-found]
@@ -54,8 +59,7 @@ def setup_once() -> None:
5459

5560
def _wrap_execute(f: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
5661
async def _inner(*args: Any, **kwargs: Any) -> T:
57-
hub = Hub.current
58-
integration = hub.get_integration(AsyncPGIntegration)
62+
integration = sentry_sdk.get_client().get_integration(AsyncPGIntegration)
5963

6064
# Avoid recording calls to _execute twice.
6165
# Calls to Connection.execute with args also call
@@ -65,13 +69,11 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
6569
return await f(*args, **kwargs)
6670

6771
query = args[1]
68-
with record_sql_queries(
69-
hub, None, query, None, None, executemany=False
70-
) as span:
72+
with record_sql_queries(None, query, None, None, executemany=False) as span:
7173
res = await f(*args, **kwargs)
7274

7375
with capture_internal_exceptions():
74-
add_query_source(hub, span)
76+
add_query_source(span)
7577

7678
return res
7779

@@ -83,21 +85,19 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
8385

8486
@contextlib.contextmanager
8587
def _record(
86-
hub: Hub,
8788
cursor: SubCursor | None,
8889
query: str,
8990
params_list: tuple[Any, ...] | None,
9091
*,
9192
executemany: bool = False,
9293
) -> Iterator[Span]:
93-
integration = hub.get_integration(AsyncPGIntegration)
94-
if not integration._record_params:
94+
integration = sentry_sdk.get_client().get_integration(AsyncPGIntegration)
95+
if integration is not None and not integration._record_params:
9596
params_list = None
9697

9798
param_style = "pyformat" if params_list else None
9899

99100
with record_sql_queries(
100-
hub,
101101
cursor,
102102
query,
103103
params_list,
@@ -111,16 +111,11 @@ def _record(
111111
def _wrap_connection_method(
112112
f: Callable[..., Awaitable[T]], *, executemany: bool = False
113113
) -> Callable[..., Awaitable[T]]:
114+
@ensure_integration_enabled_async(AsyncPGIntegration, f)
114115
async def _inner(*args: Any, **kwargs: Any) -> T:
115-
hub = Hub.current
116-
integration = hub.get_integration(AsyncPGIntegration)
117-
118-
if integration is None:
119-
return await f(*args, **kwargs)
120-
121116
query = args[1]
122117
params_list = args[2] if len(args) > 2 else None
123-
with _record(hub, None, query, params_list, executemany=executemany) as span:
118+
with _record(None, query, params_list, executemany=executemany) as span:
124119
_set_db_data(span, args[0])
125120
res = await f(*args, **kwargs)
126121

@@ -130,18 +125,12 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
130125

131126

132127
def _wrap_cursor_creation(f: Callable[..., T]) -> Callable[..., T]:
128+
@ensure_integration_enabled(AsyncPGIntegration, f)
133129
def _inner(*args: Any, **kwargs: Any) -> T: # noqa: N807
134-
hub = Hub.current
135-
integration = hub.get_integration(AsyncPGIntegration)
136-
137-
if integration is None:
138-
return f(*args, **kwargs)
139-
140130
query = args[1]
141131
params_list = args[2] if len(args) > 2 else None
142132

143133
with _record(
144-
hub,
145134
None,
146135
query,
147136
params_list,
@@ -157,17 +146,12 @@ def _inner(*args: Any, **kwargs: Any) -> T: # noqa: N807
157146

158147

159148
def _wrap_connect_addr(f: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
149+
@ensure_integration_enabled_async(AsyncPGIntegration, f)
160150
async def _inner(*args: Any, **kwargs: Any) -> T:
161-
hub = Hub.current
162-
integration = hub.get_integration(AsyncPGIntegration)
163-
164-
if integration is None:
165-
return await f(*args, **kwargs)
166-
167151
user = kwargs["params"].user
168152
database = kwargs["params"].database
169153

170-
with hub.start_span(op=OP.DB, description="connect") as span:
154+
with sentry_sdk.start_span(op=OP.DB, description="connect") as span:
171155
span.set_data(SPANDATA.DB_SYSTEM, "postgresql")
172156
addr = kwargs.get("addr")
173157
if addr:
@@ -180,7 +164,9 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
180164
span.set_data(SPANDATA.DB_USER, user)
181165

182166
with capture_internal_exceptions():
183-
hub.add_breadcrumb(message="connect", category="query", data=span._data)
167+
sentry_sdk.add_breadcrumb(
168+
message="connect", category="query", data=span._data
169+
)
184170
res = await f(*args, **kwargs)
185171

186172
return res

‎sentry_sdk/integrations/django/__init__.py‎

Lines changed: 34 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,11 @@
44
import weakref
55
from importlib import import_module
66

7+
import sentry_sdk
78
from sentry_sdk._types import TYPE_CHECKING
89
from sentry_sdk.consts import OP, SPANDATA
910
from sentry_sdk.db.explain_plan.django import attach_explain_plan_to_span
10-
from sentry_sdk.hub import Hub, _should_send_default_pii
11-
from sentry_sdk.scope import Scope, add_global_event_processor
11+
from sentry_sdk.scope import Scope, add_global_event_processor, should_send_default_pii
1212
from sentry_sdk.serializer import add_global_repr_processor
1313
from sentry_sdk.tracing import SOURCE_FOR_STYLE, TRANSACTION_SOURCE_URL
1414
from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
@@ -19,6 +19,7 @@
1919
SENSITIVE_DATA_SUBSTITUTE,
2020
logger,
2121
capture_internal_exceptions,
22+
ensure_integration_enabled,
2223
event_from_exception,
2324
transaction_from_function,
2425
walk_exception_chain,
@@ -146,11 +147,9 @@ def setup_once():
146147

147148
old_app = WSGIHandler.__call__
148149

150+
@ensure_integration_enabled(DjangoIntegration, old_app)
149151
def sentry_patched_wsgi_handler(self, environ, start_response):
150152
# type: (Any, Dict[str, str], Callable[..., Any]) -> _ScopedResponse
151-
if Hub.current.get_integration(DjangoIntegration) is None:
152-
return old_app(self, environ, start_response)
153-
154153
bound_old_app = old_app.__get__(self, WSGIHandler)
155154

156155
from django.conf import settings
@@ -229,11 +228,6 @@ def _django_queryset_repr(value, hint):
229228
if not isinstance(value, QuerySet) or value._result_cache:
230229
return NotImplemented
231230

232-
# Do not call Hub.get_integration here. It is intentional that
233-
# running under a new hub does not suddenly start executing
234-
# querysets. This might be surprising to the user but it's likely
235-
# less annoying.
236-
237231
return "<%s from %s at 0x%x>" % (
238232
value.__class__.__name__,
239233
value.__module__,
@@ -400,8 +394,8 @@ def _set_transaction_name_and_source(scope, transaction_style, request):
400394

401395
def _before_get_response(request):
402396
# type: (WSGIRequest) -> None
403-
hub = Hub.current
404-
integration = hub.get_integration(DjangoIntegration)
397+
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
398+
405399
if integration is None:
406400
return
407401

@@ -431,8 +425,7 @@ def _attempt_resolve_again(request, scope, transaction_style):
431425

432426
def _after_get_response(request):
433427
# type: (WSGIRequest) -> None
434-
hub = Hub.current
435-
integration = hub.get_integration(DjangoIntegration)
428+
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
436429
if integration is None or integration.transaction_style != "url":
437430
return
438431

@@ -490,7 +483,7 @@ def wsgi_request_event_processor(event, hint):
490483
with capture_internal_exceptions():
491484
DjangoRequestExtractor(request).extract_into_event(event)
492485

493-
if _should_send_default_pii():
486+
if should_send_default_pii():
494487
with capture_internal_exceptions():
495488
_set_user_info(request, event)
496489

@@ -501,22 +494,19 @@ def wsgi_request_event_processor(event, hint):
501494

502495
def _got_request_exception(request=None, **kwargs):
503496
# type: (WSGIRequest, **Any) -> None
504-
hub = Hub.current
505-
integration = hub.get_integration(DjangoIntegration)
497+
client = sentry_sdk.get_client()
498+
integration = client.get_integration(DjangoIntegration)
506499
if integration is not None:
507500
if request is not None and integration.transaction_style == "url":
508501
scope = Scope.get_current_scope()
509502
_attempt_resolve_again(request, scope, integration.transaction_style)
510503

511-
# If an integration is there, a client has to be there.
512-
client = hub.client # type: Any
513-
514504
event, hint = event_from_exception(
515505
sys.exc_info(),
516506
client_options=client.options,
517507
mechanism={"type": "django", "handled": False},
518508
)
519-
hub.capture_event(event, hint=hint)
509+
sentry_sdk.capture_event(event, hint=hint)
520510

521511

522512
class DjangoRequestExtractor(RequestExtractor):
@@ -612,62 +602,56 @@ def install_sql_hook():
612602
# This won't work on Django versions < 1.6
613603
return
614604

605+
@ensure_integration_enabled(DjangoIntegration, real_execute)
615606
def execute(self, sql, params=None):
616607
# type: (CursorWrapper, Any, Optional[Any]) -> Any
617-
hub = Hub.current
618-
if hub.get_integration(DjangoIntegration) is None:
619-
return real_execute(self, sql, params)
620-
621608
with record_sql_queries(
622-
hub, self.cursor, sql, params, paramstyle="format", executemany=False
609+
self.cursor, sql, params, paramstyle="format", executemany=False
623610
) as span:
624611
_set_db_data(span, self)
625-
if hub.client:
626-
options = hub.client.options["_experiments"].get("attach_explain_plans")
627-
if options is not None:
628-
attach_explain_plan_to_span(
629-
span,
630-
self.cursor.connection,
631-
sql,
632-
params,
633-
self.mogrify,
634-
options,
635-
)
612+
options = (
613+
sentry_sdk.get_client()
614+
.options["_experiments"]
615+
.get("attach_explain_plans")
616+
)
617+
if options is not None:
618+
attach_explain_plan_to_span(
619+
span,
620+
self.cursor.connection,
621+
sql,
622+
params,
623+
self.mogrify,
624+
options,
625+
)
636626
result = real_execute(self, sql, params)
637627

638628
with capture_internal_exceptions():
639-
add_query_source(hub, span)
629+
add_query_source(span)
640630

641631
return result
642632

633+
@ensure_integration_enabled(DjangoIntegration, real_executemany)
643634
def executemany(self, sql, param_list):
644635
# type: (CursorWrapper, Any, List[Any]) -> Any
645-
hub = Hub.current
646-
if hub.get_integration(DjangoIntegration) is None:
647-
return real_executemany(self, sql, param_list)
648-
649636
with record_sql_queries(
650-
hub, self.cursor, sql, param_list, paramstyle="format", executemany=True
637+
self.cursor, sql, param_list, paramstyle="format", executemany=True
651638
) as span:
652639
_set_db_data(span, self)
653640

654641
result = real_executemany(self, sql, param_list)
655642

656643
with capture_internal_exceptions():
657-
add_query_source(hub, span)
644+
add_query_source(span)
658645

659646
return result
660647

648+
@ensure_integration_enabled(DjangoIntegration, real_connect)
661649
def connect(self):
662650
# type: (BaseDatabaseWrapper) -> None
663-
hub = Hub.current
664-
if hub.get_integration(DjangoIntegration) is None:
665-
return real_connect(self)
666-
667651
with capture_internal_exceptions():
668-
hub.add_breadcrumb(message="connect", category="query")
652+
sentry_sdk.add_breadcrumb(message="connect", category="query")
669653

670-
with hub.start_span(op=OP.DB, description="connect") as span:
654+
with sentry_sdk.start_span(op=OP.DB, description="connect") as span:
671655
_set_db_data(span, self)
672656
return real_connect(self)
673657

@@ -679,7 +663,6 @@ def connect(self):
679663

680664
def _set_db_data(span, cursor_or_db):
681665
# type: (Span, Any) -> None
682-
683666
db = cursor_or_db.db if hasattr(cursor_or_db, "db") else cursor_or_db
684667
vendor = db.vendor
685668
span.set_data(SPANDATA.DB_SYSTEM, vendor)

0 commit comments

Comments
 (0)