FastAPI WebSockets
What are WebSockets in FastAPI?
WebSockets in FastAPI provide a way to establish a two-way, full-duplex communication channel between the client and the server over a single TCP connection. Unlike HTTP, WebSockets allow for real-time, bi-directional data transfer without needing to make a new request for each message. This is useful for use cases such as real-time chat applications, live updates, and notifications.
How do you implement WebSockets in FastAPI?
FastAPI has built-in support for WebSockets using the WebSocket class. To implement WebSockets, you define an endpoint using the @app.websocket() decorator and handle connections by sending and receiving messages over the WebSocket connection.
Example of a basic WebSocket implementation:
from fastapi import FastAPI, WebSocket
app = FastAPI()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
while True:
data = await websocket.receive_text()
await websocket.send_text(f"Message received: {data}")
In this example, the /ws endpoint establishes a WebSocket connection. The server continuously listens for messages from the client, processes them, and sends back a response.
How do you establish a WebSocket connection from the client?
To establish a WebSocket connection from the client, you use the JavaScript WebSocket API. The client connects to the WebSocket endpoint and can then send and receive messages in real-time.
Example of establishing a WebSocket connection from a web browser:
const ws = new WebSocket("ws://localhost:8000/ws");
ws.onopen = () => {
console.log("WebSocket connection opened");
ws.send("Hello Server");
};
ws.onmessage = (event) => {
console.log("Message from server:", event.data);
};
ws.onclose = () => {
console.log("WebSocket connection closed");
};
In this example, the client connects to the WebSocket server at ws://localhost:8000/ws, sends a message, and listens for messages from the server.
How do you send and receive messages over WebSockets in FastAPI?
In FastAPI, you can send messages to the client using websocket.send_text() for text messages and websocket.send_bytes() for binary messages. Similarly, you can receive messages from the client using websocket.receive_text() or websocket.receive_bytes().
Example of sending and receiving messages:
from fastapi import FastAPI, WebSocket
app = FastAPI()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
while True:
data = await websocket.receive_text()
await websocket.send_text(f"Message received: {data}")
In this example, the server waits for a message from the client, processes it, and then sends back a confirmation message.
How do you broadcast messages to multiple WebSocket clients in FastAPI?
To broadcast messages to multiple WebSocket clients in FastAPI, you need to manage the connected clients by storing them in a list or set. Whenever a message is received, you iterate through the connected clients and send the message to each of them.
Example of broadcasting messages to multiple clients:
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
app = FastAPI()
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 broadcast(self, message: str):
for connection in self.active_connections:
await connection.send_text(message)
manager = ConnectionManager()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await manager.connect(websocket)
try:
while True:
data = await websocket.receive_text()
await manager.broadcast(f"Client says: {data}")
except WebSocketDisconnect:
manager.disconnect(websocket)
In this example, the ConnectionManager class manages multiple WebSocket connections. When a client sends a message, it is broadcast to all connected clients.
How do you handle WebSocket disconnections in FastAPI?
To handle WebSocket disconnections in FastAPI, you can use the WebSocketDisconnect exception. This allows you to detect when a client disconnects and perform any necessary cleanup, such as removing the client from the list of active connections.
Example of handling disconnections:
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
app = FastAPI()
connected_clients = []
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
connected_clients.append(websocket)
try:
while True:
data = await websocket.receive_text()
for client in connected_clients:
await client.send_text(f"Broadcast: {data}")
except WebSocketDisconnect:
connected_clients.remove(websocket)
await websocket.close()
In this example, when a client disconnects, it is removed from the connected_clients list and the WebSocket connection is closed properly.
How do you secure WebSocket connections in FastAPI?
WebSocket connections can be secured in FastAPI by adding authentication or authorization layers, similar to how you secure HTTP endpoints. You can use query parameters, headers, or tokens to authenticate clients before accepting a WebSocket connection.
Example of securing a WebSocket connection with query parameters:
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends
app = FastAPI()
async def get_user(token: str):
# Perform token validation (pseudo code)
if token != "valid-token":
raise WebSocketDisconnect()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, token: str = Depends(get_user)):
await websocket.accept()
try:
while True:
data = await websocket.receive_text()
await websocket.send_text(f"Message received: {data}")
except WebSocketDisconnect:
await websocket.close()
In this example, the client must provide a valid token in the query parameters to establish a WebSocket connection. If the token is invalid, the connection is closed.
How do you send binary data over WebSockets in FastAPI?
FastAPI supports sending and receiving binary data over WebSockets using the websocket.send_bytes() and websocket.receive_bytes() methods. This is useful for transmitting files, images, or other binary data.
Example of sending and receiving binary data:
from fastapi import FastAPI, WebSocket
app = FastAPI()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
while True:
data = await websocket.receive_bytes()
await websocket.send_bytes(data) # Echo the binary data back to the client
In this example, binary data is received from the client and echoed back using the receive_bytes() and send_bytes() methods.
How do you close WebSocket connections in FastAPI?
WebSocket connections in FastAPI can be closed by calling the websocket.close() method. You can close a connection for various reasons, such as a client sending an invalid message, authentication failure, or when the session ends.
Example of closing a WebSocket connection:
from fastapi import FastAPI, WebSocket
app = FastAPI()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
try:
while True:
data = await websocket.receive_text()
if data == "close":
await websocket.close()
break
await websocket.send_text(f"Received: {data}")
except Exception as e:
await websocket.close(code=1000, reason=str(e))
In this example, the WebSocket connection is closed if the client sends a message with the text close. Additionally, the connection is closed with a custom reason if an error occurs.
How do you handle WebSocket events in FastAPI?
WebSockets in FastAPI support various events such as connection opening, receiving data, and connection closing. You can use these events to handle different stages of the WebSocket lifecycle and perform actions like logging or notifying clients when someone connects or disconnects.
Example of handling WebSocket events:
from fastapi import FastAPI, WebSocket
app = FastAPI()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
print("Client connected")
try:
while True:
data = await websocket.receive_text()
await websocket.send_text(f"Message received: {data}")
except Exception as e:
print(f"Error occurred: {e}")
finally:
print("Client disconnected")
await websocket.close()
In this example, WebSocket events such as client connection, message reception, and disconnection are logged.