diff --git a/apps/orgs/signal_handlers/cache.py b/apps/orgs/signal_handlers/cache.py index 1d9ca3891..525368431 100644 --- a/apps/orgs/signal_handlers/cache.py +++ b/apps/orgs/signal_handlers/cache.py @@ -1,33 +1,54 @@ +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 from assets.models import Asset, SystemUser, Domain, Gateway from orgs.caches import OrgResourceStatisticsCache +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(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 7a3b2b138..1f19e8020 100644 --- a/apps/users/models/user.py +++ b/apps/users/models/user.py @@ -22,6 +22,7 @@ from orgs.models import Organization from common.utils import date_expired_default, get_logger, lazyproperty, random_string from common import fields from django.db.models import TextChoices +from rbac.const import Scope from ..signals import post_user_change_password, post_user_leave_org, pre_user_leave_org __all__ = ['User', 'UserPasswordHistory'] @@ -204,6 +205,13 @@ class RoleManager(models.Manager): return return self.role_bindings.delete() + def refresh_user_amount_cache(self): + from orgs.signal_handlers.cache import refresh_user_amount_cache, refresh_cache + if current_org.is_root(): + refresh_cache('users_amount', current_org) + elif self.scope == Scope.org: + refresh_user_amount_cache(self.user) + def add(self, *roles): from rbac.models import RoleBinding items = [] @@ -226,6 +234,7 @@ class RoleManager(models.Manager): logger.error('Create role binding error: {}'.format(e)) finally: self.user.expire_users_rbac_perms_cache() + self.refresh_user_amount_cache() def set(self, roles): self.clear()