Skip to content

Commit a301748

Browse files
alexander-alderman-webbalexander-alderman-webb
committed
merge master
2 parents e89f581 + 2e15294 commit a301748

34 files changed

Lines changed: 1104 additions & 389 deletions

‎sentry_sdk/ai/_openai_completions_api.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ def _transform_tool_definitions(
8282
continue
8383

8484
if tool["type"] == "function":
85-
tool_definition: ToolDefinition = {
85+
tool_definition: "ToolDefinition" = {
8686
"type": "function",
8787
}
8888

‎sentry_sdk/ai/_openai_responses_api.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -72,9 +72,6 @@ def _transform_tool_definitions(tools: "Iterable[ToolParam]") -> "list[ToolDefin
7272
tool_definitions.append(tool_definition)
7373
continue
7474

75-
if "type" not in tool:
76-
continue
77-
7875
tool_definitions.append({"type": tool["type"]})
7976

8077
return tool_definitions

‎sentry_sdk/integrations/clickhouse_driver.py‎

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from sentry_sdk.traces import StreamedSpan
99
from sentry_sdk.tracing import Span
1010
from sentry_sdk.tracing_utils import has_span_streaming_enabled
11-
from sentry_sdk.utils import capture_internal_exceptions
11+
from sentry_sdk.utils import capture_internal_exceptions, has_data_collection_enabled
1212

1313
# Hack to get new Python features working in older versions
1414
# without introducing a hard dependency on `typing_extensions`
@@ -107,8 +107,12 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T":
107107
if query_id:
108108
span.set_data("db.query_id", query_id)
109109

110-
if params and should_send_default_pii():
111-
span.set_data("db.params", params)
110+
if params:
111+
if has_data_collection_enabled(client.options):
112+
if client.options["data_collection"]["database_query_data"]:
113+
span.set_data("db.params", params)
114+
elif should_send_default_pii():
115+
span.set_data("db.params", params)
112116

113117
connection._sentry_span = span # type: ignore[attr-defined]
114118

@@ -135,8 +139,13 @@ def _inner_end(*args: "P.args", **kwargs: "P.kwargs") -> "T":
135139
if isinstance(span, StreamedSpan):
136140
span.end()
137141
else:
138-
if res is not None and should_send_default_pii():
139-
span.set_data("db.result", res)
142+
if res is not None:
143+
client_options = sentry_sdk.get_client().options
144+
if has_data_collection_enabled(client_options):
145+
if client_options["data_collection"]["database_query_data"]:
146+
span.set_data("db.result", res)
147+
elif should_send_default_pii():
148+
span.set_data("db.result", res)
140149

141150
with capture_internal_exceptions():
142151
span.scope.add_breadcrumb(
@@ -167,7 +176,29 @@ def _inner_send_data( # type: ignore[no-untyped-def] # clickhouse-driver does n
167176
if span is not None:
168177
_set_db_data(span, self.connection)
169178

170-
if should_send_default_pii():
179+
client_options = sentry_sdk.get_client().options
180+
if has_data_collection_enabled(client_options):
181+
if client_options["data_collection"]["database_query_data"]:
182+
db_params = span._data.get("db.params", [])
183+
if isinstance(data, (list, tuple)):
184+
db_params.extend(data)
185+
186+
else: # data is a generic iterator
187+
orig_data = data
188+
189+
# Wrap the generator to add items to db.params as they are yielded.
190+
# This allows us to send the params to Sentry without needing to allocate
191+
# memory for the entire generator at once.
192+
def wrapped_generator() -> "Iterator[Any]":
193+
for item in orig_data:
194+
db_params.append(item)
195+
yield item
196+
197+
# Replace the original iterator with the wrapped one.
198+
data = wrapped_generator()
199+
200+
span.set_data("db.params", db_params)
201+
elif should_send_default_pii():
171202
db_params = span._data.get("db.params", [])
172203

173204
if isinstance(data, (list, tuple)):

‎sentry_sdk/integrations/openai_agents/spans/invoke_agent.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ def invoke_agent_span(
3838
SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent",
3939
},
4040
)
41+
42+
set_on_span = span.set_attribute
4143
else:
4244
start_span_function = get_start_span_function()
4345
span = start_span_function(
@@ -49,6 +51,8 @@ def invoke_agent_span(
4951

5052
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
5153

54+
set_on_span = span.set_data
55+
5256
if should_send_default_pii():
5357
messages = []
5458
if agent.instructions:
@@ -97,6 +101,12 @@ def invoke_agent_span(
97101

98102
_set_agent_data(span, agent)
99103

104+
if len(agent.tools) > 0:
105+
set_on_span(
106+
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
107+
safe_serialize([vars(tool) for tool in agent.tools]),
108+
)
109+
100110
return span
101111

102112

‎sentry_sdk/integrations/pymongo.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from sentry_sdk.traces import SpanStatus, StreamedSpan
99
from sentry_sdk.tracing import Span
1010
from sentry_sdk.tracing_utils import has_span_streaming_enabled
11-
from sentry_sdk.utils import capture_internal_exceptions
11+
from sentry_sdk.utils import capture_internal_exceptions, has_data_collection_enabled
1212

1313
try:
1414
from pymongo import monitoring
@@ -146,7 +146,10 @@ def started(self, event: "CommandStartedEvent") -> None:
146146
db_name = event.database_name
147147

148148
lsid = command.pop("lsid", None)
149-
if not should_send_default_pii():
149+
if has_data_collection_enabled(client.options):
150+
if not client.options["data_collection"]["database_query_data"]:
151+
command = _strip_pii(command)
152+
elif not should_send_default_pii():
150153
command = _strip_pii(command)
151154

152155
query = json.dumps(command, default=str)

‎sentry_sdk/integrations/redis/utils.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from typing import TYPE_CHECKING
22

3+
import sentry_sdk
34
from sentry_sdk.consts import SPANDATA
45
from sentry_sdk.integrations.redis.consts import (
56
_COMMANDS_INCLUDING_SENSITIVE_DATA,
@@ -11,7 +12,7 @@
1112
from sentry_sdk.scope import should_send_default_pii
1213
from sentry_sdk.traces import StreamedSpan
1314
from sentry_sdk.tracing import Span
14-
from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE
15+
from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE, has_data_collection_enabled
1516

1617
if TYPE_CHECKING:
1718
from typing import Any, Optional, Sequence, Union
@@ -22,6 +23,7 @@ def _get_safe_command(name: str, args: "Sequence[Any]") -> str:
2223

2324
name_low = name.lower()
2425
send_default_pii = should_send_default_pii()
26+
client_options = sentry_sdk.get_client().options
2527

2628
for i, arg in enumerate(args):
2729
if i > _MAX_NUM_ARGS:
@@ -35,7 +37,10 @@ def _get_safe_command(name: str, args: "Sequence[Any]") -> str:
3537
if arg_is_the_key:
3638
command_parts.append(repr(arg))
3739
else:
38-
if send_default_pii:
40+
if has_data_collection_enabled(client_options):
41+
if client_options["data_collection"]["database_query_data"]:
42+
command_parts.append(repr(arg))
43+
elif send_default_pii:
3944
command_parts.append(repr(arg))
4045
else:
4146
command_parts.append(SENSITIVE_DATA_SUBSTITUTE)

‎tests/integrations/aiohttp/test_aiohttp.py‎

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1904,10 +1904,7 @@ async def test_server_url_query_data_collection_span_streaming(
19041904
sentry_init(
19051905
integrations=[AioHttpIntegration()],
19061906
traces_sample_rate=1.0,
1907-
_experiments={
1908-
"trace_lifecycle": "stream",
1909-
**init_kwargs.pop("_experiments", {}),
1910-
},
1907+
trace_lifecycle="stream",
19111908
**init_kwargs,
19121909
)
19131910

@@ -1949,10 +1946,7 @@ async def test_client_url_query_data_collection_span_streaming(
19491946
sentry_init(
19501947
integrations=[AioHttpIntegration()],
19511948
traces_sample_rate=1.0,
1952-
_experiments={
1953-
"trace_lifecycle": "stream",
1954-
**init_kwargs.pop("_experiments", {}),
1955-
},
1949+
trace_lifecycle="stream",
19561950
**init_kwargs,
19571951
)
19581952

‎tests/integrations/aiomysql/test_aiomysql.py‎

Lines changed: 8 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -441,9 +441,7 @@ async def test_query_source_disabled(
441441
"traces_sample_rate": 1.0,
442442
"enable_db_query_source": False,
443443
"db_query_source_threshold_ms": 0,
444-
"_experiments": {
445-
"trace_lifecycle": "stream" if span_streaming else "static",
446-
},
444+
"trace_lifecycle": "stream" if span_streaming else "static",
447445
}
448446

449447
sentry_init(**sentry_options)
@@ -516,9 +514,7 @@ async def test_query_source_enabled(
516514
"integrations": [AioMySQLIntegration()],
517515
"traces_sample_rate": 1.0,
518516
"db_query_source_threshold_ms": 0,
519-
"_experiments": {
520-
"trace_lifecycle": "stream" if span_streaming else "static",
521-
},
517+
"trace_lifecycle": "stream" if span_streaming else "static",
522518
}
523519
if enable_db_query_source is not None:
524520
sentry_options["enable_db_query_source"] = enable_db_query_source
@@ -591,9 +587,7 @@ async def test_query_source(sentry_init, capture_events, capture_items, span_str
591587
traces_sample_rate=1.0,
592588
enable_db_query_source=True,
593589
db_query_source_threshold_ms=0,
594-
_experiments={
595-
"trace_lifecycle": "stream" if span_streaming else "static",
596-
},
590+
trace_lifecycle="stream" if span_streaming else "static",
597591
)
598592

599593
if span_streaming:
@@ -692,9 +686,7 @@ async def test_no_query_source_if_duration_too_short(
692686
traces_sample_rate=1.0,
693687
enable_db_query_source=True,
694688
db_query_source_threshold_ms=100,
695-
_experiments={
696-
"trace_lifecycle": "stream" if span_streaming else "static",
697-
},
689+
trace_lifecycle="stream" if span_streaming else "static",
698690
)
699691

700692
if span_streaming:
@@ -789,9 +781,7 @@ async def test_query_source_if_duration_over_threshold(
789781
traces_sample_rate=1.0,
790782
enable_db_query_source=True,
791783
db_query_source_threshold_ms=100,
792-
_experiments={
793-
"trace_lifecycle": "stream" if span_streaming else "static",
794-
},
784+
trace_lifecycle="stream" if span_streaming else "static",
795785
)
796786

797787
if span_streaming:
@@ -915,9 +905,7 @@ async def test_span_origin(sentry_init, capture_events, capture_items, span_stre
915905
sentry_init(
916906
integrations=[AioMySQLIntegration()],
917907
traces_sample_rate=1.0,
918-
_experiments={
919-
"trace_lifecycle": "stream" if span_streaming else "static",
920-
},
908+
trace_lifecycle="stream" if span_streaming else "static",
921909
)
922910

923911
if span_streaming:
@@ -972,9 +960,7 @@ async def test_multiline_query_description_normalized(
972960
sentry_init(
973961
integrations=[AioMySQLIntegration()],
974962
traces_sample_rate=1.0,
975-
_experiments={
976-
"trace_lifecycle": "stream" if span_streaming else "static",
977-
},
963+
trace_lifecycle="stream" if span_streaming else "static",
978964
)
979965

980966
if span_streaming:
@@ -1099,9 +1085,7 @@ async def test_db_data_on_spans(
10991085
sentry_init(
11001086
integrations=[AioMySQLIntegration()],
11011087
traces_sample_rate=1.0,
1102-
_experiments={
1103-
"trace_lifecycle": "stream" if span_streaming else "static",
1104-
},
1088+
trace_lifecycle="stream" if span_streaming else "static",
11051089
)
11061090

11071091
if span_streaming:

0 commit comments

Comments
 (0)