diff --git a/apps/orgs/signal_handlers/cache.py b/apps/orgs/signal_handlers/cache.py index 1d9ca3891..546d11311 100644 --- a/apps/orgs/signal_handlers/cache.py +++ b/apps/orgs/signal_handlers/cache.py @@ -1,33 +1,61 @@ +from functools import wraps from django.db.models.signals import post_save, pre_delete, pre_save, post_delete from django.dispatch import receiver from orgs.models import Organization from assets.models import Node -from perms.models import (AssetPermission, ApplicationPermission) +from perms.models import AssetPermission, ApplicationPermission from users.models import UserGroup, User +from users.signals import pre_user_leave_org from applications.models import Application from terminal.models import Session +from rbac.models import OrgRoleBinding, SystemRoleBinding from assets.models import Asset, SystemUser, Domain, Gateway from orgs.caches import OrgResourceStatisticsCache +from orgs.utils import current_org +from common.utils import get_logger + +logger = get_logger(__name__) -def refresh_user_amount_on_user_create_or_delete(user_id): - orgs = Organization.objects.filter(m2m_org_members__user_id=user_id).distinct() +def refresh_cache(name, org): + names = None + if isinstance(name, (str,)): + names = [name, ] + if isinstance(names, (list, tuple)): + for name in names: + OrgResourceStatisticsCache(org).expire(name) + OrgResourceStatisticsCache(Organization.root()).expire(name) + else: + logger.warning('refresh cache fail: {}'.format(name)) + + +def refresh_user_amount_cache(user): + orgs = user.orgs.distinct() for org in orgs: - org_cache = OrgResourceStatisticsCache(org) - org_cache.expire('users_amount') - OrgResourceStatisticsCache(Organization.root()).expire('users_amount') + refresh_cache('users_amount', org) -@receiver(post_save, sender=User) -def on_user_create_refresh_cache(sender, instance, created, **kwargs): +@receiver(post_save, sender=OrgRoleBinding) +def on_user_create_or_invite_refresh_cache(sender, instance, created, **kwargs): if created: - refresh_user_amount_on_user_create_or_delete(instance.id) + refresh_cache('users_amount', instance.org) + + +@receiver(post_save, sender=SystemRoleBinding) +def on_user_global_create_refresh_cache(sender, instance, created, **kwargs): + if created and current_org.is_root(): + refresh_cache('users_amount', current_org) + + +@receiver(pre_user_leave_org) +def on_user_remove_refresh_cache(sender, org=None, **kwargs): + refresh_cache('users_amount', org) @receiver(pre_delete, sender=User) def on_user_delete_refresh_cache(sender, instance, **kwargs): - refresh_user_amount_on_user_create_or_delete(instance.id) + refresh_user_amount_cache(instance) # @receiver(m2m_changed, sender=OrganizationMember) diff --git a/apps/users/models/user.py b/apps/users/models/user.py index 1a5ce634c..ede795d13 100644 --- a/apps/users/models/user.py +++ b/apps/users/models/user.py @@ -78,7 +78,7 @@ class AuthMixin: def is_history_password(self, password): allow_history_password_count = settings.OLD_PASSWORD_HISTORY_LIMIT_COUNT history_passwords = self.history_passwords.all() \ - .order_by('-date_created')[:int(allow_history_password_count)] + .order_by('-date_created')[:int(allow_history_password_count)] for history_password in history_passwords: if check_password(password, history_password.password): @@ -726,7 +726,7 @@ class User(AuthMixin, TokenMixin, RoleMixin, MFAMixin, AbstractUser): @classmethod def get_group_ids_by_user_id(cls, user_id): - group_ids = cls.groups.through.objects.filter(user_id=user_id)\ + group_ids = cls.groups.through.objects.filter(user_id=user_id) \ .distinct().values_list('usergroup_id', flat=True) group_ids = list(group_ids) return group_ids @@ -866,20 +866,20 @@ class User(AuthMixin, TokenMixin, RoleMixin, MFAMixin, AbstractUser): @property def all_orgs(self): from rbac.builtin import BuiltinRole - has_system_role = self.system_roles.all()\ - .exclude(name=BuiltinRole.system_user.name)\ + has_system_role = self.system_roles.all() \ + .exclude(name=BuiltinRole.system_user.name) \ .exists() if has_system_role: orgs = list(Organization.objects.all()) else: - orgs = list(self.orgs.all().distinct()) + orgs = list(self.orgs.distinct()) if self.has_perm('orgs.view_rootorg'): orgs = [Organization.root()] + orgs return orgs @property def my_orgs(self): - return list(self.orgs.all().distinct()) + return list(self.orgs.distinct()) class Meta: ordering = ['username']