|
| 1 | +from collections.abc import Generator |
1 | 2 | import sentry_sdk
|
2 | 3 | from sentry_sdk.consts import OP, SPANDATA
|
3 | 4 | from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
|
@@ -49,9 +50,7 @@ def setup_once() -> None:
|
49 | 50 | )
|
50 | 51 |
|
51 | 52 | # 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() |
55 | 54 |
|
56 | 55 | # Every query ends either with the Client's `receive_end_of_query` (no result expected)
|
57 | 56 | # or its `receive_result` (result expected)
|
@@ -128,23 +127,33 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
|
128 | 127 | return _inner_end
|
129 | 128 |
|
130 | 129 |
|
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) |
136 | 143 |
|
137 | 144 | if span is not None:
|
138 |
| - _set_db_data(span, instance.connection) |
| 145 | + _set_db_data(span, self.connection) |
139 | 146 |
|
140 | 147 | if should_send_default_pii():
|
141 | 148 | db_params = span._data.get("db.params", [])
|
142 | 149 | db_params.extend(data)
|
143 | 150 | span.set_data("db.params", db_params)
|
144 | 151 |
|
145 |
| - return f(*args, **kwargs) |
| 152 | + return original_send_data( |
| 153 | + self, sample_block, data, types_check, columnar, *args, **kwargs |
| 154 | + ) |
146 | 155 |
|
147 |
| - return _inner_send_data |
| 156 | + clickhouse_driver.client.Client.send_data = _inner_send_data |
148 | 157 |
|
149 | 158 |
|
150 | 159 | def _set_db_data(
|
|
0 commit comments