Skip to content

Commit

Permalink
Fix grpc aio method handling (#3873)
Browse files Browse the repository at this point in the history
  • Loading branch information
sl0thentr0py authored Dec 16, 2024
1 parent bb9e34e commit 9d20ffd
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 64 deletions.
8 changes: 6 additions & 2 deletions sentry_sdk/integrations/grpc/aio/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,10 +44,12 @@ async def intercept_unary_unary(
request: Message,
) -> Union[UnaryUnaryCall, Message]:
method = client_call_details.method
if isinstance(method, bytes):
method = method.decode()

with sentry_sdk.start_span(
op=OP.GRPC_CLIENT,
name="unary unary call to %s" % method.decode(),
name="unary unary call to %s" % method,
origin=SPAN_ORIGIN,
only_if_parent=True,
) as span:
Expand Down Expand Up @@ -75,10 +77,12 @@ async def intercept_unary_stream(
request: Message,
) -> Union[AsyncIterable[Any], UnaryStreamCall]:
method = client_call_details.method
if isinstance(method, bytes):
method = method.decode()

with sentry_sdk.start_span(
op=OP.GRPC_CLIENT,
name="unary stream call to %s" % method.decode(),
name="unary stream call to %s" % method,
origin=SPAN_ORIGIN,
only_if_parent=True,
) as span:
Expand Down
2 changes: 1 addition & 1 deletion sentry_sdk/integrations/rust_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@

import json
from enum import Enum, auto
from typing import Any, Callable, Dict, Tuple, Optional
from typing import Any, Callable, Dict, Optional

import sentry_sdk
from sentry_sdk.integrations import Integration
Expand Down
67 changes: 32 additions & 35 deletions tests/integrations/grpc/test_grpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from typing import List, Optional
from unittest.mock import Mock

from sentry_sdk import start_span, start_transaction
from sentry_sdk import start_span
from sentry_sdk.consts import OP
from sentry_sdk.integrations.grpc import GRPCIntegration
from tests.conftest import ApproxDict
Expand Down Expand Up @@ -41,7 +41,7 @@ def _tear_down(server: grpc.Server):


@pytest.mark.forked
def test_grpc_server_starts_transaction(sentry_init, capture_events_forksafe):
def test_grpc_server_starts_root_span(sentry_init, capture_events_forksafe):
sentry_init(traces_sample_rate=1.0, integrations=[GRPCIntegration()])
events = capture_events_forksafe()

Expand Down Expand Up @@ -99,7 +99,7 @@ def test_grpc_server_other_interceptors(sentry_init, capture_events_forksafe):


@pytest.mark.forked
def test_grpc_server_continues_transaction(sentry_init, capture_events_forksafe):
def test_grpc_server_continues_trace(sentry_init, capture_events_forksafe):
sentry_init(traces_sample_rate=1.0, integrations=[GRPCIntegration()])
events = capture_events_forksafe()

Expand All @@ -108,20 +108,20 @@ def test_grpc_server_continues_transaction(sentry_init, capture_events_forksafe)
with grpc.insecure_channel("localhost:{}".format(PORT)) as channel:
stub = gRPCTestServiceStub(channel)

with start_transaction() as transaction:
with start_span() as root_span:
metadata = (
(
"baggage",
"sentry-trace_id={trace_id},sentry-environment=test,"
"sentry-transaction=test-transaction,sentry-sample_rate=1.0".format(
trace_id=transaction.trace_id
trace_id=root_span.trace_id
),
),
(
"sentry-trace",
"{trace_id}-{parent_span_id}-{sampled}".format(
trace_id=transaction.trace_id,
parent_span_id=transaction.span_id,
trace_id=root_span.trace_id,
parent_span_id=root_span.span_id,
sampled=1,
),
),
Expand All @@ -139,7 +139,7 @@ def test_grpc_server_continues_transaction(sentry_init, capture_events_forksafe)
"source": "custom",
}
assert event["contexts"]["trace"]["op"] == OP.GRPC_SERVER
assert event["contexts"]["trace"]["trace_id"] == transaction.trace_id
assert event["contexts"]["trace"]["trace_id"] == root_span.trace_id
assert span["op"] == "test"


Expand All @@ -153,17 +153,17 @@ def test_grpc_client_starts_span(sentry_init, capture_events_forksafe):
with grpc.insecure_channel("localhost:{}".format(PORT)) as channel:
stub = gRPCTestServiceStub(channel)

with start_transaction():
with start_span():
stub.TestServe(gRPCTestMessage(text="test"))

_tear_down(server=server)

events.write_file.close()
events.read_event()
local_transaction = events.read_event()
span = local_transaction["spans"][0]
local_root_span = events.read_event()
span = local_root_span["spans"][0]

assert len(local_transaction["spans"]) == 1
assert len(local_root_span["spans"]) == 1
assert span["op"] == OP.GRPC_CLIENT
assert (
span["description"]
Expand All @@ -188,16 +188,16 @@ def test_grpc_client_unary_stream_starts_span(sentry_init, capture_events_forksa
with grpc.insecure_channel("localhost:{}".format(PORT)) as channel:
stub = gRPCTestServiceStub(channel)

with start_transaction():
with start_span():
[el for el in stub.TestUnaryStream(gRPCTestMessage(text="test"))]

_tear_down(server=server)

events.write_file.close()
local_transaction = events.read_event()
span = local_transaction["spans"][0]
local_root_span = events.read_event()
span = local_root_span["spans"][0]

assert len(local_transaction["spans"]) == 1
assert len(local_root_span["spans"]) == 1
assert span["op"] == OP.GRPC_CLIENT
assert (
span["description"]
Expand Down Expand Up @@ -233,7 +233,7 @@ def test_grpc_client_other_interceptor(sentry_init, capture_events_forksafe):
channel = grpc.intercept_channel(channel, MockClientInterceptor())
stub = gRPCTestServiceStub(channel)

with start_transaction():
with start_span():
stub.TestServe(gRPCTestMessage(text="test"))

_tear_down(server=server)
Expand All @@ -242,10 +242,10 @@ def test_grpc_client_other_interceptor(sentry_init, capture_events_forksafe):

events.write_file.close()
events.read_event()
local_transaction = events.read_event()
span = local_transaction["spans"][0]
local_root_span = events.read_event()
span = local_root_span["spans"][0]

assert len(local_transaction["spans"]) == 1
assert len(local_root_span["spans"]) == 1
assert span["op"] == OP.GRPC_CLIENT
assert (
span["description"]
Expand All @@ -272,18 +272,18 @@ def test_grpc_client_and_servers_interceptors_integration(
with grpc.insecure_channel("localhost:{}".format(PORT)) as channel:
stub = gRPCTestServiceStub(channel)

with start_transaction():
with start_span():
stub.TestServe(gRPCTestMessage(text="test"))

_tear_down(server=server)

events.write_file.close()
server_transaction = events.read_event()
local_transaction = events.read_event()
server_root_span = events.read_event()
local_root_span = events.read_event()

assert (
server_transaction["contexts"]["trace"]["trace_id"]
== local_transaction["contexts"]["trace"]["trace_id"]
server_root_span["contexts"]["trace"]["trace_id"]
== local_root_span["contexts"]["trace"]["trace_id"]
)


Expand Down Expand Up @@ -328,26 +328,23 @@ def test_span_origin(sentry_init, capture_events_forksafe):
with grpc.insecure_channel("localhost:{}".format(PORT)) as channel:
stub = gRPCTestServiceStub(channel)

with start_transaction(name="custom_transaction"):
with start_span(name="custom_root"):
stub.TestServe(gRPCTestMessage(text="test"))

_tear_down(server=server)

events.write_file.close()

transaction_from_integration = events.read_event()
custom_transaction = events.read_event()
root_span_from_integration = events.read_event()
custom_root_span = events.read_event()

assert root_span_from_integration["contexts"]["trace"]["origin"] == "auto.grpc.grpc"
assert (
transaction_from_integration["contexts"]["trace"]["origin"] == "auto.grpc.grpc"
)
assert (
transaction_from_integration["spans"][0]["origin"]
== "auto.grpc.grpc.TestService"
root_span_from_integration["spans"][0]["origin"] == "auto.grpc.grpc.TestService"
) # manually created in TestService, not the instrumentation

assert custom_transaction["contexts"]["trace"]["origin"] == "manual"
assert custom_transaction["spans"][0]["origin"] == "auto.grpc.grpc"
assert custom_root_span["contexts"]["trace"]["origin"] == "manual"
assert custom_root_span["spans"][0]["origin"] == "auto.grpc.grpc"


class TestService(gRPCTestServiceServicer):
Expand Down
48 changes: 23 additions & 25 deletions tests/integrations/grpc/test_grpc_aio.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import pytest_asyncio
import sentry_sdk

from sentry_sdk import start_span, start_transaction
from sentry_sdk import start_span
from sentry_sdk.consts import OP
from sentry_sdk.integrations.grpc import GRPCIntegration
from tests.conftest import ApproxDict
Expand Down Expand Up @@ -60,7 +60,7 @@ async def test_noop_for_unimplemented_method(sentry_init, capture_events):


@pytest.mark.asyncio
async def test_grpc_server_starts_transaction(grpc_server, capture_events):
async def test_grpc_server_starts_root_span(grpc_server, capture_events):
events = capture_events()

async with grpc.aio.insecure_channel("localhost:{}".format(AIO_PORT)) as channel:
Expand All @@ -79,26 +79,26 @@ async def test_grpc_server_starts_transaction(grpc_server, capture_events):


@pytest.mark.asyncio
async def test_grpc_server_continues_transaction(grpc_server, capture_events):
async def test_grpc_server_continues_trace(grpc_server, capture_events):
events = capture_events()

async with grpc.aio.insecure_channel("localhost:{}".format(AIO_PORT)) as channel:
stub = gRPCTestServiceStub(channel)

with sentry_sdk.start_transaction() as transaction:
with sentry_sdk.start_span() as root_span:
metadata = (
(
"baggage",
"sentry-trace_id={trace_id},sentry-environment=test,"
"sentry-transaction=test-transaction,sentry-sample_rate=1.0".format(
trace_id=transaction.trace_id
trace_id=root_span.trace_id
),
),
(
"sentry-trace",
"{trace_id}-{parent_span_id}-{sampled}".format(
trace_id=transaction.trace_id,
parent_span_id=transaction.span_id,
trace_id=root_span.trace_id,
parent_span_id=root_span.span_id,
sampled=1,
),
),
Expand All @@ -114,7 +114,7 @@ async def test_grpc_server_continues_transaction(grpc_server, capture_events):
"source": "custom",
}
assert event["contexts"]["trace"]["op"] == OP.GRPC_SERVER
assert event["contexts"]["trace"]["trace_id"] == transaction.trace_id
assert event["contexts"]["trace"]["trace_id"] == root_span.trace_id
assert span["op"] == "test"


Expand Down Expand Up @@ -159,15 +159,15 @@ async def test_grpc_client_starts_span(grpc_server, capture_events_forksafe):

async with grpc.aio.insecure_channel("localhost:{}".format(AIO_PORT)) as channel:
stub = gRPCTestServiceStub(channel)
with start_transaction():
with start_span():
await stub.TestServe(gRPCTestMessage(text="test"))

events.write_file.close()
events.read_event()
local_transaction = events.read_event()
span = local_transaction["spans"][0]
local_root_span = events.read_event()
span = local_root_span["spans"][0]

assert len(local_transaction["spans"]) == 1
assert len(local_root_span["spans"]) == 1
assert span["op"] == OP.GRPC_CLIENT
assert (
span["description"]
Expand All @@ -190,15 +190,15 @@ async def test_grpc_client_unary_stream_starts_span(

async with grpc.aio.insecure_channel("localhost:{}".format(AIO_PORT)) as channel:
stub = gRPCTestServiceStub(channel)
with start_transaction():
with start_span():
response = stub.TestUnaryStream(gRPCTestMessage(text="test"))
[_ async for _ in response]

events.write_file.close()
local_transaction = events.read_event()
span = local_transaction["spans"][0]
local_root_span = events.read_event()
span = local_root_span["spans"][0]

assert len(local_transaction["spans"]) == 1
assert len(local_root_span["spans"]) == 1
assert span["op"] == OP.GRPC_CLIENT
assert (
span["description"]
Expand Down Expand Up @@ -243,24 +243,22 @@ async def test_span_origin(grpc_server, capture_events_forksafe):

async with grpc.aio.insecure_channel("localhost:{}".format(AIO_PORT)) as channel:
stub = gRPCTestServiceStub(channel)
with start_transaction(name="custom_transaction"):
with start_span(name="custom_root"):
await stub.TestServe(gRPCTestMessage(text="test"))

events.write_file.close()

transaction_from_integration = events.read_event()
custom_transaction = events.read_event()
root_span_from_integration = events.read_event()
custom_root_span = events.read_event()

assert root_span_from_integration["contexts"]["trace"]["origin"] == "auto.grpc.grpc"
assert (
transaction_from_integration["contexts"]["trace"]["origin"] == "auto.grpc.grpc"
)
assert (
transaction_from_integration["spans"][0]["origin"]
root_span_from_integration["spans"][0]["origin"]
== "auto.grpc.grpc.TestService.aio"
) # manually created in TestService, not the instrumentation

assert custom_transaction["contexts"]["trace"]["origin"] == "manual"
assert custom_transaction["spans"][0]["origin"] == "auto.grpc.grpc"
assert custom_root_span["contexts"]["trace"]["origin"] == "manual"
assert custom_root_span["spans"][0]["origin"] == "auto.grpc.grpc"


class TestService(gRPCTestServiceServicer):
Expand Down
2 changes: 1 addition & 1 deletion tests/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -982,4 +982,4 @@ def test_serialize_span_attribute(value, result):
),
)
def test_datetime_from_isoformat(input_str, expected_output):
assert datetime_from_isoformat(input_str) == expected_output, input_str
assert datetime_from_isoformat(input_str) == expected_output, input_str

0 comments on commit 9d20ffd

Please sign in to comment.