@@ -49,9 +49,7 @@ def setup_once() -> None:
49
49
)
50
50
51
51
# 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 ()
55
53
56
54
# Every query ends either with the Client's `receive_end_of_query` (no result expected)
57
55
# or its `receive_result` (result expected)
@@ -128,23 +126,27 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
128
126
return _inner_end
129
127
130
128
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 )
136
136
137
137
if span is not None :
138
- _set_db_data (span , instance .connection )
138
+ _set_db_data (span , self .connection )
139
139
140
140
if should_send_default_pii ():
141
141
db_params = span ._data .get ("db.params" , [])
142
142
db_params .extend (data )
143
143
span .set_data ("db.params" , db_params )
144
144
145
- return f (* args , ** kwargs )
145
+ return original_send_data (
146
+ self , sample_block , data , types_check , columnar , * args , ** kwargs
147
+ )
146
148
147
- return _inner_send_data
149
+ clickhouse_driver . client . Client . send_data = _inner_send_data
148
150
149
151
150
152
def _set_db_data (
0 commit comments