From 1169677286439216a030fd8ab17f49b5f7f4d6ba Mon Sep 17 00:00:00 2001 From: ibuler Date: Wed, 1 Feb 2023 16:43:43 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E9=87=8D=E6=9E=84=20notifications=20si?= =?UTF-8?q?te=20msg?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/audits/signal_handlers.py | 2 +- apps/notifications/api/site_msgs.py | 29 ++++---- apps/notifications/filters.py | 4 +- .../migrations/0004_auto_20230201_1614.py | 72 ++++++++++++++++++ apps/notifications/models/site_msg.py | 16 ++-- apps/notifications/serializers/site_msgs.py | 23 ++++-- apps/notifications/signal_handlers.py | 4 +- apps/notifications/site_msg.py | 73 ++++++------------- apps/notifications/urls/api_urls.py | 10 +-- apps/terminal/notifications.py | 20 ++--- 10 files changed, 156 insertions(+), 97 deletions(-) create mode 100644 apps/notifications/migrations/0004_auto_20230201_1614.py diff --git a/apps/audits/signal_handlers.py b/apps/audits/signal_handlers.py index 8cc396496..4f8be3670 100644 --- a/apps/audits/signal_handlers.py +++ b/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', diff --git a/apps/notifications/api/site_msgs.py b/apps/notifications/api/site_msgs.py index 5c9257c25..399f332dc 100644 --- a/apps/notifications/api/site_msgs.py +++ b/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) - seri.is_valid(raise_exception=True) - ids = seri.validated_data['ids'] + s = self.get_serializer(data=request.data) + s.is_valid(raise_exception=True) + 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) - seri.is_valid(raise_exception=True) - SiteMessageUtil.send_msg(**seri.validated_data, sender=request.user) + s = self.get_serializer(data=request.data) + s.is_valid(raise_exception=True) + SiteMessageUtil.send_msg(**s.validated_data, sender=request.user) return Response({'detail': 'ok'}) diff --git a/apps/notifications/filters.py b/apps/notifications/filters.py index c28b8a3f2..ee11eca2b 100644 --- a/apps/notifications/filters.py +++ b/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',) diff --git a/apps/notifications/migrations/0004_auto_20230201_1614.py b/apps/notifications/migrations/0004_auto_20230201_1614.py new file mode 100644 index 000000000..a59aeb8c2 --- /dev/null +++ b/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), + ), + ] diff --git a/apps/notifications/models/site_msg.py b/apps/notifications/models/site_msg.py index 170f21a3f..826aa7016 100644 --- a/apps/notifications/models/site_msg.py +++ b/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): - sitemessage = models.ForeignKey('notifications.SiteMessage', on_delete=models.CASCADE, db_constraint=False, - related_name='m2m_sitemessageusers') - user = models.ForeignKey('users.User', on_delete=models.CASCADE, db_constraint=False, - related_name='m2m_sitemessageusers') +class SiteMessage(JMSBaseModel): + content = models.ForeignKey('notifications.MessageContent', on_delete=models.CASCADE, + db_constraint=False, related_name='messages') + user = models.ForeignKey('users.User', on_delete=models.CASCADE, db_constraint=False) 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) diff --git a/apps/notifications/serializers/site_msgs.py b/apps/notifications/serializers/site_msgs.py index 1f157add5..351a04eac 100644 --- a/apps/notifications/serializers/site_msgs.py +++ b/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', - 'date_created', 'date_updated', 'sender', + 'id', 'subject', 'message', + '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' ] diff --git a/apps/notifications/signal_handlers.py b/apps/notifications/signal_handlers.py index c0b1f1c1c..74b700643 100644 --- a/apps/notifications/signal_handlers.py +++ b/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: diff --git a/apps/notifications/site_msg.py b/apps/notifications/site_msg.py index faa36b5f4..d4e604a51 100644 --- a/apps/notifications/site_msg.py +++ b/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: - if group_ids: - site_msg.groups.add(*group_ids) + elif group_ids: + site_msg.groups.add(*group_ids) - user_ids_from_group = User.groups.through.objects.filter( - usergroup_id__in=group_ids - ).values_list('user_id', flat=True) - user_ids = [*user_ids, *user_ids_from_group] + user_ids_from_group = User.groups.through.objects.filter( + usergroup_id__in=group_ids + ).values_list('user_id', flat=True) + 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( - m2m_sitemessageusers__user_id=user_id - ).distinct().annotate( - has_read=F('m2m_sitemessageusers__has_read'), - read_at=F('m2m_sitemessageusers__read_at') - ).order_by('-date_created') - - return site_msgs + site_msg_rels = SiteMessage.objects \ + .filter(user=user_id) \ + .prefetch_related('content') \ + .order_by('-date_created') + return site_msg_rels @classmethod def get_user_all_msgs_count(cls, user_id): - site_msgs_count = SiteMessageModel.objects.filter( - m2m_sitemessageusers__user_id=user_id + site_msgs_count = SiteMessage.objects.filter( + 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( - 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 + return cls.get_user_all_msgs(user_id).filter(has_read=has_read) @classmethod def get_user_unread_msgs_count(cls, user_id): - site_msgs_count = SiteMessageModel.objects.filter( - m2m_sitemessageusers__user_id=user_id, - m2m_sitemessageusers__has_read=False - ).distinct().count() + site_msgs_count = SiteMessage.objects \ + .filter(user=user_id, has_read=False) \ + .values_list('content', flat=True) \ + .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) - if msg_ids is not None: - q &= Q(sitemessage_id__in=msg_ids) - site_msg_users = SiteMessageUsers.objects.filter(q) - - 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')) + site_msgs = SiteMessage.objects.filter(user_id=user_id) + if msg_ids: + site_msgs = site_msgs.filter(id__in=msg_ids) + site_msgs.update(has_read=True, read_at=local_now()) diff --git a/apps/notifications/urls/api_urls.py b/apps/notifications/urls/api_urls.py index a65dd7b79..f35d7754f 100644 --- a/apps/notifications/urls/api_urls.py +++ b/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 += [ diff --git a/apps/terminal/notifications.py b/apps/terminal/notifications.py index 700bfde9b..a6b1a69a6 100644 --- a/apps/terminal/notifications.py +++ b/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['session'] = Session.objects.first().id + command = Command.objects.first() + 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: