From 356aa336ad91e3375fc1e7dbad72d8fee43a1c3f Mon Sep 17 00:00:00 2001 From: Innokentii Konstantinov Date: Fri, 26 Aug 2022 13:46:50 +0500 Subject: [PATCH] Remove auto-recreating logic for UserNotificationPolicy (#414) * Remove auto-recreating logic for UserNotificationPolicy It's removed to get rid of select_for_update on User on each notify_user_task * Fix and add tests * remove get_user_policies method --- .../incident_log_builder.py | 4 +-- engine/apps/alerts/tasks/notify_group.py | 2 +- engine/apps/alerts/tasks/notify_user.py | 4 +-- .../api/views/user_notification_policy.py | 10 ++++--- .../base/models/user_notification_policy.py | 27 ++++++------------- .../tests/test_user_notification_policy.py | 23 ++++++++++++++++ .../views/personal_notifications.py | 6 ++++- engine/apps/user_management/models/user.py | 3 +++ engine/common/exceptions/__init__.py | 7 ++++- engine/common/exceptions/exceptions.py | 4 +++ engine/conftest.py | 3 +++ 11 files changed, 62 insertions(+), 31 deletions(-) diff --git a/engine/apps/alerts/incident_log_builder/incident_log_builder.py b/engine/apps/alerts/incident_log_builder/incident_log_builder.py index c1582551..ca5ae047 100644 --- a/engine/apps/alerts/incident_log_builder/incident_log_builder.py +++ b/engine/apps/alerts/incident_log_builder/incident_log_builder.py @@ -659,9 +659,7 @@ class IncidentLogBuilder: # last passed step order + 1 notification_policy_order = last_user_log.notification_policy.order + 1 - notification_policies = UserNotificationPolicy.objects.get_or_create_for_user( - user=user_to_notify, important=important - ) + notification_policies = UserNotificationPolicy.objects.filter(user=user_to_notify, important=important) for notification_policy in notification_policies: future_notification = notification_policy.order >= notification_policy_order diff --git a/engine/apps/alerts/tasks/notify_group.py b/engine/apps/alerts/tasks/notify_group.py index d18c31b1..9803affb 100644 --- a/engine/apps/alerts/tasks/notify_group.py +++ b/engine/apps/alerts/tasks/notify_group.py @@ -58,7 +58,7 @@ def notify_group_task(alert_group_pk, escalation_policy_snapshot_order=None): if not user.is_notification_allowed: continue - notification_policies = UserNotificationPolicy.objects.get_or_create_for_user( + notification_policies = UserNotificationPolicy.objects.filter( user=user, important=escalation_policy_step == EscalationPolicy.STEP_NOTIFY_GROUP_IMPORTANT, ) diff --git a/engine/apps/alerts/tasks/notify_user.py b/engine/apps/alerts/tasks/notify_user.py index 57d902b2..a9ba1d15 100644 --- a/engine/apps/alerts/tasks/notify_user.py +++ b/engine/apps/alerts/tasks/notify_user.py @@ -73,9 +73,7 @@ def notify_user_task( user_has_notification = UserHasNotification.objects.filter(pk=user_has_notification.pk).select_for_update()[0] if previous_notification_policy_pk is None: - notification_policy = UserNotificationPolicy.objects.get_or_create_for_user( - user=user, important=important - ).first() + notification_policy = UserNotificationPolicy.objects.filter(user=user, important=important).first() # Here we collect a brief overview of notification steps configured for user to send it to thread. collected_steps_ids = [] next_notification_policy = notification_policy.next() diff --git a/engine/apps/api/views/user_notification_policy.py b/engine/apps/api/views/user_notification_policy.py index ae7e4bee..7231bcc5 100644 --- a/engine/apps/api/views/user_notification_policy.py +++ b/engine/apps/api/views/user_notification_policy.py @@ -26,6 +26,7 @@ from apps.base.models.user_notification_policy import BUILT_IN_BACKENDS, Notific from apps.user_management.models import User from common.api_helpers.exceptions import BadRequest from common.api_helpers.mixins import UpdateSerializerMixin +from common.exceptions import UserNotificationPolicyCouldNotBeDeleted from common.insight_log import EntityEvent, write_resource_insight_log @@ -55,14 +56,14 @@ class UserNotificationPolicyView(UpdateSerializerMixin, ModelViewSet): except ValueError: raise BadRequest(detail="Invalid user param") if user_id is None or user_id == self.request.user.public_primary_key: - queryset = self.model.objects.get_or_create_for_user(user=self.request.user, important=important) + queryset = self.model.objects.filter(user=self.request.user, important=important) else: try: target_user = User.objects.get(public_primary_key=user_id) except User.DoesNotExist: raise BadRequest(detail="User does not exist") - queryset = self.model.objects.get_or_create_for_user(user=target_user, important=important) + queryset = self.model.objects.filter(user=target_user, important=important) queryset = self.serializer_class.setup_eager_loading(queryset) @@ -111,7 +112,10 @@ class UserNotificationPolicyView(UpdateSerializerMixin, ModelViewSet): def perform_destroy(self, instance): user = instance.user prev_state = user.insight_logs_serialized - instance.delete() + try: + instance.delete() + except UserNotificationPolicyCouldNotBeDeleted: + raise BadRequest(detail="Can't delete last user notification policy") new_state = user.insight_logs_serialized write_resource_insight_log( instance=user, diff --git a/engine/apps/base/models/user_notification_policy.py b/engine/apps/base/models/user_notification_policy.py index e0a62275..b6444995 100644 --- a/engine/apps/base/models/user_notification_policy.py +++ b/engine/apps/base/models/user_notification_policy.py @@ -4,13 +4,14 @@ from typing import Tuple from django.conf import settings from django.core.exceptions import ValidationError from django.core.validators import MinLengthValidator -from django.db import models, transaction +from django.db import models from django.db.models import Q, QuerySet from django.utils import timezone from ordered_model.models import OrderedModel from apps.base.messaging import get_messaging_backends from apps.user_management.models import User +from common.exceptions import UserNotificationPolicyCouldNotBeDeleted from common.public_primary_keys import generate_public_primary_key, increase_public_primary_key_length @@ -69,24 +70,6 @@ def validate_channel_choice(value): class UserNotificationPolicyQuerySet(models.QuerySet): - def get_or_create_for_user(self, user: User, important: bool) -> "QuerySet[UserNotificationPolicy]": - with transaction.atomic(): - User.objects.select_for_update().get(pk=user.pk) - return self._get_or_create_for_user(user, important) - - def _get_or_create_for_user(self, user: User, important: bool) -> "QuerySet[UserNotificationPolicy]": - notification_policies = super().filter(user=user, important=important) - - if notification_policies.exists(): - return notification_policies - - if important: - policies = self.create_important_policies_for_user(user) - else: - policies = self.create_default_policies_for_user(user) - - return policies - def create_default_policies_for_user(self, user: User) -> "QuerySet[UserNotificationPolicy]": model = self.model @@ -197,6 +180,12 @@ class UserNotificationPolicy(OrderedModel): else: return "Not set" + def delete(self): + if UserNotificationPolicy.objects.filter(important=self.important, user=self.user).count() == 1: + raise UserNotificationPolicyCouldNotBeDeleted("Can't delete last user notification policy") + else: + super().delete() + class NotificationChannelOptions: """ diff --git a/engine/apps/base/tests/test_user_notification_policy.py b/engine/apps/base/tests/test_user_notification_policy.py index 5d0e1df7..41354491 100644 --- a/engine/apps/base/tests/test_user_notification_policy.py +++ b/engine/apps/base/tests/test_user_notification_policy.py @@ -9,6 +9,7 @@ from apps.base.models.user_notification_policy import ( validate_channel_choice, ) from apps.base.tests.messaging_backend import TestOnlyBackend +from common.exceptions import UserNotificationPolicyCouldNotBeDeleted @pytest.mark.parametrize( @@ -80,3 +81,25 @@ def test_extra_messaging_backends_details(): ) assert validate_channel_choice(channel_choice) is None + + +@pytest.mark.django_db +def test_unable_to_delete_last_notification_policy( + make_organization, + make_user_for_organization, + make_user_notification_policy, +): + organization = make_organization() + user = make_user_for_organization(organization) + + first_policy = make_user_notification_policy( + user, UserNotificationPolicy.Step.NOTIFY, notify_by=UserNotificationPolicy.NotificationChannel.SLACK + ) + + second_policy = make_user_notification_policy( + user, UserNotificationPolicy.Step.WAIT, wait_delay=timedelta(minutes=5) + ) + + first_policy.delete() + with pytest.raises(UserNotificationPolicyCouldNotBeDeleted): + second_policy.delete() diff --git a/engine/apps/public_api/views/personal_notifications.py b/engine/apps/public_api/views/personal_notifications.py index 44b251a3..b1288b4b 100644 --- a/engine/apps/public_api/views/personal_notifications.py +++ b/engine/apps/public_api/views/personal_notifications.py @@ -12,6 +12,7 @@ from apps.user_management.models import User from common.api_helpers.exceptions import BadRequest from common.api_helpers.mixins import RateLimitHeadersMixin, UpdateSerializerMixin from common.api_helpers.paginators import FiftyPageSizePaginator +from common.exceptions import UserNotificationPolicyCouldNotBeDeleted from common.insight_log import EntityEvent, write_resource_insight_log @@ -74,7 +75,10 @@ class PersonalNotificationView(RateLimitHeadersMixin, UpdateSerializerMixin, Mod def perform_destroy(self, instance): user = self.request.user prev_state = user.insight_logs_serialized - instance.delete() + try: + instance.delete() + except UserNotificationPolicyCouldNotBeDeleted: + raise BadRequest(detail="Can't delete last user notification policy") new_state = user.insight_logs_serialized write_resource_insight_log( instance=user, diff --git a/engine/apps/user_management/models/user.py b/engine/apps/user_management/models/user.py index c1a00669..041a0ec5 100644 --- a/engine/apps/user_management/models/user.py +++ b/engine/apps/user_management/models/user.py @@ -260,6 +260,9 @@ class User(models.Model): # TODO: check whether this signal can be moved to save method of the model @receiver(post_save, sender=User) def listen_for_user_model_save(sender, instance, created, *args, **kwargs): + if created: + instance.notification_policies.create_default_policies_for_user(instance) + instance.notification_policies.create_important_policies_for_user(instance) drop_cached_ical_for_custom_events_for_organization.apply_async( (instance.organization_id,), ) diff --git a/engine/common/exceptions/__init__.py b/engine/common/exceptions/__init__.py index d191b8f4..ec922fb4 100644 --- a/engine/common/exceptions/__init__.py +++ b/engine/common/exceptions/__init__.py @@ -1 +1,6 @@ -from .exceptions import MaintenanceCouldNotBeStartedError, TeamCanNotBeChangedError, UnableToSendDemoAlert # noqa: F401 +from .exceptions import ( # noqa: F401 + MaintenanceCouldNotBeStartedError, + TeamCanNotBeChangedError, + UnableToSendDemoAlert, + UserNotificationPolicyCouldNotBeDeleted, +) diff --git a/engine/common/exceptions/exceptions.py b/engine/common/exceptions/exceptions.py index 69318bd5..9adf0b47 100644 --- a/engine/common/exceptions/exceptions.py +++ b/engine/common/exceptions/exceptions.py @@ -17,3 +17,7 @@ class TeamCanNotBeChangedError(OperationCouldNotBePerformedError): class UnableToSendDemoAlert(OperationCouldNotBePerformedError): pass + + +class UserNotificationPolicyCouldNotBeDeleted(OperationCouldNotBePerformedError): + pass diff --git a/engine/conftest.py b/engine/conftest.py index 85fb9a3d..68ef50d5 100644 --- a/engine/conftest.py +++ b/engine/conftest.py @@ -68,6 +68,7 @@ from apps.telegram.tests.factories import ( TelegramVerificationCodeFactory, ) from apps.twilioapp.tests.factories import PhoneCallFactory, SMSFactory +from apps.user_management.models.user import User, listen_for_user_model_save from apps.user_management.tests.factories import OrganizationFactory, TeamFactory, UserFactory from common.constants.role import Role @@ -150,7 +151,9 @@ def make_organization(): @pytest.fixture def make_user_for_organization(): def _make_user_for_organization(organization, role=Role.ADMIN, **kwargs): + post_save.disconnect(listen_for_user_model_save, sender=User) user = UserFactory(organization=organization, role=role, **kwargs) + post_save.disconnect(listen_for_user_model_save, sender=User) return user return _make_user_for_organization