feat(signalr): support on_client_connect hook
This commit is contained in:
@@ -48,7 +48,7 @@ class Client:
|
|||||||
message = await self.connection.receive()
|
message = await self.connection.receive()
|
||||||
d = message.get("bytes") or message.get("text", "").encode()
|
d = message.get("bytes") or message.get("text", "").encode()
|
||||||
if not d:
|
if not d:
|
||||||
raise WebSocketDisconnect(code=1008, reason="Empty message received.")
|
return PingPacket() # FIXME: Graceful empty message handling
|
||||||
return self.procotol.decode(d)
|
return self.procotol.decode(d)
|
||||||
|
|
||||||
async def _ping(self):
|
async def _ping(self):
|
||||||
@@ -104,6 +104,10 @@ class Hub:
|
|||||||
client._ping_task = task
|
client._ping_task = task
|
||||||
return client
|
return client
|
||||||
|
|
||||||
|
async def on_connect(self, client: Client) -> None:
|
||||||
|
if method := getattr(self, "on_client_connect", None):
|
||||||
|
await method(client)
|
||||||
|
|
||||||
async def remove_client(self, connection_id: str) -> None:
|
async def remove_client(self, connection_id: str) -> None:
|
||||||
if client := self.clients.get(connection_id):
|
if client := self.clients.get(connection_id):
|
||||||
del self.clients[connection_id]
|
del self.clients[connection_id]
|
||||||
@@ -138,6 +142,8 @@ class Hub:
|
|||||||
task = asyncio.create_task(self._handle_packet(client, packet))
|
task = asyncio.create_task(self._handle_packet(client, packet))
|
||||||
self.tasks.add(task)
|
self.tasks.add(task)
|
||||||
task.add_done_callback(self.tasks.discard)
|
task.add_done_callback(self.tasks.discard)
|
||||||
|
except StopIteration:
|
||||||
|
pass
|
||||||
except WebSocketDisconnect as e:
|
except WebSocketDisconnect as e:
|
||||||
print(
|
print(
|
||||||
f"Client {client.connection_id} disconnected: {e.code}, {e.reason}"
|
f"Client {client.connection_id} disconnected: {e.code}, {e.reason}"
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import json
|
import json
|
||||||
import time
|
import time
|
||||||
from typing import Literal
|
from typing import Literal
|
||||||
@@ -92,4 +93,7 @@ async def connect(
|
|||||||
if error or not client:
|
if error or not client:
|
||||||
await websocket.close(code=1008)
|
await websocket.close(code=1008)
|
||||||
return
|
return
|
||||||
|
task = asyncio.create_task(hub_.on_connect(client))
|
||||||
|
hub_.tasks.add(task)
|
||||||
|
task.add_done_callback(hub_.tasks.discard)
|
||||||
await hub_._listen_client(client)
|
await hub_._listen_client(client)
|
||||||
|
|||||||
Reference in New Issue
Block a user