FastAPI 学习之路（五十八）对之前的代码进行优化

我们之前登录认证的一些内容都是直接写入到代码中，我们现在统一的放到config.py中。

        首先是token相关配置

# token
SECRET_KEY = "09d25e094faa6ca2556c818166b7a9563b93f7099f6f0f4caa6cf63b88e8d3e7"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30

        修改user.py的引入token配置相关方式为

from chapter37.config import SECRET_KEY, ALGORITHM, ACCESS_TOKEN_EXPIRE_MINUTES

        redis相关配置信息，如host，port，db等也需要移到config中一起管理

# redis
redis_host = "10.30.10.36"
redis_port = "6379"
redis_db = "0"

        修改main.py，引入redis相关配置，并修改初始化redis方法

from config import redis_host, redis_port, redis_db

async def create_redis() -> Redis:
    return await create_redis_pool(f"redis://:@{redis_host}:{redis_port}/{redis_db}?encoding=utf-8")

         调整之前websocket相关，统一放到routers目录下，创建websocket.py

"""
-*- encoding=utf-8 -*-
Time: 2024/7/18 10:57
Author: lc
Email: 15101006331@163.com
File: websocket.py
"""


from typing import List, Dict
from fastapi import APIRouter
from fastapi import WebSocket, Depends, WebSocketDisconnect

from routers.users import get_cookie_or_token


class ConnectionManager:

    def __init__(self):
        """存放链接"""
        self.active_connections: List[Dict[str, WebSocket]] = []

    async def connect(self, user: str, ws: WebSocket):
        """链接"""
        await ws.accept()
        self.active_connections.append({"user": user, "ws": ws})

    async def disconnect(self, user: str, ws: WebSocket):
        """断开链接，移除"""
        self.active_connections.remove({"user": user, "ws": ws})

    @staticmethod
    async def send_personal_message(message: str, ws: WebSocket):
        """发送所有人消息"""
        await ws.send_text(message)

    async def send_other_message(self, message: dict, user: str):
        """发送个人消息"""
        for coon in self.active_connections:
            if coon["user"] == user:
                await coon["ws"].send_json(message)

    async def broadcast(self, data: str):
        """广播"""
        for conn in self.active_connections:
            await conn["ws"].send_text(data)

    async def broadcast_json(self, data: dict):
        """广播json数据"""
        for conn in self.active_connections:
            await conn["ws"].send_json(data)


ws_manager = ConnectionManager()
websocket_router = APIRouter()


@websocket_router.websocket("/ws/{user}")
async def websocket_json_data(
        websocket: WebSocket,
        user: str,
        cookie_or_token: str = Depends(get_cookie_or_token)
):
    await ws_manager.connect(user, websocket)
    try:
        while True:
            data = await websocket.receive_json()
            send_user = data["username"]
            if send_user:
                await ws_manager.send_other_message(data, send_user)
            else:
                await ws_manager.broadcast_json(data)
    except WebSocketDisconnect as e:
        await ws_manager.disconnect(user, websocket)

        在main中引用并使用

from routers.websocket import websocket_router

app.include_router(websocket_router, prefix="/socket", tags=["websocket"])

 调整完后，启动看下接口文档（注意：如果想要websocket接口展示到文档中，需要做额外的配置）