perf: 重构 notifications site msg

apps/audits/signal_handlers.py

@@ -291,7 +291,7 @@ def on_django_start_set_operate_log_monitor_models(sender, **kwargs):
     }
     exclude_models = {
         'UserPasswordHistory', 'ContentType',
-        'SiteMessage', 'SiteMessageUsers',
+        'MessageContent', 'SiteMessage',
         'PlatformAutomation', 'PlatformProtocol', 'Protocol',
         'HistoricalAccount', 'GatheredUser', 'ApprovalRule',
         'BaseAutomation', 'CeleryTask', 'Command', 'JobAuditLog',

apps/notifications/api/site_msgs.py

@@ -1,17 +1,16 @@
-from rest_framework.response import Response
-from rest_framework.mixins import ListModelMixin, RetrieveModelMixin
 from rest_framework.decorators import action
+from rest_framework.mixins import ListModelMixin, RetrieveModelMixin
+from rest_framework.response import Response
 
-from common.utils.http import is_true
-from common.permissions import IsValidUser
-from common.const.http import GET, PATCH, POST
 from common.api import JMSGenericViewSet
+from common.const.http import GET, PATCH, POST
+from common.permissions import IsValidUser
+from common.utils.http import is_true
 from ..serializers import (
-    SiteMessageDetailSerializer, SiteMessageIdsSerializer,
+    SiteMessageSerializer, SiteMessageIdsSerializer,
     SiteMessageSendSerializer,
 )
 from ..site_msg import SiteMessageUtil
-from ..filters import SiteMsgFilter
 
 __all__ = ('SiteMessageViewSet',)
 
@@ -19,11 +18,11 @@ __all__ = ('SiteMessageViewSet',)
 class SiteMessageViewSet(ListModelMixin, RetrieveModelMixin, JMSGenericViewSet):
     permission_classes = (IsValidUser,)
     serializer_classes = {
-        'default': SiteMessageDetailSerializer,
+        'default': SiteMessageSerializer,
         'mark_as_read': SiteMessageIdsSerializer,
         'send': SiteMessageSendSerializer,
     }
-    filterset_class = SiteMsgFilter
+    filterset_fields = ('has_read',)
 
     def get_queryset(self):
         user = self.request.user
@@ -44,9 +43,9 @@ class SiteMessageViewSet(ListModelMixin, RetrieveModelMixin, JMSGenericViewSet):
     @action(methods=[PATCH], detail=False, url_path='mark-as-read')
     def mark_as_read(self, request, **kwargs):
         user = request.user
-        seri = self.get_serializer(data=request.data)
+        s = self.get_serializer(data=request.data)
-        seri.is_valid(raise_exception=True)
+        s.is_valid(raise_exception=True)
-        ids = seri.validated_data['ids']
+        ids = s.validated_data['ids']
         SiteMessageUtil.mark_msgs_as_read(user.id, ids)
         return Response({'detail': 'ok'})
 
@@ -58,7 +57,7 @@ class SiteMessageViewSet(ListModelMixin, RetrieveModelMixin, JMSGenericViewSet):
 
     @action(methods=[POST], detail=False)
     def send(self, request, **kwargs):
-        seri = self.get_serializer(data=request.data)
+        s = self.get_serializer(data=request.data)
-        seri.is_valid(raise_exception=True)
+        s.is_valid(raise_exception=True)
-        SiteMessageUtil.send_msg(**seri.validated_data, sender=request.user)
+        SiteMessageUtil.send_msg(**s.validated_data, sender=request.user)
         return Response({'detail': 'ok'})

apps/notifications/filters.py

@@ -1,7 +1,7 @@
 import django_filters
 
 from common.drf.filters import BaseFilterSet
-from .models import SiteMessage
+from .models import MessageContent
 
 
 class SiteMsgFilter(BaseFilterSet):
@@ -14,5 +14,5 @@ class SiteMsgFilter(BaseFilterSet):
     has_read = django_filters.BooleanFilter(method='do_nothing')
 
     class Meta:
-        model = SiteMessage
+        model = MessageContent
         fields = ('has_read',)

apps/notifications/migrations/0004_auto_20230201_1614.py

@@ -0,0 +1,72 @@
+# Generated by Django 3.2.14 on 2023-02-01 08:14
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('users', '0041_auto_20221220_1956'),
+        ('notifications', '0003_auto_20221220_1956'),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name='sitemessageusers',
+            name='sitemessage',
+        ),
+        migrations.RemoveField(
+            model_name='sitemessageusers',
+            name='user',
+        ),
+        migrations.DeleteModel(
+            name='SiteMessage',
+        ),
+        migrations.DeleteModel(
+            name='SiteMessageUsers',
+        ),
+        migrations.CreateModel(
+            name='MessageContent',
+            fields=[
+                ('created_by', models.CharField(blank=True, max_length=128, null=True, verbose_name='Created by')),
+                ('updated_by', models.CharField(blank=True, max_length=128, null=True, verbose_name='Updated by')),
+                ('date_created', models.DateTimeField(auto_now_add=True, null=True, verbose_name='Date created')),
+                ('date_updated', models.DateTimeField(auto_now=True, verbose_name='Date updated')),
+                ('id', models.UUIDField(default=uuid.uuid4, primary_key=True, serialize=False)),
+                ('subject', models.CharField(max_length=1024)),
+                ('message', models.TextField()),
+                ('is_broadcast', models.BooleanField(default=False)),
+                ('groups', models.ManyToManyField(to='users.UserGroup')),
+                ('sender', models.ForeignKey(db_constraint=False, default=None, null=True, on_delete=django.db.models.deletion.DO_NOTHING, related_name='send_site_message', to=settings.AUTH_USER_MODEL)),
+            ],
+            options={
+                'abstract': False,
+            },
+        ),
+        migrations.CreateModel(
+            name='SiteMessage',
+            fields=[
+                ('created_by', models.CharField(blank=True, max_length=128, null=True, verbose_name='Created by')),
+                ('updated_by', models.CharField(blank=True, max_length=128, null=True, verbose_name='Updated by')),
+                ('date_created', models.DateTimeField(auto_now_add=True, null=True, verbose_name='Date created')),
+                ('date_updated', models.DateTimeField(auto_now=True, verbose_name='Date updated')),
+                ('id', models.UUIDField(default=uuid.uuid4, primary_key=True, serialize=False)),
+                ('has_read', models.BooleanField(default=False)),
+                ('read_at', models.DateTimeField(default=None, null=True)),
+                ('content', models.ForeignKey(db_constraint=False, on_delete=django.db.models.deletion.CASCADE, related_name='messages', to='notifications.messagecontent')),
+                ('user', models.ForeignKey(db_constraint=False, on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL)),
+            ],
+            options={
+                'abstract': False,
+            },
+        ),
+        migrations.AddField(
+            model_name='messagecontent',
+            name='users',
+            field=models.ManyToManyField(related_name='recv_site_messages', through='notifications.SiteMessage', to=settings.AUTH_USER_MODEL),
+        ),
+    ]

