@@ -49,9 +49,7 @@ def setup_once() -> None:
4949 )
5050
5151 # If the query contains parameters then the send_data function is used to send those parameters to clickhouse
52- clickhouse_driver .client .Client .send_data = _wrap_send_data (
53- clickhouse_driver .client .Client .send_data
54- )
52+ _wrap_send_data ()
5553
5654 # Every query ends either with the Client's `receive_end_of_query` (no result expected)
5755 # or its `receive_result` (result expected)
@@ -128,23 +126,27 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
128126 return _inner_end
129127
130128
131- def _wrap_send_data (f : Callable [P , T ]) -> Callable [P , T ]:
132- def _inner_send_data (* args : P .args , ** kwargs : P .kwargs ) -> T :
133- instance = args [0 ] # type: clickhouse_driver.client.Client
134- data = args [2 ]
135- span = getattr (instance .connection , "_sentry_span" , None )
129+ def _wrap_send_data () -> None :
130+ original_send_data = clickhouse_driver .client .Client .send_data
131+
132+ def _inner_send_data ( # type: ignore[no-untyped-def] # clickhouse-driver does not type send_data
133+ self , sample_block , data , types_check = False , columnar = False , * args , ** kwargs
134+ ):
135+ span = getattr (self .connection , "_sentry_span" , None )
136136
137137 if span is not None :
138- _set_db_data (span , instance .connection )
138+ _set_db_data (span , self .connection )
139139
140140 if should_send_default_pii ():
141141 db_params = span ._data .get ("db.params" , [])
142142 db_params .extend (data )
143143 span .set_data ("db.params" , db_params )
144144
145- return f (* args , ** kwargs )
145+ return original_send_data (
146+ self , sample_block , data , types_check , columnar , * args , ** kwargs
147+ )
146148
147- return _inner_send_data
149+ clickhouse_driver . client . Client . send_data = _inner_send_data
148150
149151
150152def _set_db_data (
0 commit comments