Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix websocket span messages #426

Merged
merged 1 commit into from
Sep 16, 2024
Merged

Fix websocket span messages #426

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions logfire/_internal/exporters/processor_wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,9 @@ def _tweak_http_spans(span: ReadableSpanDict):
except Exception: # pragma: no cover
pass

if not method and name in ('HTTP', f'HTTP {target}', f'HTTP {route}'):
method = 'HTTP'

# Build up a list of possible span names and messages in order from worst to best
names: list[str] = []
messages: list[str] = []
Expand Down
150 changes: 150 additions & 0 deletions tests/otel_integrations/test_fastapi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,14 @@ async def echo_body(request: Request):
return await request.body()


async def websocket_endpoint(websocket: WebSocket, name: str):
logfire.info('websocket_endpoint: {name}', name=name)
await websocket.accept()
assert (await websocket.receive_text()) == 'ping'
await websocket.send_text('pong')
await websocket.close()


@pytest.fixture()
def app():
# Don't define the endpoint functions in this fixture to prevent a qualname with <locals> in it
Expand All @@ -90,6 +98,7 @@ def app():
app.get('/validation_error')(validation_error)
app.get('/with_path_param/{param}')(with_path_param)
app.get('/secret/{path_param}', name='secret')(get_secret)
app.websocket('/ws/{name}')(websocket_endpoint)
first_lvl_app.get('/other', name='other_route_name', operation_id='other_route_operation_id')(other_route)
second_lvl_app.get('/other', name='other_route_name', operation_id='other_route_operation_id')(other_route)
return app
Expand Down Expand Up @@ -1818,3 +1827,144 @@ def test_request_hooks_with_send_receive_spans(exporter: TestExporter):
},
]
)


def test_websocket(client: TestClient, exporter: TestExporter) -> None:
with client.websocket_connect('/ws/foo') as websocket:
websocket.send_text('ping')
data = websocket.receive_text()
assert data == 'pong'

assert exporter.exported_spans_as_dict() == snapshot(
[
{
'name': 'FastAPI arguments',
'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 2000000000,
'end_time': 3000000000,
'attributes': {
'logfire.msg_template': 'FastAPI arguments',
'logfire.msg': 'FastAPI arguments',
'logfire.span_type': 'span',
'logfire.level_num': 5,
},
},
{
'name': 'websocket_endpoint: {name}',
'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 4000000000,
'end_time': 4000000000,
'attributes': {
'logfire.span_type': 'log',
'logfire.level_num': 9,
'logfire.msg_template': 'websocket_endpoint: {name}',
'logfire.msg': 'websocket_endpoint: foo',
'code.filepath': 'test_fastapi.py',
'code.function': 'websocket_endpoint',
'code.lineno': 123,
'name': 'foo',
'logfire.json_schema': '{"type":"object","properties":{"name":{}}}',
},
},
{
'name': 'HTTP /ws/{name} websocket receive connect',
'context': {'trace_id': 1, 'span_id': 6, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 5000000000,
'end_time': 6000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/{name} websocket receive connect',
'logfire.level_num': 5,
'asgi.event.type': 'websocket.connect',
},
},
{
'name': 'HTTP /ws/{name} websocket send accept',
'context': {'trace_id': 1, 'span_id': 8, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 7000000000,
'end_time': 8000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/{name} websocket send accept',
'logfire.level_num': 5,
'asgi.event.type': 'websocket.accept',
},
},
{
'name': 'HTTP /ws/{name} websocket receive',
'context': {'trace_id': 1, 'span_id': 10, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 9000000000,
'end_time': 10000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/{name} websocket receive',
'logfire.level_num': 5,
'http.status_code': 200,
'http.response.status_code': 200,
'asgi.event.type': 'websocket.receive',
},
},
{
'name': 'HTTP /ws/{name} websocket send',
'context': {'trace_id': 1, 'span_id': 12, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 11000000000,
'end_time': 12000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/{name} websocket send',
'logfire.level_num': 5,
'asgi.event.type': 'websocket.send',
'http.status_code': 200,
'http.response.status_code': 200,
},
},
{
'name': 'HTTP /ws/{name} websocket send close',
'context': {'trace_id': 1, 'span_id': 14, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 13000000000,
'end_time': 14000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/{name} websocket send close',
'logfire.level_num': 5,
'asgi.event.type': 'websocket.close',
},
},
{
'name': 'HTTP /ws/{name}',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': 1000000000,
'end_time': 15000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/foo',
'http.scheme': 'ws',
'url.scheme': 'ws',
'http.host': 'testserver',
'client.address': 'testserver',
'net.host.port': 80,
'server.port': 80,
'http.target': '/ws/foo',
'url.path': '/ws/foo',
'http.url': 'ws://testserver/ws/foo',
'http.server_name': 'testserver',
'http.user_agent': 'testclient',
'user_agent.original': 'testclient',
'net.peer.ip': 'testclient',
'net.peer.port': 50000,
'client.port': 50000,
'http.route': '/ws/{name}',
'http.status_code': 200,
'http.response.status_code': 200,
},
},
]
)
36 changes: 18 additions & 18 deletions tests/otel_integrations/test_starlette.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ async def websocket_endpoint(websocket: WebSocket):
def app():
routes = [
Route('/secret/{path_param}', secret),
WebSocketRoute('/ws', websocket_endpoint),
WebSocketRoute('/ws/{name}', websocket_endpoint),
]

