Merge pull request nineaiyu#13 from nineaiyu/dev

feat: 增加消息推送功能

message/notify.py

@@ -21,6 +21,8 @@
 
 from common.base.magic import MagicCacheData
 from common.celery.utils import get_celery_task_log_path
+from message.utils import async_push_message
+from system.models import UserInfo
 from system.utils.serializer import UserInfoSerializer
 
 logger = logging.getLogger(__name__)
@@ -32,6 +34,13 @@ def get_userinfo(user_obj):
     return UserInfoSerializer(instance=user_obj).data
 
 
+@sync_to_async
+def get_user_pk(username):
+    try:
+        return UserInfo.objects.filter(username=username, is_active=True).values_list('pk', flat=True).first()
+    except UserInfo.DoesNotExist:
+        return
+
 @sync_to_async
 def token_auth(scope):
     cookies = scope.get('cookies')
@@ -107,6 +116,24 @@ async def receive_json(self, content, **kwargs):
             await self.channel_layer.group_send(
                 self.room_group_name, {"type": "chat_message", "data": data}
             )
+            text = data.get('text')
+            if text.startswith('@'):
+                target = text.split(' ')[0].split('@')
+                if len(target) > 1:
+                    target = target[1]
+                    try:
+                        pk = await get_user_pk(target)
+                        if pk:
+                            push_message = {
+                                'title': f"用户 {self.user_obj.username} 发来一条消息",
+                                'message': text,
+                                'level': 'info',
+                                'notice_type': {'label': '聊天室', 'value': 0},
+                                'message_type': 'chat_message',
+                            }
+                            await async_push_message(pk, push_message)
+                    except Exception as e:
+                        logger.error(e)
         else:
             await self.channel_layer.send(self.channel_name, {"type": action, "data": data})
 

message/utils.py

@@ -5,22 +5,26 @@
 # author : ly_13
 # date : 3/6/2024
 import json
+from typing import Dict
 
 from asgiref.sync import async_to_sync
 from channels.layers import get_channel_layer
 from django.conf import settings
 from rest_framework.utils import encoders
 
 
-@async_to_sync
-async def push_message(user_obj, message):
-    room_group_name = f"{settings.CACHE_KEY_TEMPLATE.get('user_websocket_key')}_{user_obj.pk}"
+async def async_push_message(user_pk: str | int, message: Dict, message_type='push_message'):
+    room_group_name = f"{settings.CACHE_KEY_TEMPLATE.get('user_websocket_key')}_{user_pk}"
     channel_layer = get_channel_layer()
     await channel_layer.group_send(room_group_name, {
-        'type': 'push_message',
-        'data': message if not isinstance(message, dict) else json.dumps(message, cls=encoders.JSONEncoder,
-                                                                         ensure_ascii=False)
+        'type': message_type,
+        'data': json.dumps(message, cls=encoders.JSONEncoder, ensure_ascii=False)
     })
+@async_to_sync
+async def push_message(user_pk: str | int, message: Dict, message_type='push_message'):
+    return await async_push_message(user_pk, message, message_type)
+
+
 @async_to_sync
 async def check_message(user_obj, message):
     room_group_name = f"{settings.CACHE_KEY_TEMPLATE.get('user_websocket_key')}_{user_obj.pk}"

system/utils/notify.py

@@ -5,7 +5,7 @@
 # author : ly_13
 # date : 9/8/2023
 
-from typing import List, Optional, Literal, Dict
+from typing import List, Dict
 
 from django.db.models import QuerySet
 
@@ -16,8 +16,17 @@
 SYSTEM = NoticeMessage.NoticeChoices.SYSTEM
 
 