apps/notifications/models/site_msg.py

@@ -2,24 +2,24 @@ from django.db import models
 
 from common.db.models import JMSBaseModel
 
-__all__ = ('SiteMessageUsers', 'SiteMessage')
+__all__ = ('SiteMessage', 'MessageContent')
 
 
-class SiteMessageUsers(JMSBaseModel):
+class SiteMessage(JMSBaseModel):
-    sitemessage = models.ForeignKey('notifications.SiteMessage', on_delete=models.CASCADE, db_constraint=False,
+    content = models.ForeignKey('notifications.MessageContent', on_delete=models.CASCADE,
-                                    related_name='m2m_sitemessageusers')
+                                db_constraint=False, related_name='messages')
-    user = models.ForeignKey('users.User', on_delete=models.CASCADE, db_constraint=False,
+    user = models.ForeignKey('users.User', on_delete=models.CASCADE, db_constraint=False)
-                             related_name='m2m_sitemessageusers')
     has_read = models.BooleanField(default=False)
     read_at = models.DateTimeField(default=None, null=True)
     comment = ''
 
 
-class SiteMessage(JMSBaseModel):
+class MessageContent(JMSBaseModel):
     subject = models.CharField(max_length=1024)
     message = models.TextField()
     users = models.ManyToManyField(
-        'users.User', through=SiteMessageUsers, related_name='recv_site_messages'
+        'users.User', through=SiteMessage,
+        related_name='recv_site_messages'
     )
     groups = models.ManyToManyField('users.UserGroup')
     is_broadcast = models.BooleanField(default=False)

apps/notifications/serializers/site_msgs.py

