Skip to content

Commit 965904d

Browse files
alexmojakialexmojaki
authored
Fix websocket span messages (#426)
1 parent f2093b3 commit 965904d

File tree

3 files changed

+171
-18
lines changed

3 files changed

+171
-18
lines changed

‎logfire/_internal/exporters/processor_wrapper.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,9 @@ def _tweak_http_spans(span: ReadableSpanDict):
196196
except Exception: # pragma: no cover
197197
pass
198198

199+
if not method and name in ('HTTP', f'HTTP {target}', f'HTTP {route}'):
200+
method = 'HTTP'
201+
199202
# Build up a list of possible span names and messages in order from worst to best
200203
names: list[str] = []
201204
messages: list[str] = []

‎tests/otel_integrations/test_fastapi.py‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,14 @@ async def echo_body(request: Request):
7373
return await request.body()
7474

7575

76+
async def websocket_endpoint(websocket: WebSocket, name: str):
77+
logfire.info('websocket_endpoint: {name}', name=name)
78+
await websocket.accept()
79+
assert (await websocket.receive_text()) == 'ping'
80+
await websocket.send_text('pong')
81+
await websocket.close()
82+
83+
7684
@pytest.fixture()
7785
def app():
7886
# Don't define the endpoint functions in this fixture to prevent a qualname with <locals> in it
@@ -90,6 +98,7 @@ def app():
9098
app.get('/validation_error')(validation_error)
9199
app.get('/with_path_param/{param}')(with_path_param)
92100
app.get('/secret/{path_param}', name='secret')(get_secret)
101+
app.websocket('/ws/{name}')(websocket_endpoint)
93102
first_lvl_app.get('/other', name='other_route_name', operation_id='other_route_operation_id')(other_route)
94103
second_lvl_app.get('/other', name='other_route_name', operation_id='other_route_operation_id')(other_route)
95104
return app
@@ -1818,3 +1827,144 @@ def test_request_hooks_with_send_receive_spans(exporter: TestExporter):
18181827
},
18191828
]
18201829
)
1830+
1831+
1832+
def test_websocket(client: TestClient, exporter: TestExporter) -> None:
1833+
with client.websocket_connect('/ws/foo') as websocket:
1834+
websocket.send_text('ping')
1835+
data = websocket.receive_text()
1836+
assert data == 'pong'
1837+
1838+
assert exporter.exported_spans_as_dict() == snapshot(
1839+
[
1840+
{
1841+
'name': 'FastAPI arguments',
1842+
'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
1843+
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
1844+
'start_time': 2000000000,
1845+
'end_time': 3000000000,
1846+
'attributes': {
1847+
'logfire.msg_template': 'FastAPI arguments',
1848+
'logfire.msg': 'FastAPI arguments',
1849+
'logfire.span_type': 'span',
1850+
'logfire.level_num': 5,
1851+
},
1852+
},
1853+
{
1854+
'name': 'websocket_endpoint: {name}',
1855+
'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
1856+
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
1857+
'start_time': 4000000000,
1858+
'end_time': 4000000000,
1859+
'attributes': {
1860+
'logfire.span_type': 'log',
1861+
'logfire.level_num': 9,
1862+
'logfire.msg_template': 'websocket_endpoint: {name}',
1863+
'logfire.msg': 'websocket_endpoint: foo',
1864+
'code.filepath': 'test_fastapi.py',
1865+
'code.function': 'websocket_endpoint',
1866+
'code.lineno': 123,
1867+
'name': 'foo',
1868+
'logfire.json_schema': '{"type":"object","properties":{"name":{}}}',
1869+
},
1870+
},
1871+
{
1872+
'name': 'HTTP /ws/{name} websocket receive connect',
1873+
'context': {'trace_id': 1, 'span_id': 6, 'is_remote': False},
1874+
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
1875+
'start_time': 5000000000,
1876+
'end_time': 6000000000,
1877+
'attributes': {
1878+
'logfire.span_type': 'span',
1879+
'logfire.msg': 'HTTP /ws/{name} websocket receive connect',
1880+
'logfire.level_num': 5,
1881+
'asgi.event.type': 'websocket.connect',
1882+
},
1883+
},
1884+
{
1885+
'name': 'HTTP /ws/{name} websocket send accept',
1886+
'context': {'trace_id': 1, 'span_id': 8, 'is_remote': False},
1887+
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
1888+
'start_time': 7000000000,
1889+
'end_time': 8000000000,
1890+
'attributes': {
1891+
'logfire.span_type': 'span',
1892+
'logfire.msg': 'HTTP /ws/{name} websocket send accept',
1893+
'logfire.level_num': 5,
1894+
'asgi.event.type': 'websocket.accept',
1895+
},
1896+
},
1897+
{
1898+
'name': 'HTTP /ws/{name} websocket receive',
1899+
'context': {'trace_id': 1, 'span_id': 10, 'is_remote': False},
1900+
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
1901+
'start_time': 9000000000,
1902+
'end_time': 10000000000,
1903+
'attributes': {
1904+
'logfire.span_type': 'span',
1905+
'logfire.msg': 'HTTP /ws/{name} websocket receive',
1906+
'logfire.level_num': 5,
1907+
'http.status_code': 200,
1908+
'http.response.status_code': 200,
1909+
'asgi.event.type': 'websocket.receive',
1910+
},
1911+
},
1912+
{
1913+
'name': 'HTTP /ws/{name} websocket send',
1914+
'context': {'trace_id': 1, 'span_id': 12, 'is_remote': False},
1915+
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
1916+
'start_time': 11000000000,
1917+
'end_time': 12000000000,
1918+
'attributes': {
1919+
'logfire.span_type': 'span',
1920+
'logfire.msg': 'HTTP /ws/{name} websocket send',
1921+
'logfire.level_num': 5,
1922+
'asgi.event.type': 'websocket.send',
1923+
'http.status_code': 200,
1924+
'http.response.status_code': 200,
1925+
},
1926+
},
1927+
{
1928+
'name': 'HTTP /ws/{name} websocket send close',
1929+
'context': {'trace_id': 1, 'span_id': 14, 'is_remote': False},
1930+
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
1931+
'start_time': 13000000000,
1932+
'end_time': 14000000000,
1933+
'attributes': {
1934+
'logfire.span_type': 'span',
1935+
'logfire.msg': 'HTTP /ws/{name} websocket send close',
1936+
'logfire.level_num': 5,
1937+
'asgi.event.type': 'websocket.close',
1938+
},
1939+
},
1940+
{
1941+
'name': 'HTTP /ws/{name}',
1942+
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
1943+
'parent': None,
1944+
'start_time': 1000000000,
1945+
'end_time': 15000000000,
1946+
'attributes': {
1947+
'logfire.span_type': 'span',
1948+
'logfire.msg': 'HTTP /ws/foo',
1949+
'http.scheme': 'ws',
1950+
'url.scheme': 'ws',
1951+
'http.host': 'testserver',
1952+
'client.address': 'testserver',
1953+
'net.host.port': 80,
1954+
'server.port': 80,
1955+
'http.target': '/ws/foo',
1956+
'url.path': '/ws/foo',
1957+
'http.url': 'ws://testserver/ws/foo',
1958+
'http.server_name': 'testserver',
1959+
'http.user_agent': 'testclient',
1960+
'user_agent.original': 'testclient',
1961+
'net.peer.ip': 'testclient',
1962+
'net.peer.port': 50000,
1963+
'client.port': 50000,
1964+
'http.route': '/ws/{name}',
1965+
'http.status_code': 200,
1966+
'http.response.status_code': 200,
1967+
},
1968+
},
1969+
]
1970+
)