+def push_notice_messages(notify_obj, pks):
+    notice_message = NoticeMessageSerializer(
+        fields=['pk', 'level', 'title', 'notice_type', 'message'],
+        instance=notify_obj).data
+    notice_message['message_type'] = 'notify_message'
+    for pk in pks:
+        push_message(pk, notice_message)
+    return notify_obj
+
 def base_notify(users: List | QuerySet, title: str, message: str, notice_type: int,
-                level: Optional[Literal['success', '', 'warning', 'error']], extra_json: Dict = None):
+                level: NoticeMessage.LevelChoices, extra_json: Dict = None):
     if isinstance(users, (QuerySet, list)):
         recipients = users
     else:
@@ -32,27 +41,18 @@ def base_notify(users: List | QuerySet, title: str, message: str, notice_type: i
         extra_json=extra_json
     )
     notify_obj.notice_user.set(recipients)
-    notice_message = NoticeMessageSerializer(
-        fields=['level', 'title', 'notice_type', 'message'],
-        instance=notify_obj).data
-    for user in recipients:
-        push_message(user, notice_message)
+    push_notice_messages(notify_obj, [user.pk for user in recipients])
     return notify_obj
 
 
 def notify_success(users: List | QuerySet, title: str, message: str, notice_type: int = SYSTEM,
                    extra_json: Dict = None):
-    return base_notify(users, title, message, notice_type, 'success', extra_json)
+    return base_notify(users, title, message, notice_type, NoticeMessage.LevelChoices.SUCCESS, extra_json)
 
 
 def notify_info(users: List | QuerySet, title: str, message: str, notice_type: int = SYSTEM, extra_json: Dict = None):
-    return base_notify(users, title, message, notice_type, '', extra_json)
-
-
-def notify_warning(users: List | QuerySet, title: str, message: str, notice_type: int = SYSTEM,
-                   extra_json: Dict = None):
-    return base_notify(users, title, message, notice_type, 'warning', extra_json)
+    return base_notify(users, title, message, notice_type, NoticeMessage.LevelChoices.PRIMARY, extra_json)
 
 
 def notify_error(users: List | QuerySet, title: str, message: str, notice_type: int = SYSTEM, extra_json: Dict = None):
-    return base_notify(users, title, message, notice_type, 'error', extra_json)
+    return base_notify(users, title, message, notice_type, NoticeMessage.LevelChoices.DANGER, extra_json)

system/utils/signal_handler.py

@@ -18,6 +18,7 @@
 from common.core.serializers import get_sub_serializer_fields
 from system.models import Menu, NoticeMessage, UserRole, UserInfo, NoticeUserRead, DeptInfo, DataPermission, \
     SystemConfig, ModelLabelField
+from system.utils.notify import push_notice_messages
 
 logger = logging.getLogger(__name__)
 
@@ -102,6 +103,8 @@ def invalid_notify_caches(instance, pk_set):
     if instance.notice_type == NoticeMessage.NoticeChoices.DEPT:
         pks = UserInfo.objects.filter(dept__in=pk_set).values_list('pk', flat=True)
     if pks:
+        if instance.publish:
+            push_notice_messages(instance, set(pks))
         for pk in set(pks):
             invalid_notify_cache(pk)
 
@@ -158,6 +161,8 @@ def clean_cache_handler(sender, instance, **kwargs):
         pk_set = None
         if instance.notice_type == NoticeMessage.NoticeChoices.NOTICE:
             invalid_notify_cache('*')
+            if instance.publish:
+                push_notice_messages(instance, UserInfo.objects.values_list('pk', flat=True))
         elif instance.notice_type == NoticeMessage.NoticeChoices.DEPT:
             pk_set = instance.notice_dept.values_list('pk', flat=True)
         elif instance.notice_type == NoticeMessage.NoticeChoices.ROLE:

system/views/admin/user.py

@@ -79,7 +79,7 @@ def reset_password(self, request, *args, **kwargs):
             instance.set_password(password)
             instance.modifier = request.user
             instance.save(update_fields=['password', 'modifier'])
-            notify.notify_info(users=instance, title="密码重置成功",
+            notify.notify_error(users=instance, title="密码重置成功",
                                message="密码被管理员重置成功")
             return ApiResponse()
         return ApiResponse(code=1001, detail='修改失败')