app = Starlette(routes=routes)
Expand All @@ -47,107 +47,107 @@ def client(app: Starlette) -> TestClient:


def test_websocket(client: TestClient, exporter: TestExporter) -> None:
with client.websocket_connect('/ws') as websocket:
with client.websocket_connect('/ws/foo') as websocket:
websocket.send_text('ping')
data = websocket.receive_text()
assert data == 'pong'

assert exporter.exported_spans_as_dict() == snapshot(
[
{
'name': '/ws websocket receive connect',
'name': '/ws/{name} websocket receive connect',
'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 2000000000,
'end_time': 3000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws websocket receive connect',
'logfire.msg': '/ws/{name} websocket receive connect',
'asgi.event.type': 'websocket.connect',
'logfire.level_num': 5,
},
},
{
'name': '/ws websocket send accept',
'name': '/ws/{name} websocket send accept',
'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 4000000000,
'end_time': 5000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws websocket send accept',
'logfire.msg': '/ws/{name} websocket send accept',
'asgi.event.type': 'websocket.accept',
'logfire.level_num': 5,
},
},
{
'name': '/ws websocket receive',
'name': '/ws/{name} websocket receive',
'context': {'trace_id': 1, 'span_id': 7, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 6000000000,
'end_time': 7000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws websocket receive',
'logfire.msg': '/ws/{name} websocket receive',
'http.status_code': 200,
'asgi.event.type': 'websocket.receive',
'http.response.status_code': 200,
'logfire.level_num': 5,
},
},
{
'name': '/ws websocket send',
'name': '/ws/{name} websocket send',
'context': {'trace_id': 1, 'span_id': 9, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 8000000000,
'end_time': 9000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws websocket send',
'logfire.msg': '/ws/{name} websocket send',
'http.status_code': 200,
'asgi.event.type': 'websocket.send',
'logfire.level_num': 5,
'http.response.status_code': 200,
},
},
{
'name': '/ws websocket send close',
'name': '/ws/{name} websocket send close',
'context': {'trace_id': 1, 'span_id': 11, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 10000000000,
'end_time': 11000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws websocket send close',
'logfire.msg': '/ws/{name} websocket send close',
'asgi.event.type': 'websocket.close',
'logfire.level_num': 5,
},
},
{
'name': '/ws',
'name': '/ws/{name}',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': 1000000000,
'end_time': 12000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws',
'logfire.msg': '/ws/foo',
'http.scheme': 'ws',
'url.scheme': 'ws',
'http.host': 'testserver',
'net.host.port': 80,
'server.port': 80,
'http.target': '/ws',
'url.path': '/ws',
'http.url': 'ws://testserver/ws',
'http.target': '/ws/foo',
'url.path': '/ws/foo',
'http.url': 'ws://testserver/ws/foo',
'http.server_name': 'testserver',
'http.user_agent': 'testclient',
'user_agent.original': 'testclient',
'net.peer.ip': 'testclient',
'client.address': 'testserver',
'net.peer.port': 50000,
'client.port': 50000,
'http.route': '/ws',
'http.route': '/ws/{name}',
'http.request.header.host': ('testserver',),
'http.request.header.accept': ('*/*',),
'http.request.header.accept_encoding': ('gzip, deflate',),
Expand Down