68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
|
import json
|
||
|
|
from channels.generic.websocket import AsyncWebsocketConsumer
|
||
|
|
from channels.db import database_sync_to_async
|
||
|
|
from channels.exceptions import StopConsumer
|
||
|
|
import logging
|
||
|
|
from rest_framework.authtoken.models import Token
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
class NotificationConsumer(AsyncWebsocketConsumer):
|
||
|
|
async def connect(self):
|
||
|
|
"""建立WebSocket连接"""
|
||
|
|
try:
|
||
|
|
# 获取token
|
||
|
|
headers = dict(self.scope['headers'])
|
||
|
|
auth_header = headers.get(b'authorization', b'').decode()
|
||
|
|
|
||
|
|
if not auth_header.startswith('Token '):
|
||
|
|
await self.close()
|
||
|
|
return
|
||
|
|
|
||
|
|
token_key = auth_header.split(' ')[1]
|
||
|
|
|
||
|
|
# 验证token
|
||
|
|
self.user = await self.get_user_from_token(token_key)
|
||
|
|
if not self.user:
|
||
|
|
await self.close()
|
||
|
|
return
|
||
|
|
|
||
|
|
# 为用户创建专属房间
|
||
|
|
self.room_name = f"notification_user_{self.user.id}"
|
||
|
|
await self.channel_layer.group_add(
|
||
|
|
self.room_name,
|
||
|
|
self.channel_name
|
||
|
|
)
|
||
|
|
await self.accept()
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
await self.close()
|
||
|
|
|
||
|
|
@database_sync_to_async
|
||
|
|
def get_user_from_token(self, token_key):
|
||
|
|
try:
|
||
|
|
token = Token.objects.select_related('user').get(key=token_key)
|
||
|
|
return token.user
|
||
|
|
except Token.DoesNotExist:
|
||
|
|
return None
|
||
|
|
|
||
|
|
async def disconnect(self, close_code):
|
||
|
|
"""断开WebSocket连接"""
|
||
|
|
try:
|
||
|
|
if hasattr(self, 'room_name'):
|
||
|
|
await self.channel_layer.group_discard(
|
||
|
|
self.room_name,
|
||
|
|
self.channel_name
|
||
|
|
)
|
||
|
|
logger.info(f"用户 {self.user.username} 已断开连接，关闭代码: {close_code}")
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"断开连接时发生错误: {str(e)}")
|
||
|
|
|
||
|
|
async def notification(self, event):
|
||
|
|
"""处理并发送通知消息"""
|
||
|
|
try:
|
||
|
|
await self.send(text_data=json.dumps(event))
|
||
|
|
logger.info(f"已发送通知给用户 {self.user.username}")
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"发送通知消息时发生错误: {str(e)}")
|