目录
- 使用FastAPI和WebSocket构建高性能实时聊天系统
-
- [1. 系统概述与架构设计](#1. 系统概述与架构设计)
-
- [1.1 实时聊天系统的核心需求](#1.1 实时聊天系统的核心需求)
- [1.2 技术架构选型](#1.2 技术架构选型)
- [1.3 系统架构详解](#1.3 系统架构详解)
- [2. 环境搭建与项目初始化](#2. 环境搭建与项目初始化)
-
- [2.1 项目结构设计](#2.1 项目结构设计)
- [2.2 依赖管理配置](#2.2 依赖管理配置)
- [3. 数据库设计与模型](#3. 数据库设计与模型)
-
- [3.1 PostgreSQL数据库设计](#3.1 PostgreSQL数据库设计)
- [3.2 Redis数据结构设计](#3.2 Redis数据结构设计)
- [4. WebSocket实时通信实现](#4. WebSocket实时通信实现)
-
- [4.1 WebSocket连接管理器](#4.1 WebSocket连接管理器)
- [4.2 WebSocket路由端点](#4.2 WebSocket路由端点)
- [5. 业务服务层实现](#5. 业务服务层实现)
-
- [5.1 聊天服务实现](#5.1 聊天服务实现)
- [5.2 用户服务实现](#5.2 用户服务实现)
- [6. API路由和端点](#6. API路由和端点)
-
- [6.1 认证路由](#6.1 认证路由)
- [6.2 聊天路由](#6.2 聊天路由)
- [7. 主应用程序配置](#7. 主应用程序配置)
- [8. Docker部署配置](#8. Docker部署配置)
- [9. 性能优化与扩展](#9. 性能优化与扩展)
-
- [9.1 性能优化策略](#9.1 性能优化策略)
- [9.2 水平扩展策略](#9.2 水平扩展策略)
- [10. 完整项目总结](#10. 完整项目总结)
-
- [10.1 关键技术亮点](#10.1 关键技术亮点)
- [10.2 性能指标](#10.2 性能指标)
- [10.3 部署建议](#10.3 部署建议)
- [10.4 未来扩展方向](#10.4 未来扩展方向)
『宝藏代码胶囊开张啦！』------ 我的 CodeCapsule 来咯！✨写代码不再头疼！我的新站点 CodeCapsule 主打一个 "白菜价"+"量身定制 "！无论是卡脖子的毕设/课设/文献复现 ，需要灵光一现的算法改进 ，还是想给项目加个"外挂"，这里都有便宜又好用的代码方案等你发现！低成本，高适配，助你轻松通关！速来围观 👉 CodeCapsule官网
使用FastAPI和WebSocket构建高性能实时聊天系统
1. 系统概述与架构设计
1.1 实时聊天系统的核心需求
实时聊天系统在现代应用开发中扮演着至关重要的角色。一个完整的实时聊天系统需要满足以下核心需求：
- 即时消息传递：毫秒级消息延迟
- 高并发连接：支持大量用户同时在线
- 消息持久化：消息历史记录存储
- 用户状态管理：在线/离线状态跟踪
- 消息通知：离线消息推送
- 安全通信：端到端加密
- 多端同步：Web、移动端同步
1.2 技术架构选型
基础设施
数据层
业务层
网关层
客户端层
Web前端
移动端
桌面端
Nginx负载均衡
WebSocket连接管理
FastAPI服务器集群
消息分发服务
用户状态服务
Redis缓存
PostgreSQL
MongoDB
Docker容器
Kubernetes集群
监控告警系统
1.3 系统架构详解
| 层级 | 技术栈 | 作用 |
|---|---|---|
| 客户端 | Vue.js/React, Flutter, React Native | 用户界面交互 |
| 网关层 | Nginx, WebSocket | 连接管理，负载均衡 |
| 业务层 | FastAPI, WebSocket, Redis Pub/Sub | 实时消息处理，业务逻辑 |
| 数据层 | PostgreSQL, Redis, MongoDB | 数据存储，缓存，会话 |
| 基础设施 | Docker, Kubernetes, Prometheus | 容器化，编排，监控 |
2. 环境搭建与项目初始化
2.1 项目结构设计
real-time-chat-system/
├── app/
│ ├── __init__.py
│ ├── main.py
│ ├── config.py
│ ├── database/
│ │ ├── __init__.py
│ │ ├── models.py
│ │ ├── redis_client.py
│ │ └── postgres.py
│ ├── api/
│ │ ├── __init__.py
│ │ ├── endpoints/
│ │ │ ├── __init__.py
│ │ │ ├── auth.py
│ │ │ ├── users.py
│ │ │ └── chat.py
│ │ └── dependencies.py
│ ├── core/
│ │ ├── __init__.py
│ │ ├── security.py
│ │ ├── websocket.py
│ │ └── cache.py
│ ├── models/
│ │ ├── __init__.py
│ │ ├── schemas.py
│ │ └── enums.py
│ ├── services/
│ │ ├── __init__.py
│ │ ├── auth_service.py
│ │ ├── chat_service.py
│ │ ├── user_service.py
│ │ └── notification_service.py
│ └── utils/
│ ├── __init__.py
│ ├── validators.py
│ └── helpers.py
├── tests/
│ ├── __init__.py
│ ├── test_auth.py
│ ├── test_chat.py
│ └── test_websocket.py
├── alembic/
│ ├── versions/
│ └── alembic.ini
├── docker/
│ ├── Dockerfile
│ ├── docker-compose.yml
│ └── nginx/
│ └── nginx.conf
├── requirements/
│ ├── base.txt
│ ├── dev.txt
│ └── prod.txt
├── .env.example
├── .gitignore
├── README.md
└── pyproject.toml2.2 依赖管理配置
python
# pyproject.toml
[build-system]
requires = ["setuptools>=61.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "real-time-chat-system"
version = "1.0.0"
description = "高性能实时聊天系统"
readme = "README.md"
authors = [{name = "Chat System Team", email = "team@chatsystem.com"}]
license = {text = "MIT"}
requires-python = ">=3.9"
dependencies = [
"fastapi==0.104.1",
"uvicorn[standard]==0.24.0",
"python-jose[cryptography]==3.3.0",
"passlib[bcrypt]==1.7.4",
"python-multipart==0.0.6",
"sqlalchemy==2.0.23",
"alembic==1.12.1",
"asyncpg==0.29.0",
"redis==5.0.1",
"aioredis==2.0.1",
"pydantic==2.5.0",
"pydantic-settings==2.1.0",
"python-socketio==5.10.0",
"websockets==12.0",
"celery==5.3.4",
"httpx==0.25.1",
"email-validator==2.1.0",
"brotli==1.1.0",
]
[project.optional-dependencies]
dev = [
"pytest==7.4.3",
"pytest-asyncio==0.21.1",
"black==23.11.0",
"flake8==6.1.0",
"mypy==1.7.1",
"pre-commit==3.5.0",
"pytest-cov==4.1.0",
"locust==2.20.1",
]3. 数据库设计与模型
3.1 PostgreSQL数据库设计
python
# app/database/models.py
"""
数据库模型定义
"""
from sqlalchemy import (
Column, Integer, String, Text, DateTime, Boolean,
ForeignKey, Numeric, JSON, Index, UniqueConstraint
)
from sqlalchemy.sql import func
from sqlalchemy.orm import relationship, declarative_base
from sqlalchemy.dialects.postgresql import UUID, ARRAY
import uuid
from datetime import datetime
Base = declarative_base()
class User(Base):
"""用户表"""
__tablename__ = "users"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
username = Column(String(50), unique=True, nullable=False, index=True)
email = Column(String(255), unique=True, nullable=False, index=True)
phone = Column(String(20), unique=True, nullable=True)
hashed_password = Column(String(255), nullable=False)
# 用户信息
full_name = Column(String(100))
avatar_url = Column(String(500))
bio = Column(Text)
status = Column(String(20), default="offline") # online, offline, away, busy
# 设置
settings = Column(JSON, default=lambda: {
"notifications": True,
"sound": True,
"theme": "light",
"language": "zh-CN"
})
# 时间戳
created_at = Column(DateTime(timezone=True), server_default=func.now())
updated_at = Column(DateTime(timezone=True), onupdate=func.now())
last_seen = Column(DateTime(timezone=True))
# 关系
messages = relationship("Message", back_populates="sender")
chat_memberships = relationship("ChatMember", back_populates="user")
__table_args__ = (
Index('ix_users_username_lower', func.lower(username)),
Index('ix_users_email_lower', func.lower(email)),
)
def __repr__(self):
return f"<User(id={self.id}, username={self.username})>"
class ChatRoom(Base):
"""聊天室表"""
__tablename__ = "chat_rooms"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
name = Column(String(100))
description = Column(Text)
room_type = Column(String(20), nullable=False) # private, group, channel
avatar_url = Column(String(500))
# 群组特定字段
created_by = Column(UUID(as_uuid=True), ForeignKey("users.id"))
is_public = Column(Boolean, default=False)
max_members = Column(Integer, default=1000)
# 元数据
metadata = Column(JSON, default=lambda: {})
settings = Column(JSON, default=lambda: {
"allow_edit": True,
"allow_delete": True,
"allow_invites": True,
"require_approval": False
})
# 时间戳
created_at = Column(DateTime(timezone=True), server_default=func.now())
updated_at = Column(DateTime(timezone=True), onupdate=func.now())
# 关系
messages = relationship("Message", back_populates="chat_room")
members = relationship("ChatMember", back_populates="chat_room")
__table_args__ = (
Index('ix_chat_rooms_room_type', room_type),
Index('ix_chat_rooms_created_by', created_by),
)
class ChatMember(Base):
"""聊天室成员表"""
__tablename__ = "chat_members"
id = Column(Integer, primary_key=True, autoincrement=True)
chat_room_id = Column(UUID(as_uuid=True), ForeignKey("chat_rooms.id"), nullable=False)
user_id = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=False)
# 成员角色和权限
role = Column(String(20), default="member") # owner, admin, moderator, member
permissions = Column(JSON, default=lambda: {
"send_messages": True,
"invite_users": False,
"remove_users": False,
"edit_room": False
})
# 状态
is_muted = Column(Boolean, default=False)
is_banned = Column(Boolean, default=False)
joined_at = Column(DateTime(timezone=True), server_default=func.now())
# 阅读状态
last_read_message_id = Column(Integer, ForeignKey("messages.id"), nullable=True)
unread_count = Column(Integer, default=0)
# 关系
user = relationship("User", back_populates="chat_memberships")
chat_room = relationship("ChatRoom", back_populates="members")
__table_args__ = (
UniqueConstraint('chat_room_id', 'user_id', name='uq_chat_member'),
Index('ix_chat_members_user_id', user_id),
Index('ix_chat_members_chat_room_id', chat_room_id),
)
class Message(Base):
"""消息表"""
__tablename__ = "messages"
id = Column(Integer, primary_key=True, autoincrement=True)
uuid = Column(UUID(as_uuid=True), unique=True, default=uuid.uuid4, index=True)
chat_room_id = Column(UUID(as_uuid=True), ForeignKey("chat_rooms.id"), nullable=False)
sender_id = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=False)
# 消息内容
message_type = Column(String(20), default="text") # text, image, file, audio, video, location
content = Column(Text, nullable=False)
metadata = Column(JSON, default=lambda: {}) # 文件大小、时长、位置等
# 引用和回复
reply_to_id = Column(Integer, ForeignKey("messages.id"), nullable=True)
is_edited = Column(Boolean, default=False)
edited_at = Column(DateTime(timezone=True), nullable=True)
# 状态
status = Column(String(20), default="sent") # sent, delivered, read, failed
deleted_at = Column(DateTime(timezone=True), nullable=True)
# 时间戳
created_at = Column(DateTime(timezone=True), server_default=func.now())
updated_at = Column(DateTime(timezone=True), onupdate=func.now())
# 关系
chat_room = relationship("ChatRoom", back_populates="messages")
sender = relationship("User", back_populates="messages")
reply_to = relationship("Message", remote_side=[id])
__table_args__ = (
Index('ix_messages_chat_room_created', chat_room_id, created_at),
Index('ix_messages_sender_id', sender_id),
Index('ix_messages_uuid', uuid, unique=True),
)
class Attachment(Base):
"""附件表"""
__tablename__ = "attachments"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
message_id = Column(Integer, ForeignKey("messages.id"), nullable=False)
file_type = Column(String(50), nullable=False) # image/jpeg, application/pdf, etc.
file_name = Column(String(255), nullable=False)
file_size = Column(Integer, nullable=False) # 字节数
file_url = Column(String(500), nullable=False)
thumbnail_url = Column(String(500))
# 元数据
metadata = Column(JSON, default=lambda: {})
created_at = Column(DateTime(timezone=True), server_default=func.now())
# 关系
message = relationship("Message", foreign_keys=[message_id])
class UserSession(Base):
"""用户会话表"""
__tablename__ = "user_sessions"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
user_id = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=False)
device_id = Column(String(100))
device_type = Column(String(50)) # web, ios, android
ip_address = Column(String(45)) # 支持IPv6
# 会话信息
access_token = Column(Text, nullable=False)
refresh_token = Column(Text, nullable=False)
expires_at = Column(DateTime(timezone=True), nullable=False)
# 状态
is_active = Column(Boolean, default=True)
last_activity = Column(DateTime(timezone=True), server_default=func.now())
created_at = Column(DateTime(timezone=True), server_default=func.now())
__table_args__ = (
Index('ix_user_sessions_user_id', user_id),
Index('ix_user_sessions_access_token', access_token),
Index('ix_user_sessions_expires_at', expires_at),
)
class Notification(Base):
"""通知表"""
__tablename__ = "notifications"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
user_id = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=False)
notification_type = Column(String(50), nullable=False) # message, friend_request, system
title = Column(String(200))
content = Column(Text, nullable=False)
data = Column(JSON, default=lambda: {})
# 状态
is_read = Column(Boolean, default=False)
read_at = Column(DateTime(timezone=True))
created_at = Column(DateTime(timezone=True), server_default=func.now())
__table_args__ = (
Index('ix_notifications_user_id', user_id),
Index('ix_notifications_is_read', is_read),
Index('ix_notifications_created_at', created_at),
)3.2 Redis数据结构设计
python
# app/database/redis_client.py
"""
Redis客户端和数据结构管理
"""
import redis.asyncio as redis
import json
import pickle
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Any, Set
from enum import Enum
import logging
logger = logging.getLogger(__name__)
class RedisKeys:
"""Redis键前缀常量"""
# 用户相关
USER_SESSIONS = "user:sessions:{user_id}"
USER_ONLINE = "user:online:{user_id}"
USER_PRESENCE = "user:presence:{user_id}"
# 聊天室相关
CHAT_MEMBERS = "chat:members:{chat_id}"
CHAT_MESSAGES = "chat:messages:{chat_id}"
CHAT_UNREAD = "chat:unread:{chat_id}:{user_id}"
# WebSocket连接
WS_CONNECTION = "ws:connection:{connection_id}"
USER_CONNECTIONS = "user:connections:{user_id}"
# 限流和缓存
RATE_LIMIT = "rate:limit:{key}"
MESSAGE_CACHE = "cache:message:{message_id}"
USER_PROFILE_CACHE = "cache:user:{user_id}"
# 发布/订阅频道
CHANNEL_MESSAGE = "channel:message"
CHANNEL_PRESENCE = "channel:presence"
CHANNEL_NOTIFICATION = "channel:notification"
class RedisClient:
"""Redis客户端封装"""
def __init__(self, host: str = "localhost", port: int = 6379,
password: str = None, db: int = 0):
self.pool = None
self.host = host
self.port = port
self.password = password
self.db = db
async def initialize(self):
"""初始化连接池"""
self.pool = redis.ConnectionPool.from_url(
f"redis://{self.host}:{self.port}/{self.db}",
password=self.password,
decode_responses=True,
max_connections=20
)
self.client = redis.Redis(connection_pool=self.pool)
# 测试连接
try:
await self.client.ping()
logger.info("Redis连接成功")
except Exception as e:
logger.error(f"Redis连接失败: {e}")
raise
async def close(self):
"""关闭连接"""
if self.client:
await self.client.close()
if self.pool:
await self.pool.disconnect()
# ========== 用户在线状态管理 ==========
async def set_user_online(self, user_id: str, connection_id: str,
device_info: Dict = None) -> bool:
"""设置用户在线状态"""
try:
online_key = RedisKeys.USER_ONLINE.format(user_id=user_id)
presence_key = RedisKeys.USER_PRESENCE.format(user_id=user_id)
connection_key = RedisKeys.WS_CONNECTION.format(connection_id=connection_id)
# 存储连接信息
connection_data = {
"user_id": user_id,
"device_info": device_info or {},
"connected_at": datetime.now().isoformat(),
"last_activity": datetime.now().isoformat()
}
# 设置在线状态（24小时过期）
await self.client.setex(online_key, 86400, "1")
# 存储在线状态详情
presence_data = {
"status": "online",
"last_seen": datetime.now().isoformat(),
"device": device_info.get("type", "unknown") if device_info else "unknown"
}
await self.client.setex(presence_key, 86400, json.dumps(presence_data))
# 存储连接信息
await self.client.setex(connection_key, 86400, json.dumps(connection_data))
# 添加到用户连接集合
user_connections_key = RedisKeys.USER_CONNECTIONS.format(user_id=user_id)
await self.client.sadd(user_connections_key, connection_id)
await self.client.expire(user_connections_key, 86400)
return True
except Exception as e:
logger.error(f"设置用户在线状态失败: {e}")
return False
async def set_user_offline(self, user_id: str, connection_id: str) -> bool:
"""设置用户离线状态"""
try:
# 从用户连接集合中移除
user_connections_key = RedisKeys.USER_CONNECTIONS.format(user_id=user_id)
await self.client.srem(user_connections_key, connection_id)
# 删除连接信息
connection_key = RedisKeys.WS_CONNECTION.format(connection_id=connection_id)
await self.client.delete(connection_key)
# 检查是否还有其他连接
remaining_connections = await self.client.scard(user_connections_key)
if remaining_connections == 0:
# 没有其他连接，设置离线状态
online_key = RedisKeys.USER_ONLINE.format(user_id=user_id)
presence_key = RedisKeys.USER_PRESENCE.format(user_id=user_id)
presence_data = {
"status": "offline",
"last_seen": datetime.now().isoformat()
}
await self.client.delete(online_key)
await self.client.setex(presence_key, 86400, json.dumps(presence_data))
return True
except Exception as e:
logger.error(f"设置用户离线状态失败: {e}")
return False
async def get_user_presence(self, user_id: str) -> Dict:
"""获取用户在线状态"""
try:
presence_key = RedisKeys.USER_PRESENCE.format(user_id=user_id)
data = await self.client.get(presence_key)
if data:
return json.loads(data)
# 如果Redis中没有，返回默认离线状态
return {
"status": "offline",
"last_seen": None
}
except Exception as e:
logger.error(f"获取用户在线状态失败: {e}")
return {"status": "unknown"}
async def update_user_activity(self, user_id: str, connection_id: str) -> bool:
"""更新用户活动时间"""
try:
connection_key = RedisKeys.WS_CONNECTION.format(connection_id=connection_id)
data = await self.client.get(connection_key)
if data:
connection_data = json.loads(data)
connection_data["last_activity"] = datetime.now().isoformat()
await self.client.setex(connection_key, 86400, json.dumps(connection_data))
return True
except Exception as e:
logger.error(f"更新用户活动时间失败: {e}")
return False
# ========== 聊天室成员管理 ==========
async def add_chat_member(self, chat_id: str, user_id: str,
role: str = "member") -> bool:
"""添加聊天室成员到Redis"""
try:
chat_members_key = RedisKeys.CHAT_MEMBERS.format(chat_id=chat_id)
member_data = {
"user_id": user_id,
"role": role,
"joined_at": datetime.now().isoformat()
}
await self.client.hset(
chat_members_key,
user_id,
json.dumps(member_data)
)
await self.client.expire(chat_members_key, 3600) # 1小时过期
return True
except Exception as e:
logger.error(f"添加聊天室成员失败: {e}")
return False
async def remove_chat_member(self, chat_id: str, user_id: str) -> bool:
"""从Redis移除聊天室成员"""
try:
chat_members_key = RedisKeys.CHAT_MEMBERS.format(chat_id=chat_id)
await self.client.hdel(chat_members_key, user_id)
return True
except Exception as e:
logger.error(f"移除聊天室成员失败: {e}")
return False
async def get_chat_members(self, chat_id: str) -> List[Dict]:
"""获取聊天室成员列表"""
try:
chat_members_key = RedisKeys.CHAT_MEMBERS.format(chat_id=chat_id)
members_data = await self.client.hgetall(chat_members_key)
members = []
for user_id, data in members_data.items():
member_info = json.loads(data)
member_info["user_id"] = user_id
members.append(member_info)
return members
except Exception as e:
logger.error(f"获取聊天室成员失败: {e}")
return []
# ========== 消息缓存 ==========
async def cache_message(self, message: Dict, ttl: int = 3600) -> bool:
"""缓存消息到Redis"""
try:
message_id = message.get("id") or message.get("uuid")
if not message_id:
return False
cache_key = RedisKeys.MESSAGE_CACHE.format(message_id=str(message_id))
await self.client.setex(cache_key, ttl, json.dumps(message))
return True
except Exception as e:
logger.error(f"缓存消息失败: {e}")
return False
async def get_cached_message(self, message_id: str) -> Optional[Dict]:
"""从缓存获取消息"""
try:
cache_key = RedisKeys.MESSAGE_CACHE.format(message_id=message_id)
data = await self.client.get(cache_key)
if data:
return json.loads(data)
return None
except Exception as e:
logger.error(f"获取缓存消息失败: {e}")
return None
# ========== 未读消息计数 ==========
async def increment_unread_count(self, chat_id: str, user_id: str,
count: int = 1) -> int:
"""增加未读消息计数"""
try:
unread_key = RedisKeys.CHAT_UNREAD.format(
chat_id=chat_id,
user_id=user_id
)
# 增加计数
new_count = await self.client.incrby(unread_key, count)
await self.client.expire(unread_key, 604800) # 7天过期
return new_count
except Exception as e:
logger.error(f"增加未读消息计数失败: {e}")
return 0
async def reset_unread_count(self, chat_id: str, user_id: str) -> bool:
"""重置未读消息计数"""
try:
unread_key = RedisKeys.CHAT_UNREAD.format(
chat_id=chat_id,
user_id=user_id
)
await self.client.delete(unread_key)
return True
except Exception as e:
logger.error(f"重置未读消息计数失败: {e}")
return False
async def get_unread_count(self, chat_id: str, user_id: str) -> int:
"""获取未读消息计数"""
try:
unread_key = RedisKeys.CHAT_UNREAD.format(
chat_id=chat_id,
user_id=user_id
)
count = await self.client.get(unread_key)
return int(count) if count else 0
except Exception as e:
logger.error(f"获取未读消息计数失败: {e}")
return 0
# ========== 发布/订阅 ==========
async def publish_message(self, channel: str, message: Dict) -> int:
"""发布消息到频道"""
try:
return await self.client.publish(
channel,
json.dumps(message)
)
except Exception as e:
logger.error(f"发布消息失败: {e}")
return 0
async def create_pubsub(self):
"""创建发布/订阅对象"""
return self.client.pubsub()4. WebSocket实时通信实现
4.1 WebSocket连接管理器
python
# app/core/websocket.py
"""
WebSocket连接管理和消息处理
"""
import asyncio
import json
import uuid
from typing import Dict, List, Set, Optional, Any
from datetime import datetime
from collections import defaultdict
import logging
from fastapi import WebSocket, WebSocketDisconnect, status
from starlette.websockets import WebSocketState
from app.database.redis_client import RedisClient
from app.services.chat_service import ChatService
from app.services.user_service import UserService
logger = logging.getLogger(__name__)
class ConnectionManager:
"""WebSocket连接管理器"""
def __init__(self, redis_client: RedisClient,
chat_service: ChatService,
user_service: UserService):
self.redis_client = redis_client
self.chat_service = chat_service
self.user_service = user_service
# 内存中的连接映射
self.active_connections: Dict[str, WebSocket] = {}
self.user_connections: Dict[str, Set[str]] = defaultdict(set)
self.connection_user_map: Dict[str, str] = {}
async def connect(self, websocket: WebSocket, user_id: str,
token: str = None) -> str:
"""建立WebSocket连接"""
# 生成连接ID
connection_id = str(uuid.uuid4())
# 接受连接
await websocket.accept()
# 存储连接
self.active_connections[connection_id] = websocket
self.user_connections[user_id].add(connection_id)
self.connection_user_map[connection_id] = user_id
# 设置用户在线状态
device_info = {
"type": "websocket",
"user_agent": websocket.headers.get("user-agent", ""),
"ip": websocket.client.host if websocket.client else "unknown"
}
await self.redis_client.set_user_online(
user_id, connection_id, device_info
)
# 发送连接成功消息
await self.send_personal_message(
connection_id,
{
"type": "connection_established",
"connection_id": connection_id,
"timestamp": datetime.now().isoformat()
}
)
# 广播用户上线通知
await self.broadcast_user_presence(user_id, "online")
logger.info(f"用户 {user_id} 建立WebSocket连接: {connection_id}")
return connection_id
async def disconnect(self, connection_id: str):
"""断开WebSocket连接"""
if connection_id not in self.active_connections:
return
user_id = self.connection_user_map.get(connection_id)
# 清理内存中的连接
websocket = self.active_connections.pop(connection_id, None)
if websocket and websocket.client_state != WebSocketState.DISCONNECTED:
try:
await websocket.close(code=status.WS_1000_NORMAL_CLOSURE)
except Exception:
pass
# 清理用户连接映射
if user_id:
self.user_connections[user_id].discard(connection_id)
if not self.user_connections[user_id]:
del self.user_connections[user_id]
self.connection_user_map.pop(connection_id, None)
# 更新Redis中的在线状态
if user_id:
await self.redis_client.set_user_offline(user_id, connection_id)
# 检查用户是否完全离线
remaining_user_connections = self.user_connections.get(user_id)
if not remaining_user_connections:
await self.broadcast_user_presence(user_id, "offline")
logger.info(f"WebSocket连接断开: {connection_id}")
async def send_personal_message(self, connection_id: str, message: Dict):
"""向特定连接发送消息"""
if connection_id not in self.active_connections:
return
websocket = self.active_connections[connection_id]
try:
await websocket.send_json(message)
except Exception as e:
logger.error(f"发送个人消息失败: {e}")
await self.disconnect(connection_id)
async def send_to_user(self, user_id: str, message: Dict):
"""向特定用户的所有连接发送消息"""
connection_ids = self.user_connections.get(user_id, set()).copy()
for connection_id in connection_ids:
await self.send_personal_message(connection_id, message)
async def broadcast_to_chat(self, chat_id: str, message: Dict,
exclude_user_ids: Set[str] = None):
"""向聊天室所有成员广播消息"""
try:
# 从Redis获取聊天室成员
members = await self.redis_client.get_chat_members(chat_id)
if not members:
# 如果Redis中没有，从数据库获取
members = await self.chat_service.get_chat_members_from_db(chat_id)
# 缓存到Redis
for member in members:
await self.redis_client.add_chat_member(
chat_id, member["user_id"], member.get("role", "member")
)
# 发送给每个在线成员
for member in members:
user_id = member["user_id"]
# 排除指定用户
if exclude_user_ids and user_id in exclude_user_ids:
continue
# 检查用户是否在线
user_connections = self.user_connections.get(user_id)
if user_connections:
# 构建消息
chat_message = message.copy()
chat_message["chat_id"] = chat_id
# 发送给用户的所有连接
await self.send_to_user(user_id, chat_message)
except Exception as e:
logger.error(f"广播消息到聊天室失败: {e}")
async def broadcast_user_presence(self, user_id: str, status: str):
"""广播用户在线状态变更"""
try:
# 获取用户信息
user = await self.user_service.get_user_by_id(user_id)
if not user:
return
# 构建状态消息
presence_message = {
"type": "presence_update",
"user_id": user_id,
"username": user.username,
"status": status,
"timestamp": datetime.now().isoformat()
}
# 获取用户的所有聊天室
user_chats = await self.chat_service.get_user_chats(user_id)
# 向每个聊天室广播状态
for chat in user_chats:
await self.broadcast_to_chat(
chat["id"],
presence_message,
exclude_user_ids={user_id} # 不发送给本人
)
except Exception as e:
logger.error(f"广播用户状态失败: {e}")
async def handle_message(self, connection_id: str, message: Dict):
"""处理收到的消息"""
try:
message_type = message.get("type")
user_id = self.connection_user_map.get(connection_id)
if not user_id:
logger.warning(f"未知连接的消息: {connection_id}")
return
# 更新用户活动时间
await self.redis_client.update_user_activity(user_id, connection_id)
# 根据消息类型处理
if message_type == "ping":
# 心跳响应
await self.send_personal_message(
connection_id,
{
"type": "pong",
"timestamp": datetime.now().isoformat()
}
)
elif message_type == "typing":
# 正在输入状态
chat_id = message.get("chat_id")
if chat_id:
typing_message = {
"type": "user_typing",
"user_id": user_id,
"chat_id": chat_id,
"timestamp": datetime.now().isoformat()
}
await self.broadcast_to_chat(
chat_id,
typing_message,
exclude_user_ids={user_id}
)
elif message_type == "chat_message":
# 聊天消息
await self._handle_chat_message(user_id, message)
elif message_type == "read_receipt":
# 已读回执
await self._handle_read_receipt(user_id, message)
elif message_type == "call_signal":
# 通话信号
await self._handle_call_signal(user_id, message)
else:
logger.warning(f"未知消息类型: {message_type}")
except Exception as e:
logger.error(f"处理消息失败: {e}")
async def _handle_chat_message(self, user_id: str, message: Dict):
"""处理聊天消息"""
try:
chat_id = message.get("chat_id")
content = message.get("content")
message_type = message.get("message_type", "text")
reply_to = message.get("reply_to")
if not chat_id or not content:
logger.warning(f"无效的聊天消息: {message}")
return
# 检查用户是否在聊天室中
is_member = await self.chat_service.is_user_in_chat(user_id, chat_id)
if not is_member:
await self.send_to_user(user_id, {
"type": "error",
"code": "NOT_MEMBER",
"message": "您不是该聊天室的成员"
})
return
# 创建消息记录
new_message = await self.chat_service.create_message(
user_id=user_id,
chat_id=chat_id,
content=content,
message_type=message_type,
reply_to=reply_to,
metadata=message.get("metadata", {})
)
# 构建广播消息
broadcast_message = {
"type": "new_message",
"message": new_message,
"timestamp": datetime.now().isoformat()
}
# 广播消息给聊天室所有成员（除了发送者）
await self.broadcast_to_chat(
chat_id,
broadcast_message,
exclude_user_ids={user_id}
)
# 增加未读计数（除了发送者）
members = await self.redis_client.get_chat_members(chat_id)
for member in members:
member_id = member["user_id"]
if member_id != user_id:
await self.redis_client.increment_unread_count(
chat_id, member_id
)
# 发送成功回执给发送者
await self.send_to_user(user_id, {
"type": "message_sent",
"message_id": new_message["id"],
"timestamp": datetime.now().isoformat()
})
except Exception as e:
logger.error(f"处理聊天消息失败: {e}")
# 发送错误回执给发送者
await self.send_to_user(user_id, {
"type": "error",
"code": "MESSAGE_FAILED",
"message": "消息发送失败"
})
async def _handle_read_receipt(self, user_id: str, message: Dict):
"""处理已读回执"""
try:
chat_id = message.get("chat_id")
message_id = message.get("message_id")
if not chat_id or not message_id:
return
# 更新消息状态为已读
await self.chat_service.mark_message_as_read(
message_id, user_id
)
# 重置未读计数
await self.redis_client.reset_unread_count(chat_id, user_id)
# 广播已读回执
receipt_message = {
"type": "message_read",
"user_id": user_id,
"chat_id": chat_id,
"message_id": message_id,
"timestamp": datetime.now().isoformat()
}
await self.broadcast_to_chat(
chat_id,
receipt_message,
exclude_user_ids={user_id}
)
except Exception as e:
logger.error(f"处理已读回执失败: {e}")
async def _handle_call_signal(self, user_id: str, message: Dict):
"""处理通话信号"""
try:
call_type = message.get("call_type") # offer, answer, candidate, hangup
target_user_id = message.get("target_user_id")
signal_data = message.get("signal_data")
if not target_user_id or not signal_data:
return
# 转发信号给目标用户
signal_message = {
"type": "call_signal",
"call_type": call_type,
"from_user_id": user_id,
"signal_data": signal_data,
"timestamp": datetime.now().isoformat()
}
await self.send_to_user(target_user_id, signal_message)
except Exception as e:
logger.error(f"处理通话信号失败: {e}")
async def cleanup_inactive_connections(self, timeout_seconds: int = 300):
"""清理不活跃的连接"""
try:
current_time = datetime.now()
connections_to_remove = []
for connection_id, websocket in self.active_connections.items():
user_id = self.connection_user_map.get(connection_id)
if not user_id:
continue
# 从Redis获取最后活动时间
connection_key = f"ws:connection:{connection_id}"
data = await self.redis_client.client.get(connection_key)
if data:
connection_data = json.loads(data)
last_activity_str = connection_data.get("last_activity")
if last_activity_str:
last_activity = datetime.fromisoformat(last_activity_str)
inactive_duration = (current_time - last_activity).total_seconds()
if inactive_duration > timeout_seconds:
connections_to_remove.append(connection_id)
# 断开不活跃的连接
for connection_id in connections_to_remove:
logger.info(f"断开不活跃连接: {connection_id}")
await self.disconnect(connection_id)
except Exception as e:
logger.error(f"清理不活跃连接失败: {e}")4.2 WebSocket路由端点
python
# app/api/endpoints/websocket.py
"""
WebSocket路由端点
"""
import json
import asyncio
from datetime import datetime
from typing import Optional, Dict, Any
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query
from fastapi.responses import HTMLResponse
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.websocket import ConnectionManager
from app.core.security import get_current_user_ws
from app.database.postgres import get_db
from app.database.redis_client import get_redis_client
from app.services.chat_service import get_chat_service
from app.services.user_service import get_user_service
import logging
logger = logging.getLogger(__name__)
router = APIRouter()
# 全局连接管理器实例
_connection_manager: Optional[ConnectionManager] = None
async def get_connection_manager(
redis_client = Depends(get_redis_client),
chat_service = Depends(get_chat_service),
user_service = Depends(get_user_service)
) -> ConnectionManager:
"""获取连接管理器实例"""
global _connection_manager
if _connection_manager is None:
_connection_manager = ConnectionManager(
redis_client=redis_client,
chat_service=chat_service,
user_service=user_service
)
return _connection_manager
@router.websocket("/ws")
async def websocket_endpoint(
websocket: WebSocket,
token: Optional[str] = Query(None),
manager: ConnectionManager = Depends(get_connection_manager)
):
"""WebSocket主连接端点"""
connection_id = None
try:
# 验证用户
if not token:
await websocket.close(code=4001, reason="缺少认证令牌")
return
user = await get_current_user_ws(token)
if not user:
await websocket.close(code=4003, reason="认证失败")
return
# 建立连接
connection_id = await manager.connect(websocket, str(user.id), token)
# 发送初始数据
await manager.send_personal_message(connection_id, {
"type": "initial_data",
"user_id": str(user.id),
"username": user.username,
"timestamp": datetime.now().isoformat()
})
# 消息处理循环
while True:
try:
# 接收消息
data = await websocket.receive_json(timeout=300)
if data:
# 处理消息
await manager.handle_message(connection_id, data)
except asyncio.TimeoutError:
# 发送心跳
await manager.send_personal_message(connection_id, {
"type": "ping",
"timestamp": datetime.now().isoformat()
})
except WebSocketDisconnect:
break
except WebSocketDisconnect as e:
logger.info(f"WebSocket正常断开: {connection_id}, 代码: {e.code}")
except Exception as e:
logger.error(f"WebSocket连接异常: {e}")
finally:
# 清理连接
if connection_id:
await manager.disconnect(connection_id)
@router.websocket("/ws/chat/{chat_id}")
async def chat_websocket_endpoint(
websocket: WebSocket,
chat_id: str,
token: Optional[str] = Query(None),
manager: ConnectionManager = Depends(get_connection_manager),
db: AsyncSession = Depends(get_db)
):
"""聊天室专用WebSocket端点"""
connection_id = None
try:
# 验证用户
if not token:
await websocket.close(code=4001, reason="缺少认证令牌")
return
user = await get_current_user_ws(token)
if not user:
await websocket.close(code=4003, reason="认证失败")
return
# 检查用户是否在聊天室中
from app.services.chat_service import ChatService
chat_service = ChatService(db)
is_member = await chat_service.is_user_in_chat(str(user.id), chat_id)
if not is_member:
await websocket.close(code=4004, reason="无访问权限")
return
# 建立连接
connection_id = await manager.connect(websocket, str(user.id), token)
# 发送聊天室信息
chat_info = await chat_service.get_chat_room(chat_id)
await manager.send_personal_message(connection_id, {
"type": "chat_joined",
"chat_id": chat_id,
"chat_info": chat_info,
"timestamp": datetime.now().isoformat()
})
# 发送历史消息
messages = await chat_service.get_chat_messages(
chat_id,
limit=50,
user_id=str(user.id)
)
await manager.send_personal_message(connection_id, {
"type": "chat_history",
"chat_id": chat_id,
"messages": messages,
"timestamp": datetime.now().isoformat()
})
# 消息处理循环
while True:
try:
data = await websocket.receive_json(timeout=300)
if data:
# 确保消息属于当前聊天室
if data.get("chat_id") != chat_id:
data["chat_id"] = chat_id
await manager.handle_message(connection_id, data)
except asyncio.TimeoutError:
# 心跳
await manager.send_personal_message(connection_id, {
"type": "ping",
"chat_id": chat_id,
"timestamp": datetime.now().isoformat()
})
except WebSocketDisconnect:
break
except WebSocketDisconnect as e:
logger.info(f"聊天WebSocket断开: {connection_id}")
except Exception as e:
logger.error(f"聊天WebSocket异常: {e}")
finally:
if connection_id:
await manager.disconnect(connection_id)
@router.get("/ws/test")
async def websocket_test_page():
"""WebSocket测试页面"""
html = """
<!DOCTYPE html>
<html>
<head>
<title>WebSocket测试</title>
<style>
body {
font-family: Arial, sans-serif;
max-width: 800px;
margin: 0 auto;
padding: 20px;
}
#messages {
border: 1px solid #ddd;
height: 400px;
overflow-y: auto;
padding: 10px;
margin-bottom: 10px;
}
#messageInput {
width: 70%;
padding: 10px;
}
button {
padding: 10px 20px;
background: #007bff;
color: white;
border: none;
cursor: pointer;
}
.message {
margin: 5px 0;
padding: 8px;
border-radius: 4px;
}
.sent {
background: #e3f2fd;
text-align: right;
}
.received {
background: #f5f5f5;
}
</style>
</head>
<body>
<h1>WebSocket聊天测试</h1>
<div>
<input type="text" id="tokenInput" placeholder="输入JWT令牌" style="width: 80%; margin-bottom: 10px;">
<button onclick="connect()">连接</button>
<button onclick="disconnect()">断开</button>
</div>
<div>
<input type="text" id="chatIdInput" placeholder="聊天室ID" style="width: 80%; margin-bottom: 10px;">
<button onclick="joinChat()">加入聊天室</button>
</div>
<div id="messages"></div>
<div>
<input type="text" id="messageInput" placeholder="输入消息...">
<button onclick="sendMessage()">发送</button>
</div>
<script>
let ws = null;
let currentChatId = null;
function connect() {
const token = document.getElementById('tokenInput').value;
if (!token) {
alert('请输入令牌');
return;
}
// 关闭现有连接
if (ws) {
ws.close();
}
ws = new WebSocket(`ws://localhost:8000/ws?token=${token}`);
ws.onopen = function() {
addMessage('系统', '连接已建立', 'received');
};
ws.onmessage = function(event) {
const data = JSON.parse(event.data);
handleWebSocketMessage(data);
};
ws.onclose = function(event) {
addMessage('系统', `连接已断开: ${event.reason || event.code}`, 'received');
ws = null;
};
ws.onerror = function(error) {
addMessage('系统', `连接错误: ${error}`, 'received');
};
}
function disconnect() {
if (ws) {
ws.close();
}
}
function joinChat() {
const chatId = document.getElementById('chatIdInput').value;
if (!chatId) {
alert('请输入聊天室ID');
return;
}
const token = document.getElementById('tokenInput').value;
if (!token) {
alert('请先连接');
return;
}
// 关闭现有连接
if (ws) {
ws.close();
}
currentChatId = chatId;
ws = new WebSocket(`ws://localhost:8000/ws/chat/${chatId}?token=${token}`);
ws.onopen = function() {
addMessage('系统', `已加入聊天室: ${chatId}`, 'received');
};
ws.onmessage = function(event) {
const data = JSON.parse(event.data);
handleWebSocketMessage(data);
};
ws.onclose = function(event) {
addMessage('系统', `已离开聊天室: ${event.reason || event.code}`, 'received');
ws = null;
currentChatId = null;
};
}
function sendMessage() {
if (!ws || ws.readyState !== WebSocket.OPEN) {
alert('请先建立连接');
return;
}
const input = document.getElementById('messageInput');
const message = input.value.trim();
if (!message) {
return;
}
const messageData = {
type: 'chat_message',
content: message,
timestamp: new Date().toISOString()
};
if (currentChatId) {
messageData.chat_id = currentChatId;
}
ws.send(JSON.stringify(messageData));
addMessage('我', message, 'sent');
input.value = '';
}
function handleWebSocketMessage(data) {
switch(data.type) {
case 'new_message':
addMessage(data.message.sender_name || '用户', data.message.content, 'received');
break;
case 'user_typing':
addMessage('系统', `用户 ${data.user_id} 正在输入...`, 'received');
break;
case 'presence_update':
addMessage('系统', `用户 ${data.username} 状态变更为: ${data.status}`, 'received');
break;
case 'ping':
// 自动回复pong
ws.send(JSON.stringify({type: 'pong', timestamp: new Date().toISOString()}));
break;
default:
addMessage('系统', `收到消息: ${JSON.stringify(data)}`, 'received');
}
}
function addMessage(sender, content, className) {
const messagesDiv = document.getElementById('messages');
const messageDiv = document.createElement('div');
messageDiv.className = `message ${className}`;
messageDiv.innerHTML = `<strong>${sender}:</strong> ${content}`;
messagesDiv.appendChild(messageDiv);
messagesDiv.scrollTop = messagesDiv.scrollHeight;
}
// 输入框回车发送
document.getElementById('messageInput').addEventListener('keypress', function(e) {
if (e.key === 'Enter') {
sendMessage();
}
});
</script>
</body>
</html>
"""
return HTMLResponse(html)5. 业务服务层实现
5.1 聊天服务实现
python
# app/services/chat_service.py
"""
聊天业务服务
"""
from typing import List, Dict, Optional, Any, Union
from datetime import datetime, timedelta
from uuid import UUID
import logging
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, update, delete, and_, or_, func, desc
from sqlalchemy.orm import selectinload, joinedload
from app.database.models import (
User, ChatRoom, ChatMember, Message, Attachment,
UserSession, Notification
)
from app.database.redis_client import RedisClient
from app.models.schemas import (
ChatRoomCreate, ChatRoomUpdate, MessageCreate,
ChatMemberUpdate, ChatRoomResponse, MessageResponse
)
logger = logging.getLogger(__name__)
class ChatService:
"""聊天服务"""
def __init__(self, db: AsyncSession, redis_client: RedisClient = None):
self.db = db
self.redis_client = redis_client
# ========== 聊天室管理 ==========
async def create_chat_room(self, create_data: ChatRoomCreate,
creator_id: str) -> Dict[str, Any]:
"""创建聊天室"""
try:
# 创建聊天室记录
chat_room = ChatRoom(
name=create_data.name,
description=create_data.description,
room_type=create_data.room_type,
created_by=UUID(creator_id),
is_public=create_data.is_public,
max_members=create_data.max_members,
metadata=create_data.metadata or {},
settings=create_data.settings or {}
)
self.db.add(chat_room)
await self.db.flush() # 获取ID但不提交
# 添加创建者为成员（所有者）
chat_member = ChatMember(
chat_room_id=chat_room.id,
user_id=UUID(creator_id),
role="owner",
permissions={
"send_messages": True,
"invite_users": True,
"remove_users": True,
"edit_room": True
}
)
self.db.add(chat_member)
# 添加其他成员
for member_id in create_data.member_ids or []:
if str(member_id) != creator_id:
member = ChatMember(
chat_room_id=chat_room.id,
user_id=UUID(member_id),
role="member"
)
self.db.add(member)
await self.db.commit()
await self.db.refresh(chat_room)
# 缓存到Redis
if self.redis_client:
await self.redis_client.add_chat_member(
str(chat_room.id), creator_id, "owner"
)
for member_id in create_data.member_ids or []:
if str(member_id) != creator_id:
await self.redis_client.add_chat_member(
str(chat_room.id), str(member_id), "member"
)
# 转换为响应格式
result = await self._chat_room_to_dict(chat_room)
return result
except Exception as e:
await self.db.rollback()
logger.error(f"创建聊天室失败: {e}")
raise
async def get_chat_room(self, chat_id: str) -> Optional[Dict[str, Any]]:
"""获取聊天室详情"""
try:
# 尝试从Redis获取
if self.redis_client:
# 这里可以添加Redis缓存逻辑
pass
# 从数据库获取
stmt = select(ChatRoom).where(ChatRoom.id == UUID(chat_id))
result = await self.db.execute(stmt)
chat_room = result.scalar_one_or_none()
if not chat_room:
return None
return await self._chat_room_to_dict(chat_room)
except Exception as e:
logger.error(f"获取聊天室失败: {e}")
return None
async def update_chat_room(self, chat_id: str, update_data: ChatRoomUpdate,
user_id: str) -> Optional[Dict[str, Any]]:
"""更新聊天室信息"""
try:
# 检查用户权限
has_permission = await self._check_chat_permission(
chat_id, user_id, "edit_room"
)
if not has_permission:
raise PermissionError("没有编辑聊天室的权限")
# 构建更新数据
update_dict = update_data.dict(exclude_unset=True)
stmt = (
update(ChatRoom)
.where(ChatRoom.id == UUID(chat_id))
.values(**update_dict, updated_at=datetime.now())
.returning(ChatRoom)
)
result = await self.db.execute(stmt)
chat_room = result.scalar_one_or_none()
await self.db.commit()
if chat_room:
await self.db.refresh(chat_room)
return await self._chat_room_to_dict(chat_room)
return None
except Exception as e:
await self.db.rollback()
logger.error(f"更新聊天室失败: {e}")
raise
async def delete_chat_room(self, chat_id: str, user_id: str) -> bool:
"""删除聊天室"""
try:
# 检查用户是否为所有者
stmt = select(ChatMember).where(
and_(
ChatMember.chat_room_id == UUID(chat_id),
ChatMember.user_id == UUID(user_id),
ChatMember.role == "owner"
)
)
result = await self.db.execute(stmt)
member = result.scalar_one_or_none()
if not member:
raise PermissionError("只有所有者可以删除聊天室")
# 删除聊天室（级联删除消息和成员）
stmt = delete(ChatRoom).where(ChatRoom.id == UUID(chat_id))
await self.db.execute(stmt)
await self.db.commit()
# 清理Redis缓存
if self.redis_client:
# 清理相关缓存
pass
return True
except Exception as e:
await self.db.rollback()
logger.error(f"删除聊天室失败: {e}")
raise
# ========== 消息管理 ==========
async def create_message(self, user_id: str, chat_id: str,
content: str, message_type: str = "text",
reply_to: int = None, metadata: Dict = None) -> Dict[str, Any]:
"""创建消息"""
try:
# 检查用户是否在聊天室中
is_member = await self.is_user_in_chat(user_id, chat_id)
if not is_member:
raise PermissionError("用户不在聊天室中")
# 检查用户是否被禁言
is_muted = await self._is_user_muted(chat_id, user_id)
if is_muted:
raise PermissionError("用户已被禁言")
# 创建消息记录
message = Message(
chat_room_id=UUID(chat_id),
sender_id=UUID(user_id),
message_type=message_type,
content=content,
reply_to=reply_to,
metadata=metadata or {}
)
self.db.add(message)
await self.db.flush()
# 更新聊天室最后活动时间
stmt = (
update(ChatRoom)
.where(ChatRoom.id == UUID(chat_id))
.values(updated_at=datetime.now())
)
await self.db.execute(stmt)
await self.db.commit()
await self.db.refresh(message)
# 缓存消息到Redis
if self.redis_client:
message_dict = await self._message_to_dict(message)
await self.redis_client.cache_message(message_dict)
# 转换为响应格式
return await self._message_to_dict(message)
except Exception as e:
await self.db.rollback()
logger.error(f"创建消息失败: {e}")
raise
async def get_chat_messages(self, chat_id: str, limit: int = 50,
offset: int = 0, before: datetime = None,
after: datetime = None, user_id: str = None) -> List[Dict[str, Any]]:
"""获取聊天室消息"""
try:
# 构建查询条件
conditions = [Message.chat_room_id == UUID(chat_id)]
if before:
conditions.append(Message.created_at < before)
if after:
conditions.append(Message.created_at > after)
# 排除已删除的消息
conditions.append(Message.deleted_at.is_(None))
# 构建查询
stmt = (
select(Message)
.where(and_(*conditions))
.order_by(desc(Message.created_at))
.limit(limit)
.offset(offset)
.options(
selectinload(Message.sender),
selectinload(Message.reply_to)
)
)
result = await self.db.execute(stmt)
messages = result.scalars().all()
# 转换为字典列表
messages_list = []
for message in messages:
message_dict = await self._message_to_dict(message)
messages_list.append(message_dict)
# 如果提供了用户ID，标记已读消息
if user_id and messages_list:
await self._mark_messages_as_read(
chat_id, user_id, [msg["id"] for msg in messages_list]
)
return messages_list[::-1] # 按时间正序返回
except Exception as e:
logger.error(f"获取消息失败: {e}")
return []
async def update_message(self, message_id: int, user_id: str,
content: str) -> Optional[Dict[str, Any]]:
"""更新消息"""
try:
# 检查消息是否存在且属于用户
stmt = select(Message).where(
and_(
Message.id == message_id,
Message.sender_id == UUID(user_id),
Message.deleted_at.is_(None)
)
)
result = await self.db.execute(stmt)
message = result.scalar_one_or_none()
if not message:
raise ValueError("消息不存在或没有权限编辑")
# 检查是否允许编辑
chat_member_stmt = select(ChatMember).where(
and_(
ChatMember.chat_room_id == message.chat_room_id,
ChatMember.user_id == UUID(user_id)
)
)
result = await self.db.execute(chat_member_stmt)
member = result.scalar_one_or_none()
if not member or not member.permissions.get("allow_edit", True):
raise PermissionError("不允许编辑消息")
# 更新消息
message.content = content
message.is_edited = True
message.edited_at = datetime.now()
await self.db.commit()
await self.db.refresh(message)
# 更新Redis缓存
if self.redis_client:
message_dict = await self._message_to_dict(message)
await self.redis_client.cache_message(message_dict)
return await self._message_to_dict(message)
except Exception as e:
await self.db.rollback()
logger.error(f"更新消息失败: {e}")
raise
async def delete_message(self, message_id: int, user_id: str,
delete_type: str = "soft") -> bool:
"""删除消息"""
try:
# 检查消息是否存在
stmt = select(Message).where(
and_(
Message.id == message_id,
Message.deleted_at.is_(None)
)
)
result = await self.db.execute(stmt)
message = result.scalar_one_or_none()
if not message:
raise ValueError("消息不存在")
# 检查权限（发送者或管理员）
is_sender = str(message.sender_id) == user_id
if not is_sender:
# 检查是否为管理员
has_permission = await self._check_chat_permission(
str(message.chat_room_id), user_id, "remove_messages"
)
if not has_permission:
raise PermissionError("没有删除消息的权限")
if delete_type == "soft":
# 软删除
message.deleted_at = datetime.now()
else:
# 硬删除
await self.db.delete(message)
await self.db.commit()
# 清理Redis缓存
if self.redis_client:
cache_key = f"cache:message:{message_id}"
await self.redis_client.client.delete(cache_key)
return True
except Exception as e:
await self.db.rollback()
logger.error(f"删除消息失败: {e}")
raise
async def mark_message_as_read(self, message_id: int, user_id: str) -> bool:
"""标记消息为已读"""
try:
# 获取消息
stmt = select(Message).where(Message.id == message_id)
result = await self.db.execute(stmt)
message = result.scalar_one_or_none()
if not message:
return False
# 更新用户的最后阅读消息
stmt = (
update(ChatMember)
.where(
and_(
ChatMember.chat_room_id == message.chat_room_id,
ChatMember.user_id == UUID(user_id)
)
)
.values(last_read_message_id=message_id)
)
await self.db.execute(stmt)
await self.db.commit()
return True
except Exception as e:
await self.db.rollback()
logger.error(f"标记消息为已读失败: {e}")
return False
# ========== 成员管理 ==========
async def add_chat_member(self, chat_id: str, user_id: str,
inviter_id: str, role: str = "member") -> bool:
"""添加聊天室成员"""
try:
# 检查邀请者权限
has_permission = await self._check_chat_permission(
chat_id, inviter_id, "invite_users"
)
if not has_permission:
raise PermissionError("没有邀请用户的权限")
# 检查用户是否已在聊天室中
existing_stmt = select(ChatMember).where(
and_(
ChatMember.chat_room_id == UUID(chat_id),
ChatMember.user_id == UUID(user_id)
)
)
result = await self.db.execute(existing_stmt)
existing_member = result.scalar_one_or_none()
if existing_member:
if existing_member.is_banned:
raise PermissionError("用户已被封禁")
return True # 已经存在
# 添加成员
chat_member = ChatMember(
chat_room_id=UUID(chat_id),
user_id=UUID(user_id),
role=role
)
self.db.add(chat_member)
await self.db.commit()
# 更新Redis缓存
if self.redis_client:
await self.redis_client.add_chat_member(chat_id, user_id, role)
# 发送通知
await self._send_notification(
user_id,
"chat_invite",
f"您已被邀请加入聊天室",
{"chat_id": chat_id, "inviter_id": inviter_id}
)
return True
except Exception as e:
await self.db.rollback()
logger.error(f"添加聊天室成员失败: {e}")
raise
async def remove_chat_member(self, chat_id: str, user_id: str,
remover_id: str) -> bool:
"""移除聊天室成员"""
try:
# 检查移除者权限
has_permission = await self._check_chat_permission(
chat_id, remover_id, "remove_users"
)
if not has_permission:
raise PermissionError("没有移除用户的权限")
# 检查是否在移除自己
if user_id == remover_id:
# 允许用户自己离开
pass
else:
# 检查目标用户的角色
target_stmt = select(ChatMember).where(
and_(
ChatMember.chat_room_id == UUID(chat_id),
ChatMember.user_id == UUID(user_id)
)
)
result = await self.db.execute(target_stmt)
target_member = result.scalar_one_or_none()
if target_member and target_member.role == "owner":
raise PermissionError("不能移除所有者")
# 移除成员
stmt = delete(ChatMember).where(
and_(
ChatMember.chat_room_id == UUID(chat_id),
ChatMember.user_id == UUID(user_id)
)
)
await self.db.execute(stmt)
await self.db.commit()
# 更新Redis缓存
if self.redis_client:
await self.redis_client.remove_chat_member(chat_id, user_id)
return True
except Exception as e:
await self.db.rollback()
logger.error(f"移除聊天室成员失败: {e}")
raise
async def update_chat_member(self, chat_id: str, user_id: str,
update_data: ChatMemberUpdate, updater_id: str) -> bool:
"""更新聊天室成员"""
try:
# 检查更新者权限
has_permission = await self._check_chat_permission(
chat_id, updater_id, "edit_members"
)
if not has_permission:
raise PermissionError("没有更新成员的权限")
# 构建更新数据
update_dict = update_data.dict(exclude_unset=True)
stmt = (
update(ChatMember)
.where(
and_(
ChatMember.chat_room_id == UUID(chat_id),
ChatMember.user_id == UUID(user_id)
)
)
.values(**update_dict)
)
await self.db.execute(stmt)
await self.db.commit()
return True
except Exception as e:
await self.db.rollback()
logger.error(f"更新聊天室成员失败: {e}")
raise
# ========== 辅助方法 ==========
async def is_user_in_chat(self, user_id: str, chat_id: str) -> bool:
"""检查用户是否在聊天室中"""
try:
# 先检查Redis缓存
if self.redis_client:
members = await self.redis_client.get_chat_members(chat_id)
for member in members:
if member["user_id"] == user_id:
return True
# 检查数据库
stmt = select(ChatMember).where(
and_(
ChatMember.chat_room_id == UUID(chat_id),
ChatMember.user_id == UUID(user_id),
ChatMember.is_banned == False
)
)
result = await self.db.execute(stmt)
member = result.scalar_one_or_none()
return member is not None
except Exception as e:
logger.error(f"检查用户是否在聊天室中失败: {e}")
return False
async def get_user_chats(self, user_id: str, limit: int = 100,
offset: int = 0) -> List[Dict[str, Any]]:
"""获取用户的所有聊天室"""
try:
stmt = (
select(ChatRoom)
.join(ChatMember, ChatMember.chat_room_id == ChatRoom.id)
.where(
and_(
ChatMember.user_id == UUID(user_id),
ChatMember.is_banned == False
)
)
.order_by(desc(ChatRoom.updated_at))
.limit(limit)
.offset(offset)
.options(selectinload(ChatRoom.members))
)
result = await self.db.execute(stmt)
chat_rooms = result.scalars().all()
chats_list = []
for chat_room in chat_rooms:
chat_dict = await self._chat_room_to_dict(chat_room)
chats_list.append(chat_dict)
return chats_list
except Exception as e:
logger.error(f"获取用户聊天室失败: {e}")
return []
async def get_chat_members_from_db(self, chat_id: str) -> List[Dict[str, Any]]:
"""从数据库获取聊天室成员"""
try:
stmt = (
select(ChatMember)
.where(ChatMember.chat_room_id == UUID(chat_id))
.options(selectinload(ChatMember.user))
)
result = await self.db.execute(stmt)
members = result.scalars().all()
members_list = []
for member in members:
member_dict = {
"user_id": str(member.user_id),
"username": member.user.username if member.user else None,
"role": member.role,
"permissions": member.permissions,
"joined_at": member.joined_at.isoformat() if member.joined_at else None
}
members_list.append(member_dict)
return members_list
except Exception as e:
logger.error(f"获取聊天室成员失败: {e}")
return []
async def _check_chat_permission(self, chat_id: str, user_id: str,
permission: str) -> bool:
"""检查用户在聊天室中的权限"""
try:
stmt = select(ChatMember).where(
and_(
ChatMember.chat_room_id == UUID(chat_id),
ChatMember.user_id == UUID(user_id)
)
)
result = await self.db.execute(stmt)
member = result.scalar_one_or_none()
if not member or member.is_banned:
return False
# 角色基础权限
role_permissions = {
"owner": ["send_messages", "invite_users", "remove_users",
"edit_room", "edit_members", "remove_messages"],
"admin": ["send_messages", "invite_users", "remove_users",
"edit_members", "remove_messages"],
"moderator": ["send_messages", "remove_messages"],
"member": ["send_messages"]
}
# 检查角色权限
if permission in role_permissions.get(member.role, []):
return True
# 检查自定义权限
if member.permissions and member.permissions.get(permission, False):
return True
return False
except Exception as e:
logger.error(f"检查权限失败: {e}")
return False
async def _is_user_muted(self, chat_id: str, user_id: str) -> bool:
"""检查用户是否被禁言"""
try:
stmt = select(ChatMember).where(
and_(
ChatMember.chat_room_id == UUID(chat_id),
ChatMember.user_id == UUID(user_id)
)
)
result = await self.db.execute(stmt)
member = result.scalar_one_or_none()
return member and member.is_muted
except Exception as e:
logger.error(f"检查用户禁言状态失败: {e}")
return False
async def _mark_messages_as_read(self, chat_id: str, user_id: str,
message_ids: List[int]) -> bool:
"""批量标记消息为已读"""
try:
# 获取最新的消息ID
if not message_ids:
return False
latest_message_id = max(message_ids)
# 更新用户的最后阅读消息
stmt = (
update(ChatMember)
.where(
and_(
ChatMember.chat_room_id == UUID(chat_id),
ChatMember.user_id == UUID(user_id)
)
)
.values(last_read_message_id=latest_message_id)
)
await self.db.execute(stmt)
await self.db.commit()
# 重置未读计数
if self.redis_client:
await self.redis_client.reset_unread_count(chat_id, user_id)
return True
except Exception as e:
await self.db.rollback()
logger.error(f"标记消息为已读失败: {e}")
return False
async def _send_notification(self, user_id: str, notification_type: str,
content: str, data: Dict = None) -> bool:
"""发送通知"""
try:
notification = Notification(
user_id=UUID(user_id),
notification_type=notification_type,
content=content,
data=data or {}
)
self.db.add(notification)
await self.db.commit()
return True
except Exception as e:
await self.db.rollback()
logger.error(f"发送通知失败: {e}")
return False
async def _chat_room_to_dict(self, chat_room: ChatRoom) -> Dict[str, Any]:
"""将ChatRoom对象转换为字典"""
return {
"id": str(chat_room.id),
"name": chat_room.name,
"description": chat_room.description,
"room_type": chat_room.room_type,
"avatar_url": chat_room.avatar_url,
"created_by": str(chat_room.created_by) if chat_room.created_by else None,
"is_public": chat_room.is_public,
"max_members": chat_room.max_members,
"metadata": chat_room.metadata,
"settings": chat_room.settings,
"created_at": chat_room.created_at.isoformat() if chat_room.created_at else None,
"updated_at": chat_room.updated_at.isoformat() if chat_room.updated_at else None,
"member_count": len(chat_room.members) if chat_room.members else 0,
"last_message": await self._get_last_message(chat_room.id) if chat_room.messages else None
}
async def _get_last_message(self, chat_room_id: UUID) -> Optional[Dict[str, Any]]:
"""获取最后一条消息"""
try:
stmt = (
select(Message)
.where(
and_(
Message.chat_room_id == chat_room_id,
Message.deleted_at.is_(None)
)
)
.order_by(desc(Message.created_at))
.limit(1)
.options(selectinload(Message.sender))
)
result = await self.db.execute(stmt)
message = result.scalar_one_or_none()
if message:
return await self._message_to_dict(message)
return None
except Exception as e:
logger.error(f"获取最后一条消息失败: {e}")
return None
async def _message_to_dict(self, message: Message) -> Dict[str, Any]:
"""将Message对象转换为字典"""
return {
"id": message.id,
"uuid": str(message.uuid),
"chat_room_id": str(message.chat_room_id),
"sender_id": str(message.sender_id),
"sender_name": message.sender.username if message.sender else None,
"sender_avatar": message.sender.avatar_url if message.sender else None,
"message_type": message.message_type,
"content": message.content,
"metadata": message.metadata,
"reply_to_id": message.reply_to_id,
"reply_to": await self._message_to_dict(message.reply_to) if message.reply_to else None,
"is_edited": message.is_edited,
"edited_at": message.edited_at.isoformat() if message.edited_at else None,
"status": message.status,
"deleted_at": message.deleted_at.isoformat() if message.deleted_at else None,
"created_at": message.created_at.isoformat() if message.created_at else None,
"updated_at": message.updated_at.isoformat() if message.updated_at else None
}5.2 用户服务实现
python
# app/services/user_service.py
"""
用户管理服务
"""
from typing import List, Dict, Optional, Any
from datetime import datetime, timedelta
from uuid import UUID
import logging
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, update, and_, or_, func
from sqlalchemy.orm import selectinload
from app.database.models import User, ChatRoom, ChatMember, UserSession
from app.core.security import (
get_password_hash, verify_password,
create_access_token, create_refresh_token
)
from app.models.schemas import (
UserCreate, UserUpdate, UserLogin,
UserResponse, TokenResponse
)
logger = logging.getLogger(__name__)
class UserService:
"""用户服务"""
def __init__(self, db: AsyncSession):
self.db = db
async def create_user(self, user_data: UserCreate) -> Dict[str, Any]:
"""创建用户"""
try:
# 检查用户名和邮箱是否已存在
existing_stmt = select(User).where(
or_(
User.username == user_data.username,
User.email == user_data.email
)
)
result = await self.db.execute(existing_stmt)
existing_user = result.scalar_one_or_none()
if existing_user:
if existing_user.username == user_data.username:
raise ValueError("用户名已存在")
if existing_user.email == user_data.email:
raise ValueError("邮箱已存在")
# 创建用户
hashed_password = get_password_hash(user_data.password)
user = User(
username=user_data.username,
email=user_data.email,
phone=user_data.phone,
hashed_password=hashed_password,
full_name=user_data.full_name,
avatar_url=user_data.avatar_url,
bio=user_data.bio
)
self.db.add(user)
await self.db.commit()
await self.db.refresh(user)
# 创建默认会话（可选）
# 这里可以添加创建默认聊天室等逻辑
return await self._user_to_dict(user)
except Exception as e:
await self.db.rollback()
logger.error(f"创建用户失败: {e}")
raise
async def authenticate_user(self, login_data: UserLogin) -> Optional[Dict[str, Any]]:
"""用户认证"""
try:
# 查找用户（支持用户名或邮箱登录）
stmt = select(User).where(
or_(
User.username == login_data.username,
User.email == login_data.username
)
)
result = await self.db.execute(stmt)
user = result.scalar_one_or_none()
if not user:
return None
# 验证密码
if not verify_password(login_data.password, user.hashed_password):
return None
# 更新最后登录时间
user.last_seen = datetime.now()
await self.db.commit()
return await self._user_to_dict(user)
except Exception as e:
await self.db.rollback()
logger.error(f"用户认证失败: {e}")
return None
async def create_user_session(self, user_id: str, device_info: Dict = None) -> Dict[str, Any]:
"""创建用户会话"""
try:
# 创建访问令牌和刷新令牌
access_token = create_access_token(data={"sub": user_id})
refresh_token = create_refresh_token(data={"sub": user_id})
# 计算过期时间
expires_at = datetime.now() + timedelta(days=7)
# 创建会话记录
session = UserSession(
user_id=UUID(user_id),
device_id=device_info.get("device_id") if device_info else None,
device_type=device_info.get("device_type") if device_info else None,
ip_address=device_info.get("ip_address") if device_info else None,
access_token=access_token,
refresh_token=refresh_token,
expires_at=expires_at
)
self.db.add(session)
await self.db.commit()
await self.db.refresh(session)
return {
"access_token": access_token,
"refresh_token": refresh_token,
"token_type": "bearer",
"expires_at": expires_at.isoformat(),
"session_id": str(session.id)
}
except Exception as e:
await self.db.rollback()
logger.error(f"创建用户会话失败: {e}")
raise
async def refresh_token(self, refresh_token: str) -> Optional[Dict[str, Any]]:
"""刷新访问令牌"""
try:
# 查找会话
stmt = select(UserSession).where(
and_(
UserSession.refresh_token == refresh_token,
UserSession.is_active == True,
UserSession.expires_at > datetime.now()
)
)
result = await self.db.execute(stmt)
session = result.scalar_one_or_none()
if not session:
return None
# 创建新的访问令牌
new_access_token = create_access_token(data={"sub": str(session.user_id)})
# 更新会话
session.access_token = new_access_token
session.last_activity = datetime.now()
await self.db.commit()
return {
"access_token": new_access_token,
"token_type": "bearer",
"session_id": str(session.id)
}
except Exception as e:
await self.db.rollback()
logger.error(f"刷新令牌失败: {e}")
return None
async def logout(self, session_id: str, user_id: str) -> bool:
"""用户登出"""
try:
stmt = select(UserSession).where(
and_(
UserSession.id == UUID(session_id),
UserSession.user_id == UUID(user_id),
UserSession.is_active == True
)
)
result = await self.db.execute(stmt)
session = result.scalar_one_or_none()
if not session:
return False
# 标记会话为不活跃
session.is_active = False
await self.db.commit()
return True
except Exception as e:
await self.db.rollback()
logger.error(f"用户登出失败: {e}")
return False
async def logout_all_sessions(self, user_id: str) -> bool:
"""登出所有会话"""
try:
stmt = (
update(UserSession)
.where(
and_(
UserSession.user_id == UUID(user_id),
UserSession.is_active == True
)
)
.values(is_active=False)
)
await self.db.execute(stmt)
await self.db.commit()
return True
except Exception as e:
await self.db.rollback()
logger.error(f"登出所有会话失败: {e}")
return False
async def get_user_by_id(self, user_id: str) -> Optional[Dict[str, Any]]:
"""根据ID获取用户"""
try:
stmt = select(User).where(User.id == UUID(user_id))
result = await self.db.execute(stmt)
user = result.scalar_one_or_none()
if not user:
return None
return await self._user_to_dict(user)
except Exception as e:
logger.error(f"获取用户失败: {e}")
return None
async def get_user_by_username(self, username: str) -> Optional[Dict[str, Any]]:
"""根据用户名获取用户"""
try:
stmt = select(User).where(User.username == username)
result = await self.db.execute(stmt)
user = result.scalar_one_or_none()
if not user:
return None
return await self._user_to_dict(user)
except Exception as e:
logger.error(f"获取用户失败: {e}")
return None
async def search_users(self, query: str, limit: int = 20,
offset: int = 0) -> List[Dict[str, Any]]:
"""搜索用户"""
try:
# 构建搜索条件
search_pattern = f"%{query}%"
stmt = (
select(User)
.where(
or_(
User.username.ilike(search_pattern),
User.full_name.ilike(search_pattern),
User.email.ilike(search_pattern)
)
)
.order_by(User.username)
.limit(limit)
.offset(offset)
)
result = await self.db.execute(stmt)
users = result.scalars().all()
users_list = []
for user in users:
users_list.append(await self._user_to_dict(user))
return users_list
except Exception as e:
logger.error(f"搜索用户失败: {e}")
return []
async def update_user(self, user_id: str, update_data: UserUpdate) -> Optional[Dict[str, Any]]:
"""更新用户信息"""
try:
# 构建更新数据
update_dict = update_data.dict(exclude_unset=True)
# 如果包含密码，需要哈希处理
if "password" in update_dict:
update_dict["hashed_password"] = get_password_hash(update_dict.pop("password"))
# 如果更新用户名或邮箱，需要检查是否已存在
if "username" in update_dict or "email" in update_dict:
conditions = []
if "username" in update_dict:
conditions.append(User.username == update_dict["username"])
if "email" in update_dict:
conditions.append(User.email == update_dict["email"])
# 排除当前用户
existing_stmt = select(User).where(
and_(
or_(*conditions),
User.id != UUID(user_id)
)
)
result = await self.db.execute(existing_stmt)
existing_user = result.scalar_one_or_none()
if existing_user:
raise ValueError("用户名或邮箱已存在")
# 更新用户
stmt = (
update(User)
.where(User.id == UUID(user_id))
.values(**update_dict, updated_at=datetime.now())
.returning(User)
)
result = await self.db.execute(stmt)
user = result.scalar_one_or_none()
await self.db.commit()
if user:
await self.db.refresh(user)
return await self._user_to_dict(user)
return None
except Exception as e:
await self.db.rollback()
logger.error(f"更新用户失败: {e}")
raise
async def delete_user(self, user_id: str) -> bool:
"""删除用户"""
try:
# 注意：实际应用中可能需要软删除或归档
stmt = delete(User).where(User.id == UUID(user_id))
await self.db.execute(stmt)
await self.db.commit()
return True
except Exception as e:
await self.db.rollback()
logger.error(f"删除用户失败: {e}")
return False
async def update_user_status(self, user_id: str, status: str) -> bool:
"""更新用户状态"""
try:
stmt = (
update(User)
.where(User.id == UUID(user_id))
.values(status=status, last_seen=datetime.now())
)
await self.db.execute(stmt)
await self.db.commit()
return True
except Exception as e:
await self.db.rollback()
logger.error(f"更新用户状态失败: {e}")
return False
async def get_user_friends(self, user_id: str) -> List[Dict[str, Any]]:
"""获取用户好友列表"""
try:
# 这里假设好友关系通过共同聊天室或专门的好友表建立
# 这是一个简化实现，实际应用中需要根据业务逻辑调整
# 查找用户所在的私聊聊天室
stmt = (
select(ChatRoom)
.join(ChatMember, ChatMember.chat_room_id == ChatRoom.id)
.where(
and_(
ChatRoom.room_type == "private",
ChatMember.user_id == UUID(user_id)
)
)
.options(selectinload(ChatRoom.members).selectinload(ChatMember.user))
)
result = await self.db.execute(stmt)
private_chats = result.scalars().all()
friends = []
for chat in private_chats:
for member in chat.members:
if str(member.user_id) != user_id:
friend_data = await self._user_to_dict(member.user)
friend_data["chat_id"] = str(chat.id)
friends.append(friend_data)
return friends
except Exception as e:
logger.error(f"获取好友列表失败: {e}")
return []
async def _user_to_dict(self, user: User) -> Dict[str, Any]:
"""将User对象转换为字典"""
return {
"id": str(user.id),
"username": user.username,
"email": user.email,
"phone": user.phone,
"full_name": user.full_name,
"avatar_url": user.avatar_url,
"bio": user.bio,
"status": user.status,
"settings": user.settings,
"created_at": user.created_at.isoformat() if user.created_at else None,
"updated_at": user.updated_at.isoformat() if user.updated_at else None,
"last_seen": user.last_seen.isoformat() if user.last_seen else None
}6. API路由和端点
6.1 认证路由
python
# app/api/endpoints/auth.py
"""
认证相关API端点
"""
from datetime import timedelta
from typing import Any, Dict
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordRequestForm
from app.core.security import get_current_user
from app.models.schemas import (
UserCreate, UserResponse, TokenResponse,
UserLogin, RefreshTokenRequest
)
from app.services.user_service import UserService
from app.database.postgres import get_db
from sqlalchemy.ext.asyncio import AsyncSession
router = APIRouter()
@router.post("/register", response_model=UserResponse)
async def register(
user_data: UserCreate,
db: AsyncSession = Depends(get_db)
) -> Any:
"""用户注册"""
try:
user_service = UserService(db)
user = await user_service.create_user(user_data)
return user
except ValueError as e:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=str(e)
)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="注册失败，请稍后重试"
)
@router.post("/login", response_model=TokenResponse)
async def login(
form_data: OAuth2PasswordRequestForm = Depends(),
db: AsyncSession = Depends(get_db)
) -> Any:
"""用户登录"""
# 将OAuth2表单数据转换为我们的登录格式
login_data = UserLogin(
username=form_data.username,
password=form_data.password
)
user_service = UserService(db)
user = await user_service.authenticate_user(login_data)
if not user:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="用户名或密码错误",
headers={"WWW-Authenticate": "Bearer"},
)
# 创建设备信息
device_info = {
"device_type": "web",
"user_agent": "API Login"
}
# 创建会话
session = await user_service.create_user_session(
user["id"], device_info
)
return session
@router.post("/login/json", response_model=TokenResponse)
async def login_json(
login_data: UserLogin,
db: AsyncSession = Depends(get_db)
) -> Any:
"""JSON格式用户登录"""
user_service = UserService(db)
user = await user_service.authenticate_user(login_data)
if not user:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="用户名或密码错误",
headers={"WWW-Authenticate": "Bearer"},
)
device_info = {
"device_type": "web",
"user_agent": "API Login"
}
session = await user_service.create_user_session(
user["id"], device_info
)
return session
@router.post("/refresh", response_model=TokenResponse)
async def refresh_token(
refresh_request: RefreshTokenRequest,
db: AsyncSession = Depends(get_db)
) -> Any:
"""刷新访问令牌"""
user_service = UserService(db)
tokens = await user_service.refresh_token(refresh_request.refresh_token)
if not tokens:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="无效的刷新令牌",
headers={"WWW-Authenticate": "Bearer"},
)
return tokens
@router.post("/logout")
async def logout(
session_id: str,
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db)
) -> Any:
"""用户登出"""
user_service = UserService(db)
success = await user_service.logout(session_id, current_user["id"])
if not success:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="登出失败"
)
return {"message": "登出成功"}
@router.post("/logout/all")
async def logout_all(
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db)
) -> Any:
"""登出所有会话"""
user_service = UserService(db)
success = await user_service.logout_all_sessions(current_user["id"])
if not success:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="登出所有会话失败"
)
return {"message": "已登出所有会话"}
@router.get("/me", response_model=UserResponse)
async def get_current_user_info(
current_user: Dict = Depends(get_current_user)
) -> Any:
"""获取当前用户信息"""
return current_user6.2 聊天路由
python
# app/api/endpoints/chat.py
"""
聊天相关API端点
"""
from typing import List, Dict, Any, Optional
from fastapi import APIRouter, Depends, HTTPException, status, Query, Path
from app.core.security import get_current_user
from app.models.schemas import (
ChatRoomCreate, ChatRoomUpdate, ChatRoomResponse,
MessageCreate, MessageResponse, ChatMemberUpdate
)
from app.services.chat_service import ChatService
from app.database.postgres import get_db
from app.database.redis_client import get_redis_client
from sqlalchemy.ext.asyncio import AsyncSession
router = APIRouter()
@router.get("/chats", response_model=List[ChatRoomResponse])
async def get_user_chats(
limit: int = Query(100, ge=1, le=500),
offset: int = Query(0, ge=0),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""获取用户的聊天室列表"""
try:
chat_service = ChatService(db, redis_client)
chats = await chat_service.get_user_chats(
current_user["id"], limit, offset
)
return chats
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"获取聊天室列表失败: {str(e)}"
)
@router.post("/chats", response_model=ChatRoomResponse)
async def create_chat(
chat_data: ChatRoomCreate,
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""创建聊天室"""
try:
chat_service = ChatService(db, redis_client)
chat = await chat_service.create_chat_room(
chat_data, current_user["id"]
)
return chat
except ValueError as e:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=str(e)
)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"创建聊天室失败: {str(e)}"
)
@router.get("/chats/{chat_id}", response_model=ChatRoomResponse)
async def get_chat(
chat_id: str = Path(..., description="聊天室ID"),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""获取聊天室详情"""
try:
chat_service = ChatService(db, redis_client)
# 检查用户是否在聊天室中
is_member = await chat_service.is_user_in_chat(current_user["id"], chat_id)
if not is_member:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="无权访问该聊天室"
)
chat = await chat_service.get_chat_room(chat_id)
if not chat:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="聊天室不存在"
)
return chat
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"获取聊天室详情失败: {str(e)}"
)
@router.put("/chats/{chat_id}", response_model=ChatRoomResponse)
async def update_chat(
chat_data: ChatRoomUpdate,
chat_id: str = Path(..., description="聊天室ID"),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""更新聊天室信息"""
try:
chat_service = ChatService(db, redis_client)
chat = await chat_service.update_chat_room(
chat_id, chat_data, current_user["id"]
)
if not chat:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="聊天室不存在"
)
return chat
except PermissionError as e:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=str(e)
)
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"更新聊天室失败: {str(e)}"
)
@router.delete("/chats/{chat_id}")
async def delete_chat(
chat_id: str = Path(..., description="聊天室ID"),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""删除聊天室"""
try:
chat_service = ChatService(db, redis_client)
success = await chat_service.delete_chat_room(chat_id, current_user["id"])
if not success:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="聊天室不存在"
)
return {"message": "聊天室删除成功"}
except PermissionError as e:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=str(e)
)
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"删除聊天室失败: {str(e)}"
)
@router.get("/chats/{chat_id}/messages", response_model=List[MessageResponse])
async def get_chat_messages(
chat_id: str = Path(..., description="聊天室ID"),
limit: int = Query(50, ge=1, le=200),
offset: int = Query(0, ge=0),
before: Optional[str] = Query(None, description="在此时间之前的消息，ISO格式"),
after: Optional[str] = Query(None, description="在此时间之后的消息，ISO格式"),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""获取聊天室消息"""
try:
chat_service = ChatService(db, redis_client)
# 检查用户是否在聊天室中
is_member = await chat_service.is_user_in_chat(current_user["id"], chat_id)
if not is_member:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="无权访问该聊天室"
)
# 解析时间参数
from datetime import datetime
before_dt = datetime.fromisoformat(before) if before else None
after_dt = datetime.fromisoformat(after) if after else None
messages = await chat_service.get_chat_messages(
chat_id, limit, offset, before_dt, after_dt, current_user["id"]
)
return messages
except ValueError as e:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"时间格式错误: {str(e)}"
)
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"获取消息失败: {str(e)}"
)
@router.post("/chats/{chat_id}/messages", response_model=MessageResponse)
async def send_message(
message_data: MessageCreate,
chat_id: str = Path(..., description="聊天室ID"),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""发送消息"""
try:
chat_service = ChatService(db, redis_client)
message = await chat_service.create_message(
user_id=current_user["id"],
chat_id=chat_id,
content=message_data.content,
message_type=message_data.message_type,
reply_to=message_data.reply_to,
metadata=message_data.metadata
)
return message
except PermissionError as e:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=str(e)
)
except ValueError as e:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=str(e)
)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"发送消息失败: {str(e)}"
)
@router.put("/messages/{message_id}", response_model=MessageResponse)
async def update_message(
content: str = Query(..., description="消息内容"),
message_id: int = Path(..., ge=1, description="消息ID"),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""更新消息"""
try:
chat_service = ChatService(db, redis_client)
message = await chat_service.update_message(
message_id, current_user["id"], content
)
return message
except (PermissionError, ValueError) as e:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=str(e)
)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"更新消息失败: {str(e)}"
)
@router.delete("/messages/{message_id}")
async def delete_message(
delete_type: str = Query("soft", description="删除类型：soft（软删除）或hard（硬删除）"),
message_id: int = Path(..., ge=1, description="消息ID"),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""删除消息"""
try:
if delete_type not in ["soft", "hard"]:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="删除类型必须是soft或hard"
)
chat_service = ChatService(db, redis_client)
success = await chat_service.delete_message(
message_id, current_user["id"], delete_type
)
if not success:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="消息不存在"
)
return {"message": "消息删除成功"}
except (PermissionError, ValueError) as e:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=str(e)
)
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"删除消息失败: {str(e)}"
)
@router.get("/chats/{chat_id}/members")
async def get_chat_members(
chat_id: str = Path(..., description="聊天室ID"),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""获取聊天室成员列表"""
try:
chat_service = ChatService(db, redis_client)
# 检查用户是否在聊天室中
is_member = await chat_service.is_user_in_chat(current_user["id"], chat_id)
if not is_member:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="无权访问该聊天室"
)
members = await chat_service.get_chat_members_from_db(chat_id)
return {"members": members}
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"获取成员列表失败: {str(e)}"
)
@router.post("/chats/{chat_id}/members/{user_id}")
async def add_chat_member(
chat_id: str = Path(..., description="聊天室ID"),
user_id: str = Path(..., description="用户ID"),
role: str = Query("member", description="成员角色"),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""添加聊天室成员"""
try:
chat_service = ChatService(db, redis_client)
success = await chat_service.add_chat_member(
chat_id, user_id, current_user["id"], role
)
if not success:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="添加成员失败"
)
return {"message": "成员添加成功"}
except PermissionError as e:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=str(e)
)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"添加成员失败: {str(e)}"
)
@router.delete("/chats/{chat_id}/members/{user_id}")
async def remove_chat_member(
chat_id: str = Path(..., description="聊天室ID"),
user_id: str = Path(..., description="用户ID"),
current_user: Dict = Depends(get_current_user),
db: AsyncSession = Depends(get_db),
redis_client = Depends(get_redis_client)
) -> Any:
"""移除聊天室成员"""
try:
chat_service = ChatService(db, redis_client)
success = await chat_service.remove_chat_member(
chat_id, user_id, current_user["id"]
)
if not success:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="移除成员失败"
)
return {"message": "成员移除成功"}
except PermissionError as e:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=str(e)
)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"移除成员失败: {str(e)}"
)7. 主应用程序配置
python
# app/main.py
"""
主应用程序入口
"""
import asyncio
from contextlib import asynccontextmanager
from datetime import datetime
import logging
from fastapi import FastAPI, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.exceptions import RequestValidationError
from fastapi.staticfiles import StaticFiles
from app.core.config import settings
from app.api.endpoints import auth, chat, users, websocket
from app.database.postgres import init_db, close_db
from app.database.redis_client import RedisClient
from app.core.websocket import ConnectionManager
from app.services.chat_service import ChatService
from app.services.user_service import UserService
# 配置日志
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.FileHandler('app.log'),
logging.StreamHandler()
]
)
logger = logging.getLogger(__name__)
# Redis客户端
redis_client = RedisClient(
host=settings.REDIS_HOST,
port=settings.REDIS_PORT,
password=settings.REDIS_PASSWORD,
db=settings.REDIS_DB
)
@asynccontextmanager
async def lifespan(app: FastAPI):
"""应用生命周期管理"""
# 启动时
logger.info("应用程序启动中...")
# 初始化数据库
await init_db()
logger.info("数据库初始化完成")
# 初始化Redis
await redis_client.initialize()
logger.info("Redis初始化完成")
# 启动后台任务
task = asyncio.create_task(background_task())
yield
# 关闭时
logger.info("应用程序关闭中...")
# 取消后台任务
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
# 关闭数据库连接
await close_db()
# 关闭Redis连接
await redis_client.close()
logger.info("应用程序已关闭")
# 创建FastAPI应用
app = FastAPI(
title=settings.PROJECT_NAME,
version=settings.VERSION,
description="高性能实时聊天系统API",
docs_url="/api/docs",
redoc_url="/api/redoc",
openapi_url="/api/openapi.json",
lifespan=lifespan
)
# 添加CORS中间件
app.add_middleware(
CORSMiddleware,
allow_origins=settings.BACKEND_CORS_ORIGINS,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 异常处理
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
"""请求验证异常处理"""
errors = []
for error in exc.errors():
errors.append({
"field": ".".join(str(loc) for loc in error["loc"]),
"message": error["msg"],
"type": error["type"]
})
return JSONResponse(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
content={
"detail": "请求参数验证失败",
"errors": errors
}
)
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
"""全局异常处理"""
logger.error(f"未处理的异常: {exc}", exc_info=True)
return JSONResponse(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
content={
"detail": "服务器内部错误",
"message": str(exc)
}
)
# 挂载静态文件
app.mount("/static", StaticFiles(directory="static"), name="static")
# 注册路由
app.include_router(auth.router, prefix="/api/auth", tags=["认证"])
app.include_router(users.router, prefix="/api/users", tags=["用户"])
app.include_router(chat.router, prefix="/api/chat", tags=["聊天"])
app.include_router(websocket.router, prefix="/api/ws", tags=["WebSocket"])
# 健康检查端点
@app.get("/health")
async def health_check():
"""健康检查"""
return {
"status": "healthy",
"timestamp": datetime.now().isoformat(),
"service": settings.PROJECT_NAME,
"version": settings.VERSION
}
@app.get("/")
async def root():
"""根端点"""
return {
"message": "欢迎使用实时聊天系统API",
"docs": "/api/docs",
"version": settings.VERSION
}
# 后台任务
async def background_task():
"""后台任务"""
try:
while True:
# 清理不活跃的WebSocket连接
# 这里需要访问connection_manager，实际应用中需要适当的依赖注入
# 其他后台任务...
await asyncio.sleep(60) # 每分钟执行一次
except asyncio.CancelledError:
logger.info("后台任务已取消")
except Exception as e:
logger.error(f"后台任务异常: {e}")
# 依赖注入函数（用于WebSocket路由）
async def get_chat_service():
"""获取聊天服务实例"""
from app.database.postgres import async_session
async with async_session() as session:
chat_service = ChatService(session, redis_client)
yield chat_service
async def get_user_service():
"""获取用户服务实例"""
from app.database.postgres import async_session
async with async_session() as session:
user_service = UserService(session)
yield user_service
# 将依赖函数添加到应用状态中
app.dependency_overrides.update({
"get_chat_service": get_chat_service,
"get_user_service": get_user_service,
"get_redis_client": lambda: redis_client
})
if __name__ == "__main__":
import uvicorn
uvicorn.run(
"app.main:app",
host=settings.HOST,
port=settings.PORT,
reload=settings.DEBUG,
log_level="info"
)8. Docker部署配置
yaml
# docker/docker-compose.yml
version: '3.8'
services:
# PostgreSQL数据库
postgres:
image: postgres:15-alpine
container_name: chat_postgres
environment:
POSTGRES_DB: ${POSTGRES_DB:-chatdb}
POSTGRES_USER: ${POSTGRES_USER:-chat_user}
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-chat_password}
volumes:
- postgres_data:/var/lib/postgresql/data
- ./init.sql:/docker-entrypoint-initdb.d/init.sql
ports:
- "5432:5432"
networks:
- chat_network
healthcheck:
test: ["CMD-SHELL", "pg_isready -U ${POSTGRES_USER}"]
interval: 10s
timeout: 5s
retries: 5
# Redis缓存
redis:
image: redis:7-alpine
container_name: chat_redis
command: redis-server --requirepass ${REDIS_PASSWORD:-redis_password}
volumes:
- redis_data:/data
ports:
- "6379:6379"
networks:
- chat_network
healthcheck:
test: ["CMD", "redis-cli", "ping"]
interval: 10s
timeout: 5s
retries: 5
# 主应用服务
api:
build:
context: ../..
dockerfile: docker/Dockerfile
container_name: chat_api
environment:
- DATABASE_URL=postgresql+asyncpg://${POSTGRES_USER}:${POSTGRES_PASSWORD}@postgres:5432/${POSTGRES_DB}
- REDIS_HOST=redis
- REDIS_PORT=6379
- REDIS_PASSWORD=${REDIS_PASSWORD}
- SECRET_KEY=${SECRET_KEY}
- DEBUG=${DEBUG:-false}
ports:
- "8000:8000"
depends_on:
postgres:
condition: service_healthy
redis:
condition: service_healthy
networks:
- chat_network
volumes:
- ./logs:/app/logs
restart: unless-stopped
# Nginx反向代理
nginx:
image: nginx:alpine
container_name: chat_nginx
ports:
- "80:80"
- "443:443"
volumes:
- ./nginx/nginx.conf:/etc/nginx/nginx.conf
- ./nginx/ssl:/etc/nginx/ssl
- ./logs/nginx:/var/log/nginx
depends_on:
- api
networks:
- chat_network
restart: unless-stopped
# 监控服务 (可选)
monitoring:
image: grafana/grafana:latest
container_name: chat_monitoring
environment:
- GF_SECURITY_ADMIN_PASSWORD=${GRAFANA_PASSWORD:-admin}
ports:
- "3000:3000"
volumes:
- grafana_data:/var/lib/grafana
networks:
- chat_network
restart: unless-stopped
volumes:
postgres_data:
redis_data:
grafana_data:
networks:
chat_network:
driver: bridge
dockerfile
# docker/Dockerfile
FROM python:3.11-slim
WORKDIR /app
# 安装系统依赖
RUN apt-get update && apt-get install -y \
gcc \
postgresql-client \
&& rm -rf /var/lib/apt/lists/*
# 复制依赖文件
COPY requirements/prod.txt .
# 安装Python依赖
RUN pip install --no-cache-dir -r prod.txt
# 复制应用代码
COPY app ./app
COPY alembic ./alembic
COPY alembic.ini .
# 创建日志目录
RUN mkdir -p /app/logs
# 创建非root用户
RUN useradd -m -u 1000 appuser && chown -R appuser:appuser /app
USER appuser
# 健康检查
HEALTHCHECK --interval=30s --timeout=30s --start-period=5s --retries=3 \
CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"
# 启动命令
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]9. 性能优化与扩展
9.1 性能优化策略
python
# app/core/optimization.py
"""
性能优化策略
"""
import asyncio
from functools import wraps
from typing import Callable, Any
import time
import logging
from cachetools import TTLCache, cached
from redis.asyncio import Redis
logger = logging.getLogger(__name__)
# 内存缓存
message_cache = TTLCache(maxsize=10000, ttl=300) # 缓存10000条消息，5分钟过期
user_cache = TTLCache(maxsize=5000, ttl=600) # 缓存5000个用户，10分钟过期
class RateLimiter:
"""速率限制器"""
def __init__(self, redis_client: Redis, key_prefix: str = "rate_limit"):
self.redis = redis_client
self.key_prefix = key_prefix
async def is_rate_limited(self, key: str, limit: int, window: int) -> bool:
"""检查是否达到速率限制"""
redis_key = f"{self.key_prefix}:{key}"
try:
# 使用Redis的INCR和EXPIRE实现滑动窗口
current = await self.redis.incr(redis_key)
if current == 1:
await self.redis.expire(redis_key, window)
return current > limit
except Exception as e:
logger.error(f"速率限制检查失败: {e}")
return False
class ConnectionPoolManager:
"""连接池管理器"""
def __init__(self):
self.pools = {}
async def get_redis_pool(self, config: dict) -> Redis:
"""获取Redis连接池"""
pool_key = f"{config.get('host')}:{config.get('port')}:{config.get('db')}"
if pool_key not in self.pools:
self.pools[pool_key] = Redis.from_url(
f"redis://{config.get('host')}:{config.get('port')}/{config.get('db')}",
password=config.get('password'),
decode_responses=True,
max_connections=50,
socket_keepalive=True
)
return self.pools[pool_key]
def async_timed_cache(ttl: int = 300):
"""异步定时缓存装饰器"""
def decorator(func: Callable) -> Callable:
cache = TTLCache(maxsize=100, ttl=ttl)
@wraps(func)
async def wrapper(*args, **kwargs):
# 生成缓存键
cache_key = f"{func.__name__}:{str(args)}:{str(kwargs)}"
# 检查缓存
if cache_key in cache:
return cache[cache_key]
# 执行函数
result = await func(*args, **kwargs)
# 缓存结果
cache[cache_key] = result
return result
return wrapper
return decorator
def performance_monitor(func: Callable) -> Callable:
"""性能监控装饰器"""
@wraps(func)
async def async_wrapper(*args, **kwargs):
start_time = time.time()
try:
result = await func(*args, **kwargs)
execution_time = time.time() - start_time
# 记录慢查询
if execution_time > 1.0: # 超过1秒
logger.warning(
f"慢查询: {func.__name__} 执行时间: {execution_time:.3f}秒"
)
return result
except Exception as e:
execution_time = time.time() - start_time
logger.error(
f"执行失败: {func.__name__} 执行时间: {execution_time:.3f}秒, 错误: {e}"
)
raise
@wraps(func)
def sync_wrapper(*args, **kwargs):
start_time = time.time()
try:
result = func(*args, **kwargs)
execution_time = time.time() - start_time
if execution_time > 1.0:
logger.warning(
f"慢查询: {func.__name__} 执行时间: {execution_time:.3f}秒"
)
return result
except Exception as e:
execution_time = time.time() - start_time
logger.error(
f"执行失败: {func.__name__} 执行时间: {execution_time:.3f}秒, 错误: {e}"
)
raise
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
class MessageBatcher:
"""消息批量处理器"""
def __init__(self, batch_size: int = 100, flush_interval: float = 1.0):
self.batch_size = batch_size
self.flush_interval = flush_interval
self.batch = []
self.lock = asyncio.Lock()
self.flush_task = None
async def add_message(self, message: dict):
"""添加消息到批量处理器"""
async with self.lock:
self.batch.append(message)
# 达到批量大小立即刷新
if len(self.batch) >= self.batch_size:
await self.flush()
elif not self.flush_task:
# 设置定时刷新任务
self.flush_task = asyncio.create_task(self.scheduled_flush())
async def scheduled_flush(self):
"""定时刷新"""
await asyncio.sleep(self.flush_interval)
await self.flush()
self.flush_task = None
async def flush(self):
"""刷新批量消息"""
if not self.batch:
return
async with self.lock:
batch_to_send = self.batch.copy()
self.batch.clear()
# 这里实现批量处理逻辑
# 例如：批量保存到数据库，批量发送通知等
try:
# 示例：批量保存消息到数据库
logger.info(f"批量处理 {len(batch_to_send)} 条消息")
# 实际实现需要根据业务逻辑调整
except Exception as e:
logger.error(f"批量处理失败: {e}")
# 可以考虑重试或记录错误
async def close(self):
"""关闭批量处理器"""
if self.flush_task:
self.flush_task.cancel()
try:
await self.flush_task
except asyncio.CancelledError:
pass
# 刷新剩余的消息
await self.flush()
# 使用示例
@async_timed_cache(ttl=60)
@performance_monitor
async def get_cached_user_profile(user_id: str, db_session):
"""获取缓存的用户资料（带性能监控）"""
# 这里是实际的数据库查询逻辑
# ...
pass9.2 水平扩展策略
数据存储层
消息队列层
API服务层
WebSocket网关层
负载均衡层
负载均衡器 1
负载均衡器 2
WebSocket服务器 1
WebSocket服务器 2
WebSocket服务器 3
API服务器 1
API服务器 2
API服务器 3
消息队列 1
消息队列 2
PostgreSQL主库
PostgreSQL从库 1
PostgreSQL从库 2
Redis集群节点 1
Redis集群节点 2
Redis集群节点 3
客户端
客户端
10. 完整项目总结
10.1 关键技术亮点
- 异步架构：全面使用async/await，支持高并发
- WebSocket实时通信：低延迟消息传递，支持心跳检测
- Redis缓存优化：用户状态、消息、会话缓存
- 连接池管理：数据库和Redis连接复用
- 消息批量处理：减少数据库写入压力
- 速率限制：防止API滥用
- 监控和日志：全面的性能监控和错误追踪
10.2 性能指标
| 指标 | 目标值 | 实现方法 |
|---|---|---|
| 消息延迟 | < 100ms | WebSocket直连，Redis Pub/Sub |
| 并发连接 | > 10,000 | 连接池，异步I/O |
| 消息吞吐量 | > 10,000 msg/s | 批量处理，Redis缓存 |
| 数据库查询 | < 50ms | 索引优化，查询缓存 |
| 内存使用 | < 2GB/节点 | 连接复用，缓存策略 |
10.3 部署建议
- 开发环境：使用docker-compose快速部署
- 测试环境：添加压力测试和集成测试
- 生产环境：使用Kubernetes集群部署，配置自动扩缩容
- 监控报警：集成Prometheus和Grafana监控
- 日志收集：使用ELK栈集中管理日志
10.4 未来扩展方向
- 视频/语音通话：集成WebRTC技术
- 消息加密：端到端加密支持
- AI助手：集成聊天机器人
- 文件存储：集成对象存储服务
- 移动推送：集成Firebase/APNS推送
- 国际化：多语言支持
- 插件系统：可扩展的功能模块
这个实时聊天系统提供了一个完整的解决方案，涵盖了从技术架构设计到具体实现的各个方面。通过合理的分层设计、异步编程模型和性能优化策略，系统能够支持高并发、低延迟的实时通信需求。开发者可以根据实际需求进行调整和扩展，构建适合自己业务场景的聊天系统。