Skip to content

Commit f76b786

Browse files
ref(clickhouse): List send_data parameters
Explicitly list the `send_data` parameters in the wrapped function. The parameters are coming from [here](https://github.com/mymarilyn/clickhouse-driver/blob/8a4e7c5b99b532df2b015651d893a6f36288a22c/clickhouse_driver/client.py#L634). Continue also providing `*args` and `**kwargs`, but only for forwards-compatibility.
1 parent 84adbb7 commit f76b786

File tree

1 file changed

+20
-11
lines changed

1 file changed

+20
-11
lines changed

sentry_sdk/integrations/clickhouse_driver.py

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from collections.abc import Generator
12
import sentry_sdk
23
from sentry_sdk.consts import OP, SPANDATA
34
from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
@@ -49,9 +50,7 @@ def setup_once() -> None:
4950
)
5051

5152
# 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-
)
53+
_wrap_send_data()
5554

5655
# Every query ends either with the Client's `receive_end_of_query` (no result expected)
5756
# or its `receive_result` (result expected)
@@ -128,23 +127,33 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
128127
return _inner_end
129128

130129

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)
130+
def _wrap_send_data():
131+
original_send_data = clickhouse_driver.client.Client.send_data
132+
133+
def _inner_send_data(
134+
self: clickhouse_driver.client.Client,
135+
sample_block: object,
136+
data: list | tuple | Generator,
137+
types_check: bool = False,
138+
columnar: bool = False,
139+
*args,
140+
**kwargs,
141+
) -> int:
142+
span = getattr(self.connection, "_sentry_span", None)
136143

137144
if span is not None:
138-
_set_db_data(span, instance.connection)
145+
_set_db_data(span, self.connection)
139146

140147
if should_send_default_pii():
141148
db_params = span._data.get("db.params", [])
142149
db_params.extend(data)
143150
span.set_data("db.params", db_params)
144151

145-
return f(*args, **kwargs)
152+
return original_send_data(
153+
self, sample_block, data, types_check, columnar, *args, **kwargs
154+
)
146155

147-
return _inner_send_data
156+
clickhouse_driver.client.Client.send_data = _inner_send_data
148157

149158

150159
def _set_db_data(

0 commit comments

Comments
 (0)