‎tests/otel_integrations/test_starlette.py‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ async def websocket_endpoint(websocket: WebSocket):
2828
def app():
2929
routes = [
3030
Route('/secret/{path_param}', secret),
31-
WebSocketRoute('/ws', websocket_endpoint),
31+
WebSocketRoute('/ws/{name}', websocket_endpoint),
3232
]
3333

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

4848

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

5555
assert exporter.exported_spans_as_dict() == snapshot(
5656
[
5757
{
58-
'name': '/ws websocket receive connect',
58+
'name': '/ws/{name} websocket receive connect',
5959
'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
6060
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
6161
'start_time': 2000000000,
6262
'end_time': 3000000000,
6363
'attributes': {
6464
'logfire.span_type': 'span',
65-
'logfire.msg': '/ws websocket receive connect',
65+
'logfire.msg': '/ws/{name} websocket receive connect',
6666
'asgi.event.type': 'websocket.connect',
6767
'logfire.level_num': 5,
6868
},
6969
},
7070
{
71-
'name': '/ws websocket send accept',
71+
'name': '/ws/{name} websocket send accept',
7272
'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
7373
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
7474
'start_time': 4000000000,
7575
'end_time': 5000000000,
7676
'attributes': {
7777
'logfire.span_type': 'span',
78-
'logfire.msg': '/ws websocket send accept',
78+
'logfire.msg': '/ws/{name} websocket send accept',
7979
'asgi.event.type': 'websocket.accept',
8080
'logfire.level_num': 5,
8181
},
8282
},
8383
{
84-
'name': '/ws websocket receive',
84+
'name': '/ws/{name} websocket receive',
8585
'context': {'trace_id': 1, 'span_id': 7, 'is_remote': False},
8686
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
8787
'start_time': 6000000000,
8888
'end_time': 7000000000,
8989
'attributes': {
9090
'logfire.span_type': 'span',
91-
'logfire.msg': '/ws websocket receive',
91+
'logfire.msg': '/ws/{name} websocket receive',
9292
'http.status_code': 200,
9393
'asgi.event.type': 'websocket.receive',
9494
'http.response.status_code': 200,
9595
'logfire.level_num': 5,
9696
},
9797
},
9898
{
99-
'name': '/ws websocket send',
99+
'name': '/ws/{name} websocket send',
100100
'context': {'trace_id': 1, 'span_id': 9, 'is_remote': False},
101101
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
102102
'start_time': 8000000000,
103103
'end_time': 9000000000,
104104
'attributes': {
105105
'logfire.span_type': 'span',
106-
'logfire.msg': '/ws websocket send',
106+
'logfire.msg': '/ws/{name} websocket send',
107107
'http.status_code': 200,
108108
'asgi.event.type': 'websocket.send',
109109
'logfire.level_num': 5,
110110
'http.response.status_code': 200,
111111
},
112112
},
113113
{
114-
'name': '/ws websocket send close',
114+
'name': '/ws/{name} websocket send close',
115115
'context': {'trace_id': 1, 'span_id': 11, 'is_remote': False},
116116
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
117117
'start_time': 10000000000,
118118
'end_time': 11000000000,
119119
'attributes': {
120120
'logfire.span_type': 'span',
121-
'logfire.msg': '/ws websocket send close',
121+
'logfire.msg': '/ws/{name} websocket send close',
122122
'asgi.event.type': 'websocket.close',
123123
'logfire.level_num': 5,
124124
},
125125
},
126126
{
127-
'name': '/ws',
127+
'name': '/ws/{name}',
128128
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
129129
'parent': None,
130130
'start_time': 1000000000,
131131
'end_time': 12000000000,
132132
'attributes': {
133133
'logfire.span_type': 'span',
134-
'logfire.msg': '/ws',
134+
'logfire.msg': '/ws/foo',
135135
'http.scheme': 'ws',
136136
'url.scheme': 'ws',
137137
'http.host': 'testserver',
138138
'net.host.port': 80,
139139
'server.port': 80,
140-
'http.target': '/ws',
141-
'url.path': '/ws',
142-
'http.url': 'ws://testserver/ws',
140+
'http.target': '/ws/foo',
141+
'url.path': '/ws/foo',
142+
'http.url': 'ws://testserver/ws/foo',
143143
'http.server_name': 'testserver',
144144
'http.user_agent': 'testclient',
145145
'user_agent.original': 'testclient',
146146
'net.peer.ip': 'testclient',
147147
'client.address': 'testserver',
148148
'net.peer.port': 50000,
149149
'client.port': 50000,
150-
'http.route': '/ws',
150+
'http.route': '/ws/{name}',
151151
'http.request.header.host': ('testserver',),
152152
'http.request.header.accept': ('*/*',),
153153
'http.request.header.accept_encoding': ('gzip, deflate',),

0 commit comments

Comments
 (0)