From 580acc72689c387a24a5cbe74e07dcf2f35da1bf Mon Sep 17 00:00:00 2001 From: Alex Hall Date: Mon, 16 Sep 2024 13:26:25 +0200 Subject: [PATCH] Fix websocket span messages --- .../_internal/exporters/processor_wrapper.py | 3 + tests/otel_integrations/test_fastapi.py | 150 ++++++++++++++++++ tests/otel_integrations/test_starlette.py | 36 ++--- 3 files changed, 171 insertions(+), 18 deletions(-) diff --git a/logfire/_internal/exporters/processor_wrapper.py b/logfire/_internal/exporters/processor_wrapper.py index fd7a0d408..2d77e7fe3 100644 --- a/logfire/_internal/exporters/processor_wrapper.py +++ b/logfire/_internal/exporters/processor_wrapper.py @@ -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] = [] diff --git a/tests/otel_integrations/test_fastapi.py b/tests/otel_integrations/test_fastapi.py index b3b60f375..a7d37f2a8 100644 --- a/tests/otel_integrations/test_fastapi.py +++ b/tests/otel_integrations/test_fastapi.py @@ -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 in it @@ -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 @@ -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, + }, + }, + ] + ) diff --git a/tests/otel_integrations/test_starlette.py b/tests/otel_integrations/test_starlette.py index a69ed3ca8..674a3959c 100644 --- a/tests/otel_integrations/test_starlette.py +++ b/tests/otel_integrations/test_starlette.py @@ -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) @@ -47,7 +47,7 @@ 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' @@ -55,40 +55,40 @@ def test_websocket(client: TestClient, exporter: TestExporter) -> None: 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, @@ -96,14 +96,14 @@ def test_websocket(client: TestClient, exporter: TestExporter) -> None: }, }, { - '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, @@ -111,35 +111,35 @@ def test_websocket(client: TestClient, exporter: TestExporter) -> None: }, }, { - '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', @@ -147,7 +147,7 @@ def test_websocket(client: TestClient, exporter: TestExporter) -> None: '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',),