@@ -1,7 +1,7 @@
-from rest_framework.serializers import ModelSerializer
 from rest_framework import serializers
+from rest_framework.serializers import ModelSerializer
 
-from ..models import SiteMessage
+from ..models import MessageContent
 
 
 class SenderMixin(ModelSerializer):
@@ -15,12 +15,23 @@ class SenderMixin(ModelSerializer):
             return ''
 
 
-class SiteMessageDetailSerializer(SenderMixin, ModelSerializer):
+class MessageContentSerializer(SenderMixin, ModelSerializer):
     class Meta:
-        model = SiteMessage
+        model = MessageContent
         fields = [
-            'id', 'subject', 'message', 'has_read', 'read_at',
+            'id', 'subject', 'message',
-            'date_created', 'date_updated', 'sender',
+            'date_created', 'date_updated',
+            'sender',
+        ]
+
+
+class SiteMessageSerializer(SenderMixin, ModelSerializer):
+    content = MessageContentSerializer(read_only=True)
+
+    class Meta:
+        model = MessageContent
+        fields = [
+            'id', 'has_read', 'read_at', 'content', 'date_created'
         ]
 
 

apps/notifications/signal_handlers.py

@@ -12,7 +12,7 @@ from common.utils import get_logger
 from common.utils.connection import RedisPubSub
 from notifications.backends import BACKEND
 from users.models import User
-from .models import SiteMessage, SystemMsgSubscription, UserMsgSubscription
+from .models import MessageContent, SystemMsgSubscription, UserMsgSubscription
 from .notifications import SystemMessage
 
 logger = get_logger(__name__)
@@ -26,7 +26,7 @@ class NewSiteMsgSubPub(LazyObject):
 new_site_msg_chan = NewSiteMsgSubPub()
 
 
-@receiver(post_save, sender=SiteMessage)
+@receiver(post_save, sender=MessageContent)
 @on_transaction_commit
 def on_site_message_create(sender, instance, created, **kwargs):
     if not created:

apps/notifications/site_msg.py

@@ -1,10 +1,9 @@
-from django.db.models import F, Q
 from django.db import transaction
 
-from common.utils.timezone import local_now
 from common.utils import get_logger
+from common.utils.timezone import local_now
 from users.models import User
-from .models import SiteMessage as SiteMessageModel, SiteMessageUsers
+from .models import MessageContent as SiteMessageModel, SiteMessage
 
 logger = get_logger(__file__)
 
@@ -17,11 +16,6 @@ class SiteMessageUtil:
         if not any((user_ids, group_ids, is_broadcast)):
             raise ValueError('No recipient is specified')
 
