1515# from: https://stackoverflow.com/a/71944042/300572
1616if TYPE_CHECKING :
1717 from collections .abc import Iterator
18- from typing import Any , Callable , ParamSpec , Union
18+ from typing import Any , Callable , Optional , ParamSpec , Union
1919else :
2020 # Fake ParamSpec
2121 class ParamSpec :
@@ -79,7 +79,7 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T":
7979 if client .get_integration (ClickhouseDriverIntegration ) is None :
8080 return f (* args , ** kwargs )
8181
82- connection = args [0 ]
82+ connection : "Connection" = args [0 ]
8383 query = args [1 ]
8484 query_id = args [2 ] if len (args ) > 2 else kwargs .get ("query_id" )
8585 params = args [3 ] if len (args ) > 3 else kwargs .get ("params" )
@@ -95,6 +95,16 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T":
9595 SPANDATA .DB_QUERY_TEXT : str (query ),
9696 },
9797 )
98+
99+ connection ._query = query
100+ connection ._breadcrumb_data = {
101+ SPANDATA .DB_SYSTEM : "clickhouse" ,
102+ SPANDATA .DB_NAME : connection .database ,
103+ SPANDATA .DB_DRIVER_NAME : "clickhouse-driver" ,
104+ SPANDATA .SERVER_ADDRESS : connection .host ,
105+ SPANDATA .SERVER_PORT : connection .port ,
106+ SPANDATA .DB_USER : connection .user ,
107+ }
98108 else :
99109 span = sentry_sdk .start_span (
100110 op = OP .DB ,
@@ -114,7 +124,7 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T":
114124 elif should_send_default_pii ():
115125 span .set_data ("db.params" , params )
116126
117- connection ._sentry_span = span # type: ignore[attr-defined]
127+ connection ._sentry_span = span
118128
119129 if span is not None :
120130 _set_db_data (span , connection )
@@ -130,8 +140,31 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T":
130140def _wrap_end (f : "Callable[P, T]" ) -> "Callable[P, T]" :
131141 def _inner_end (* args : "P.args" , ** kwargs : "P.kwargs" ) -> "T" :
132142 res = f (* args , ** kwargs )
133- instance = args [0 ]
134- span = getattr (instance .connection , "_sentry_span" , None ) # type: ignore[attr-defined]
143+ instance : "Client" = args [0 ]
144+
145+ query = getattr (instance .connection , "_query" , None )
146+ breadcrumb_data : "Optional[dict[str, Any]]" = getattr (
147+ instance .connection , "_breadcrumb_data" , None
148+ )
149+
150+ if query is not None and breadcrumb_data is not None :
151+ client_options = sentry_sdk .get_client ().options
152+ if (
153+ has_data_collection_enabled (client_options )
154+ and client_options ["data_collection" ]["database_query_data" ]
155+ ) or (
156+ not has_data_collection_enabled (client_options )
157+ and should_send_default_pii ()
158+ ):
159+ breadcrumb_data = {"db.result" : res , ** breadcrumb_data }
160+
161+ sentry_sdk .get_isolation_scope ().add_breadcrumb (
162+ message = query ,
163+ category = "query" ,
164+ data = {"db.result" : res , ** breadcrumb_data },
165+ )
166+
167+ span = getattr (instance .connection , "_sentry_span" , None )
135168
136169 if span is None :
137170 return res
0 commit comments