happyz
/
fastapi
mirror of https://github.com/tiangolo/fastapi.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Export WebSocketDisconnect and add example handling disconnections to docs (#1822) 

Co-authored-by: Sebastián Ramírez <tiangolo@gmail.com>
This commit is contained in:
Rupsi Kaushik 2020-08-09 09:52:19 -04:00 committed by GitHub
parent f32d67c8b9
commit 5ed48ccdc8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 133 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
docs/en/docs/advanced/websockets.md
View File

@ -137,6 +137,33 @@ With that you can connect the WebSocket and then send and receive messages:
<img src="/img/tutorial/websockets/image05.png"> <img src="/img/tutorial/websockets/image05.png">
## Handling disconnections and multiple clients
When a WebSocket connection is closed, the `await websocket.receive_text()` will raise a `WebSocketDisconnect` exception, which you can then catch and handle like in this example.
```Python hl_lines="81-83"
{!../../../docs_src/websockets/tutorial003.py!}
```
To try it out:
* Open the app with several browser tabs.
* Write messages from them.
* Then close one of the tabs.
That will raise the `WebSocketDisconnect` exception, and all the other clients will receive a message like:
```
Client #1596980209979 left the chat
```
!!! tip
The app above is a minimal and simple example to demonstrate how to handle and broadcast messages to several WebSocket connections.
But have in mind that, as everything is handled in memory, in a single list, it will only work while the process is running, and will only work with a single process.
If you need something easy to integrate with FastAPI but that is more robust, supported by Redis, PostgreSQL or others, check <a href="https://github.com/encode/broadcaster" class="external-link" target="_blank">encode/broadcaster</a>.
## More info ## More info
To learn more about the options, check Starlette's documentation for: To learn more about the options, check Starlette's documentation for:

83
docs_src/websockets/tutorial003.py Normal file
View File

@ -0,0 +1,83 @@
from typing import List
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
app = FastAPI()
html = """
<!DOCTYPE html>
<html>
<head>
<title>Chat</title>
</head>
<body>
<h1>WebSocket Chat</h1>
<h2>Your ID: <span id="ws-id"></span></h2>
<form action="" onsubmit="sendMessage(event)">
<input type="text" id="messageText" autocomplete="off"/>
<button>Send</button>
</form>
<ul id='messages'>
</ul>
<script>
var client_id = Date.now()
document.querySelector("#ws-id").textContent = client_id;
var ws = new WebSocket(`ws://localhost:8000/ws/${client_id}`);
ws.onmessage = function(event) {
var messages = document.getElementById('messages')
var message = document.createElement('li')
var content = document.createTextNode(event.data)
message.appendChild(content)
messages.appendChild(message)
};
function sendMessage(event) {
var input = document.getElementById("messageText")
ws.send(input.value)
input.value = ''
event.preventDefault()
}
</script>
</body>
</html>
"""
class ConnectionManager:
def __init__(self):
self.active_connections: List[WebSocket] = []
async def connect(self, websocket: WebSocket):
await websocket.accept()
self.active_connections.append(websocket)
def disconnect(self, websocket: WebSocket):
self.active_connections.remove(websocket)
async def send_personal_message(self, message: str, websocket: WebSocket):
await websocket.send_text(message)
async def broadcast(self, message: str):
for connection in self.active_connections:
await connection.send_text(message)
manager = ConnectionManager()
@app.get("/")
async def get():
return HTMLResponse(html)
@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: int):
await manager.connect(websocket)
try:
while True:
data = await websocket.receive_text()
await manager.send_personal_message(f"You wrote: {data}", websocket)
await manager.broadcast(f"Client #{client_id} says: {data}")
except WebSocketDisconnect:
manager.disconnect(websocket)
await manager.broadcast(f"Client #{client_id} left the chat")

2
fastapi/__init__.py
View File

@ -22,4 +22,4 @@ from .param_functions import (
from .requests import Request from .requests import Request
from .responses import Response from .responses import Response
from .routing import APIRouter from .routing import APIRouter
from .websockets import WebSocket from .websockets import WebSocket, WebSocketDisconnect

22
tests/test_tutorial/test_websockets/test_tutorial003.py Normal file
View File

@ -0,0 +1,22 @@
from fastapi.testclient import TestClient
from docs_src.websockets.tutorial003 import app
client = TestClient(app)
def test_websocket_handle_disconnection():
with client.websocket_connect("/ws/1234") as connection, client.websocket_connect(
"/ws/5678"
) as connection_two:
connection.send_text("Hello from 1234")
data1 = connection.receive_text()
assert data1 == "You wrote: Hello from 1234"
data2 = connection_two.receive_text()
client1_says = "Client #1234 says: Hello from 1234"
assert data2 == client1_says
data1 = connection.receive_text()
assert data1 == client1_says
connection_two.close()
data1 = connection.receive_text()
assert data1 == "Client #5678 left the chat"