-        logger.info(f'Site message send: '
-                    f'user_ids={user_ids} '
-                    f'group_ids={group_ids} '
-                    f'subject={subject} '
-                    f'message={message}')
         with transaction.atomic():
             site_msg = SiteMessageModel.objects.create(
                 subject=subject, message=message,
@@ -30,65 +24,46 @@ class SiteMessageUtil:
 
             if is_broadcast:
                 user_ids = User.objects.all().values_list('id', flat=True)
-            else:
+            elif group_ids:
-                if group_ids:
+                site_msg.groups.add(*group_ids)
-                    site_msg.groups.add(*group_ids)
 
-                    user_ids_from_group = User.groups.through.objects.filter(
+                user_ids_from_group = User.groups.through.objects.filter(
-                        usergroup_id__in=group_ids
+                    usergroup_id__in=group_ids
-                    ).values_list('user_id', flat=True)
+                ).values_list('user_id', flat=True)
-                    user_ids = [*user_ids, *user_ids_from_group]
+                user_ids = [*user_ids, *user_ids_from_group]
 
             site_msg.users.add(*user_ids)
 
     @classmethod
     def get_user_all_msgs(cls, user_id):
-        site_msgs = SiteMessageModel.objects.filter(
+        site_msg_rels = SiteMessage.objects \
-            m2m_sitemessageusers__user_id=user_id
+            .filter(user=user_id) \
-        ).distinct().annotate(
+            .prefetch_related('content') \
-            has_read=F('m2m_sitemessageusers__has_read'),
+            .order_by('-date_created')
-            read_at=F('m2m_sitemessageusers__read_at')
+        return site_msg_rels
-        ).order_by('-date_created')
-
-        return site_msgs
 
     @classmethod
     def get_user_all_msgs_count(cls, user_id):
-        site_msgs_count = SiteMessageModel.objects.filter(
+        site_msgs_count = SiteMessage.objects.filter(
-            m2m_sitemessageusers__user_id=user_id
+            user_id=user_id
         ).distinct().count()
         return site_msgs_count
 
     @classmethod
     def filter_user_msgs(cls, user_id, has_read=False):
-        site_msgs = SiteMessageModel.objects.filter(
+        return cls.get_user_all_msgs(user_id).filter(has_read=has_read)
-            m2m_sitemessageusers__user_id=user_id,
-            m2m_sitemessageusers__has_read=has_read
-        ).distinct().annotate(
-            has_read=F('m2m_sitemessageusers__has_read'),
-            read_at=F('m2m_sitemessageusers__read_at')
-        ).order_by('-date_created')
-
-        return site_msgs
 
     @classmethod
     def get_user_unread_msgs_count(cls, user_id):
-        site_msgs_count = SiteMessageModel.objects.filter(
+        site_msgs_count = SiteMessage.objects \
-            m2m_sitemessageusers__user_id=user_id,
+            .filter(user=user_id, has_read=False) \
-            m2m_sitemessageusers__has_read=False
+            .values_list('content', flat=True) \
-        ).distinct().count()
+            .distinct().count()
         return site_msgs_count
 
     @classmethod
     def mark_msgs_as_read(cls, user_id, msg_ids=None):
-        q = Q(user_id=user_id) & Q(has_read=False)
+        site_msgs = SiteMessage.objects.filter(user_id=user_id)
-        if msg_ids is not None:
+        if msg_ids:
-            q &= Q(sitemessage_id__in=msg_ids)
+            site_msgs = site_msgs.filter(id__in=msg_ids)
-        site_msg_users = SiteMessageUsers.objects.filter(q)
+        site_msgs.update(has_read=True, read_at=local_now())
-
-        for site_msg_user in site_msg_users:
-            site_msg_user.has_read = True
-            site_msg_user.read_at = local_now()
-
-        SiteMessageUsers.objects.bulk_update(
-            site_msg_users, fields=('has_read', 'read_at'))

apps/notifications/urls/api_urls.py

@@ -1,7 +1,6 @@
-
-from rest_framework_bulk.routes import BulkRouter
-from django.urls import path
 from django.conf import settings
+from django.urls import path
+from rest_framework_bulk.routes import BulkRouter
 
 from notifications import api
 
@@ -10,11 +9,12 @@ app_name = 'notifications'
 router = BulkRouter()
 router.register('system-msg-subscription', api.SystemMsgSubscriptionViewSet, 'system-msg-subscription')
 router.register('user-msg-subscription', api.UserMsgSubscriptionViewSet, 'user-msg-subscription')
-router.register('site-message', api.SiteMessageViewSet, 'site-message')
+router.register('site-messages', api.SiteMessageViewSet, 'site-message')
 
 urlpatterns = [
     path('backends/', api.BackendListView.as_view(), name='backends')
-] + router.urls
+]
+urlpatterns += router.urls
 
 if settings.DEBUG:
     urlpatterns += [

apps/terminal/notifications.py

@@ -1,18 +1,17 @@
 from typing import Callable
 
-from django.utils.translation import ugettext_lazy as _
 from django.conf import settings
 from django.template.loader import render_to_string
+from django.utils.translation import ugettext_lazy as _
 
-from users.models import User
 from common.utils import get_logger, reverse
-from notifications.notifications import SystemMessage
-from terminal.models import Session, Command
-from notifications.models import SystemMsgSubscription
-from notifications.backends import BACKEND
 from common.utils import lazyproperty
 from common.utils.timezone import local_now_display
-
+from notifications.backends import BACKEND
+from notifications.models import SystemMsgSubscription
+from notifications.notifications import SystemMessage
+from terminal.models import Session, Command
+from users.models import User
 
 logger = get_logger(__name__)
 
@@ -75,8 +74,11 @@ class CommandAlertMessage(CommandAlertMixin, SystemMessage):
 
     @classmethod
     def gen_test_msg(cls):
-        command = Command.objects.first().to_dict()
+        command = Command.objects.first()
-        command['session'] = Session.objects.first().id
+        if not command:
+            command = Command(user='test', asset='test', input='test', session='111111111')
+        else:
+            command['session'] = Session.objects.first().id
         return cls(command)
 
     def get_html_msg(self) -> dict: