Skip to content

Commit

Permalink
Merge branch 'potel-base' into potel-base-run-all-tests
Browse files Browse the repository at this point in the history
  • Loading branch information
sl0thentr0py committed Dec 4, 2024
2 parents 18ec5cc + c4f2a2a commit d7140f0
Show file tree
Hide file tree
Showing 6 changed files with 26 additions and 18 deletions.
13 changes: 8 additions & 5 deletions sentry_sdk/integrations/boto3.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,20 +116,19 @@ def _sentry_after_call(context, parsed, **kwargs):
data=span_data,
)

span.__exit__(None, None, None)

body = parsed.get("Body")
if not isinstance(body, StreamingBody):
span.__exit__(None, None, None)
return

streaming_span = span.start_child(
streaming_span = sentry_sdk.start_span(
op=OP.HTTP_CLIENT_STREAM,
name=span.description,
name=span.name,
origin=Boto3Integration.origin,
only_if_parent=True,
)

orig_read = body.read
orig_close = body.close

def sentry_streaming_body_read(*args, **kwargs):
# type: (*Any, **Any) -> bytes
Expand All @@ -144,13 +143,17 @@ def sentry_streaming_body_read(*args, **kwargs):

body.read = sentry_streaming_body_read

orig_close = body.close

def sentry_streaming_body_close(*args, **kwargs):
# type: (*Any, **Any) -> None
streaming_span.finish()
orig_close(*args, **kwargs)

body.close = sentry_streaming_body_close

span.__exit__(None, None, None)


def _sentry_after_call_error(context, exception, **kwargs):
# type: (Dict[str, Any], Type[BaseException], **Any) -> None
Expand Down
2 changes: 1 addition & 1 deletion sentry_sdk/integrations/huey.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,7 @@ def _sentry_execute(self, task, timestamp=None):

sentry_headers = task.kwargs.pop("sentry_headers", {})
with sentry_sdk.continue_trace(sentry_headers):
with sentry_sdk.start_transaction(
with sentry_sdk.start_span(
name=task.name,
op=OP.QUEUE_TASK_HUEY,
source=TRANSACTION_SOURCE_TASK,
Expand Down
7 changes: 6 additions & 1 deletion sentry_sdk/tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -1313,7 +1313,12 @@ def __exit__(self, ty, value, tb):
if value is not None:
self.set_status(SPANSTATUS.INTERNAL_ERROR)
else:
self.set_status(SPANSTATUS.OK)
status_unset = (
hasattr(self._otel_span, "status")
and self._otel_span.status.status_code == StatusCode.UNSET
)
if status_unset:
self.set_status(SPANSTATUS.OK)

self.finish()
context.detach(self._ctx_token)
Expand Down
12 changes: 6 additions & 6 deletions tests/integrations/boto3/test_s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ def test_basic(sentry_init, capture_events):
events = capture_events()

s3 = session.resource("s3")
with sentry_sdk.start_transaction() as transaction, MockResponse(
with sentry_sdk.start_span() as transaction, MockResponse(
s3.meta.client, 200, {}, read_fixture("s3_list.xml")
):
bucket = s3.Bucket("bucket")
Expand All @@ -45,7 +45,7 @@ def test_breadcrumb(sentry_init, capture_events):

try:
s3 = session.resource("s3")
with sentry_sdk.start_transaction(), MockResponse(
with sentry_sdk.start_span(), MockResponse(
s3.meta.client, 200, {}, read_fixture("s3_list.xml")
):
bucket = s3.Bucket("bucket")
Expand Down Expand Up @@ -75,7 +75,7 @@ def test_streaming(sentry_init, capture_events):
events = capture_events()

s3 = session.resource("s3")
with sentry_sdk.start_transaction() as transaction, MockResponse(
with sentry_sdk.start_span() as transaction, MockResponse(
s3.meta.client, 200, {}, b"hello"
):
obj = s3.Bucket("bucket").Object("foo.pdf")
Expand Down Expand Up @@ -113,7 +113,7 @@ def test_streaming_close(sentry_init, capture_events):
events = capture_events()

s3 = session.resource("s3")
with sentry_sdk.start_transaction() as transaction, MockResponse(
with sentry_sdk.start_span() as transaction, MockResponse(
s3.meta.client, 200, {}, b"hello"
):
obj = s3.Bucket("bucket").Object("foo.pdf")
Expand Down Expand Up @@ -142,7 +142,7 @@ def test_omit_url_data_if_parsing_fails(sentry_init, capture_events):
"sentry_sdk.integrations.boto3.parse_url",
side_effect=ValueError,
):
with sentry_sdk.start_transaction() as transaction, MockResponse(
with sentry_sdk.start_span() as transaction, MockResponse(
s3.meta.client, 200, {}, read_fixture("s3_list.xml")
):
bucket = s3.Bucket("bucket")
Expand Down Expand Up @@ -170,7 +170,7 @@ def test_span_origin(sentry_init, capture_events):
events = capture_events()

s3 = session.resource("s3")
with sentry_sdk.start_transaction(), MockResponse(
with sentry_sdk.start_span(), MockResponse(
s3.meta.client, 200, {}, read_fixture("s3_list.xml")
):
bucket = s3.Bucket("bucket")
Expand Down
2 changes: 1 addition & 1 deletion tests/integrations/cohere/test_cohere.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ def test_bad_chat(sentry_init, capture_events):
with pytest.raises(httpx.HTTPError):
client.chat(model="some-model", message="hello")

(event, _) = events
(event,) = events
assert event["level"] == "error"


Expand Down
8 changes: 4 additions & 4 deletions tests/integrations/huey/test_huey.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import pytest
from decimal import DivisionByZero

from sentry_sdk import start_transaction
import sentry_sdk
from sentry_sdk.integrations.huey import HueyIntegration
from sentry_sdk.utils import parse_version

Expand Down Expand Up @@ -160,7 +160,7 @@ def dummy_task():

events = capture_events()

with start_transaction() as transaction:
with sentry_sdk.start_span() as transaction:
dummy_task()

(event,) = events
Expand All @@ -182,7 +182,7 @@ def test_huey_propagate_trace(init_huey, capture_events):
def propagated_trace_task():
pass

with start_transaction() as outer_transaction:
with sentry_sdk.start_span() as outer_transaction:
execute_huey_task(huey, propagated_trace_task)

assert (
Expand All @@ -200,7 +200,7 @@ def dummy_task():

events = capture_events()

with start_transaction():
with sentry_sdk.start_span():
dummy_task()

(event,) = events
Expand Down

0 comments on commit d7140f0

Please sign in to comment.