fastAPI Tortois-ORM websocket 多个客户端实现_fastapi websocket多客户端并发

实现 FastAPI、Tortoise-ORM 和 WebSocket 多个客户端的功能，您需要完成以下步骤：

1. 安装 FastAPI、Tortoise-ORM 和 WebSocket 库 您可以使用 pip 安装这些库，例如：

   pip install fastapi
   pip install uvicorn
   pip install tortoise-orm
   pip install websockets
   

2. 实现 Tortoise-ORM 数据库模型 首先，您需要定义数据库模型，例如：

   from tortoise.models import Model
   from tortoise import fields
    
   class User(Model):
       id = fields.UUIDField(pk=True)
       name = fields.CharField(max_length=50)
       age = fields.IntField()
   

3. 实现 FastAPI 路由和 WebSocket 处理程序 您可以在 FastAPI 应用程序中定义路由和 WebSocket 处理程序，例如：

   from fastapi import FastAPI, WebSocket
   from tortoise.contrib.fastapi import register_tortoise
    
   app = FastAPI()
    
   @app.websocket("/ws/{client_id}")
   async def websocket_endpoint(websocket: WebSocket, client_id: str):
       await websocket.accept()
       while True:
           try:
               data = await websocket.receive_json()
               # 在这里处理客户端发送的数据
               # 这里可以用 Tortoise-ORM 查询数据库，然后将结果返回给客户端
           except:
               break
    
   register_tortoise(
       app,
       db_url="sqlite://db.sqlite3",
       modules={"models": ["app.models"]},
       generate_schemas=True,
   )
    
   if __name__ == "__main__":
       import uvicorn
       uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
   

4. 启动应用程序并测试 WebSocket 您可以使用浏览器中的 WebSocket 客户端插件测试 WebSocket，例如 Chrome 中的 Simple WebSocket Client 插件。在插件中输入 WebSocket URL ws://localhost:8000/ws/client-1，然后发送 JSON 数据，就可以将数据发送给服务器并接收来自服务器的响应。

5. 添加多个客户端支持 您可以为每个客户端创建一个新的 WebSocket 连接，并使用不同的客户端 ID 标识它们。然后，在服务器端，您可以使用客户端 ID 区分不同的客户端并将数据发送回正确的客户端。

下面是一个具体实现的示例代码：

from fastapi import FastAPI, WebSocket
from typing import Dict
from tortoise.contrib.fastapi import register_tortoise
from tortoise.transactions import in_transaction
import uuid
 
app = FastAPI()
clients: Dict[str, WebSocket] = {}
 
@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await websocket.accept()
    clients[client_id] = websocket
    try:
        while True:
            data = await websocket.receive_json()
            await send_message_to_all_clients(data)
    except Exception as e:
        print(f"WebSocket error: {e}")
    finally:
        del clients[client_id]
 
async def send_message_to_all_clients(data: dict):
    async with in_transaction():
        client_id = data["client_id"]
        message = data["message"]
        await save_message_to_database(client_id, message)
        for client, ws in clients.items():
            if client != client_id:
                await ws.send_json({"client_id": client_id, "message": message})
 
async def save_message_to_database(client_id: str, message: str):
    # save the message to database using Tortoise-ORM
    ...
 
register_tortoise(
    app,
    db_url="sqlite://db.sqlite3",
    modules={"models": ["app.models"]},
    generate_schemas=True,
)
 
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)

在这个示例中，我们添加了 clients 字典来存储所有连接的客户端，并使用客户端 ID 作为字典的键。在客户端连接时，我们将客户端添加到 clients 字典中。在客户端断开连接时，我们将客户端从 clients 字典中删除。每次收到一个客户端的消息时，我们将消息发送给所有连接的客户端，除了当前客户端。我们还使用 Tortoise-ORM 将消息保存到数据库中。

注意，这是一个简单的示例实现，您可以根据实际需求进行修改和完善。

下面是一个简单的封装示例：

from typing import Dict
from fastapi import FastAPI, WebSocket
from tortoise.contrib.fastapi import register_tortoise
from tortoise.transactions import in_transaction
 
app = FastAPI()
 
class WebSocketManager:
    def __init__(self):
        self.clients: Dict[str, WebSocket] = {}
 
    async def connect(self, websocket: WebSocket, client_id: str):
        await websocket.accept()
        self.clients[client_id] = websocket
 
    async def disconnect(self, client_id: str):
        del self.clients[client_id]
 
    async def send_message_to_all(self, data: dict):
        async with in_transaction():
            client_id = data["client_id"]
            message = data["message"]
            await self.save_message_to_database(client_id, message)
            for client, ws in self.clients.items():
                if client != client_id:
                    await ws.send_json({"client_id": client_id, "message": message})
 
    async def save_message_to_database(self, client_id: str, message: str):
        # save the message to database using Tortoise-ORM
        ...
 
manager = WebSocketManager()
 
@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(websocket, client_id)
    try:
        while True:
            data = await websocket.receive_json()
            await manager.send_message_to_all(data)
    except Exception as e:
        print(f"WebSocket error: {e}")
    finally:
        await manager.disconnect(client_id)
 
register_tortoise(
    app,
    db_url="sqlite://db.sqlite3",
    modules={"models": ["app.models"]},
    generate_schemas=True,
)
 
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)

在这个示例中，我们创建了一个 WebSocketManager 类来封装 WebSocket 的管理逻辑。在 WebSocketManager 类中，我们将 clients 字典和一些 WebSocket 相关的方法封装在一起。在 websocket_endpoint 函数中，我们只需要调用 WebSocketManager 的相应方法就可以处理 WebSocket 连接和消息发送的逻辑。

这个示例是一个简单的封装示例，请根据实际需求进行修改和完善。