|
1 | 1 | import re
|
2 |
| -import contextlib |
3 | 2 | import json
|
4 | 3 | import math
|
5 | 4 |
|
@@ -106,6 +105,58 @@ def __iter__(self):
|
106 | 105 | yield k[len(self.prefix) :]
|
107 | 106 |
|
108 | 107 |
|
| 108 | +class RecordSqlQueries: |
| 109 | + def __init__( |
| 110 | + self, |
| 111 | + hub, # type: sentry_sdk.Hub |
| 112 | + cursor, # type: Any |
| 113 | + query, # type: Any |
| 114 | + params_list, # type: Any |
| 115 | + paramstyle, # type: Optional[str] |
| 116 | + executemany, # type: bool |
| 117 | + ): |
| 118 | + # type: (...) -> None |
| 119 | + # TODO: Bring back capturing of params by default |
| 120 | + self._hub = hub |
| 121 | + if self._hub.client and self._hub.client.options["_experiments"].get( |
| 122 | + "record_sql_params", False |
| 123 | + ): |
| 124 | + if not params_list or params_list == [None]: |
| 125 | + params_list = None |
| 126 | + |
| 127 | + if paramstyle == "pyformat": |
| 128 | + paramstyle = "format" |
| 129 | + else: |
| 130 | + params_list = None |
| 131 | + paramstyle = None |
| 132 | + |
| 133 | + self._query = _format_sql(cursor, query) |
| 134 | + |
| 135 | + self._data = {} |
| 136 | + if params_list is not None: |
| 137 | + self._data["db.params"] = params_list |
| 138 | + if paramstyle is not None: |
| 139 | + self._data["db.paramstyle"] = paramstyle |
| 140 | + if executemany: |
| 141 | + self._data["db.executemany"] = True |
| 142 | + |
| 143 | + def __enter__(self): |
| 144 | + # type: () -> Span |
| 145 | + with capture_internal_exceptions(): |
| 146 | + self._hub.add_breadcrumb( |
| 147 | + message=self._query, category="query", data=self._data |
| 148 | + ) |
| 149 | + |
| 150 | + with self._hub.start_span(op="db", description=self._query) as span: |
| 151 | + for k, v in self._data.items(): |
| 152 | + span.set_data(k, v) |
| 153 | + return span |
| 154 | + |
| 155 | + def __exit__(self, exc_type, exc_val, exc_tb): |
| 156 | + # type: (Any, Any, Any) -> None |
| 157 | + pass |
| 158 | + |
| 159 | + |
109 | 160 | def has_tracing_enabled(options):
|
110 | 161 | # type: (Dict[str, Any]) -> bool
|
111 | 162 | """
|
@@ -150,49 +201,6 @@ def is_valid_sample_rate(rate):
|
150 | 201 | return True
|
151 | 202 |
|
152 | 203 |
|
153 |
| -@contextlib.contextmanager |
154 |
| -def record_sql_queries( |
155 |
| - hub, # type: sentry_sdk.Hub |
156 |
| - cursor, # type: Any |
157 |
| - query, # type: Any |
158 |
| - params_list, # type: Any |
159 |
| - paramstyle, # type: Optional[str] |
160 |
| - executemany, # type: bool |
161 |
| -): |
162 |
| - # type: (...) -> Generator[Span, None, None] |
163 |
| - |
164 |
| - # TODO: Bring back capturing of params by default |
165 |
| - if hub.client and hub.client.options["_experiments"].get( |
166 |
| - "record_sql_params", False |
167 |
| - ): |
168 |
| - if not params_list or params_list == [None]: |
169 |
| - params_list = None |
170 |
| - |
171 |
| - if paramstyle == "pyformat": |
172 |
| - paramstyle = "format" |
173 |
| - else: |
174 |
| - params_list = None |
175 |
| - paramstyle = None |
176 |
| - |
177 |
| - query = _format_sql(cursor, query) |
178 |
| - |
179 |
| - data = {} |
180 |
| - if params_list is not None: |
181 |
| - data["db.params"] = params_list |
182 |
| - if paramstyle is not None: |
183 |
| - data["db.paramstyle"] = paramstyle |
184 |
| - if executemany: |
185 |
| - data["db.executemany"] = True |
186 |
| - |
187 |
| - with capture_internal_exceptions(): |
188 |
| - hub.add_breadcrumb(message=query, category="query", data=data) |
189 |
| - |
190 |
| - with hub.start_span(op="db", description=query) as span: |
191 |
| - for k, v in data.items(): |
192 |
| - span.set_data(k, v) |
193 |
| - yield span |
194 |
| - |
195 |
| - |
196 | 204 | def maybe_create_breadcrumbs_from_span(hub, span):
|
197 | 205 | # type: (sentry_sdk.Hub, Span) -> None
|
198 | 206 | if span.op == "redis":
|
|
0 commit comments