1- from sentry_sdk import Hub
1+ import sentry_sdk
22from sentry_sdk .consts import OP , SPANDATA
3- from sentry_sdk .hub import _should_send_default_pii
43from sentry_sdk .integrations import Integration , DidNotEnable
54from sentry_sdk .tracing import Span
65from sentry_sdk ._types import TYPE_CHECKING
6+ from sentry_sdk .scope import should_send_default_pii
77from sentry_sdk .utils import capture_internal_exceptions
88
99from typing import TypeVar
@@ -75,15 +75,14 @@ def setup_once() -> None:
7575
7676def _wrap_start (f : Callable [P , T ]) -> Callable [P , T ]:
7777 def _inner (* args : P .args , ** kwargs : P .kwargs ) -> T :
78- hub = Hub .current
79- if hub .get_integration (ClickhouseDriverIntegration ) is None :
78+ if sentry_sdk .get_client ().get_integration (ClickhouseDriverIntegration ) is None :
8079 return f (* args , ** kwargs )
8180 connection = args [0 ]
8281 query = args [1 ]
8382 query_id = args [2 ] if len (args ) > 2 else kwargs .get ("query_id" )
8483 params = args [3 ] if len (args ) > 3 else kwargs .get ("params" )
8584
86- span = hub .start_span (op = OP .DB , description = query )
85+ span = sentry_sdk .start_span (op = OP .DB , description = query )
8786
8887 connection ._sentry_span = span # type: ignore[attr-defined]
8988
@@ -94,7 +93,7 @@ def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
9493 if query_id :
9594 span .set_data ("db.query_id" , query_id )
9695
97- if params and _should_send_default_pii ():
96+ if params and should_send_default_pii ():
9897 span .set_data ("db.params" , params )
9998
10099 # run the original code
@@ -112,7 +111,7 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
112111 span = instance .connection ._sentry_span # type: ignore[attr-defined]
113112
114113 if span is not None :
115- if res is not None and _should_send_default_pii ():
114+ if res is not None and should_send_default_pii ():
116115 span .set_data ("db.result" , res )
117116
118117 with capture_internal_exceptions ():
@@ -135,7 +134,7 @@ def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T:
135134
136135 _set_db_data (span , instance .connection )
137136
138- if _should_send_default_pii ():
137+ if should_send_default_pii ():
139138 db_params = span ._data .get ("db.params" , [])
140139 db_params .extend (data )
141140 span .set_data ("db.params" , db_params )
0 commit comments