jumpserver/apps/orgs/signal_handlers/cache.py

133 lines
4.6 KiB
Python

from django.db.models.signals import post_save, pre_delete, pre_save, post_delete
from django.dispatch import receiver
from accounts.models import Account
from assets.models import Asset, Domain
from assets.models import Node
from common.decorators import merge_delay_run
from common.utils import get_logger
from orgs.caches import OrgResourceStatisticsCache
from orgs.models import Organization
from orgs.utils import current_org
from perms.models import AssetPermission
from rbac.models import OrgRoleBinding, SystemRoleBinding, RoleBinding
from terminal.models import Session
from users.models import UserGroup, User
from users.signals import pre_user_leave_org
logger = get_logger(__name__)
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_all_orgs_user_amount_cache(user):
orgs = user.orgs.distinct()
for org in orgs:
refresh_cache('users_amount', org)
refresh_cache('new_users_amount_this_week', org)
@receiver(post_save, sender=OrgRoleBinding)
def on_user_create_or_invite_refresh_cache(sender, instance, created, **kwargs):
if created:
refresh_cache('users_amount', instance.org)
refresh_cache('new_users_amount_this_week', 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)
refresh_cache('new_users_amount_this_week', current_org)
@receiver(pre_user_leave_org)
def on_user_remove_refresh_cache(sender, org=None, **kwargs):
refresh_cache('users_amount', org)
refresh_cache('new_users_amount_this_week', org)
@receiver(pre_delete, sender=User)
def on_user_delete_refresh_cache(sender, instance, **kwargs):
refresh_all_orgs_user_amount_cache(instance)
model_cache_field_mapper = {
Node: ['nodes_amount'],
Domain: ['domains_amount'],
UserGroup: ['groups_amount'],
Account: ['accounts_amount'],
RoleBinding: ['users_amount', 'new_users_amount_this_week'],
Asset: ['assets_amount', 'new_assets_amount_this_week'],
AssetPermission: ['asset_perms_amount'],
}
class OrgResourceStatisticsRefreshUtil:
@staticmethod
@merge_delay_run(ttl=5)
def refresh_org_fields(org_fields=()):
for org, cache_field_name in org_fields:
OrgResourceStatisticsCache(org).expire(*cache_field_name)
OrgResourceStatisticsCache(Organization.root()).expire(*cache_field_name)
@classmethod
def refresh_if_need(cls, instance):
cache_field_name = model_cache_field_mapper.get(type(instance))
if not cache_field_name:
return
org = getattr(instance, 'org', None)
cls.refresh_org_fields(((org, cache_field_name),))
@receiver(post_save)
def on_post_save_refresh_org_resource_statistics_cache(sender, instance, created, **kwargs):
if created:
OrgResourceStatisticsRefreshUtil.refresh_if_need(instance)
@receiver(post_delete)
def on_post_delete_refresh_org_resource_statistics_cache(sender, instance, **kwargs):
OrgResourceStatisticsRefreshUtil.refresh_if_need(instance)
def _refresh_session_org_resource_statistics_cache(instance: Session):
cache_field_name = [
'total_count_online_users', 'total_count_online_sessions',
'total_count_today_active_assets','total_count_today_failed_sessions'
]
org_cache = OrgResourceStatisticsCache(instance.org)
org_cache.expire(*cache_field_name)
OrgResourceStatisticsCache(Organization.root()).expire(*cache_field_name)
@receiver(pre_save, sender=Session)
def on_session_pre_save(sender, instance: Session, **kwargs):
old = Session.objects.filter(id=instance.id).values_list('is_finished', flat=True)
if old:
instance._signal_old_is_finished = old[0]
else:
instance._signal_old_is_finished = None
@receiver(post_save, sender=Session)
def on_session_changed_refresh_org_resource_statistics_cache(sender, instance, created, **kwargs):
if created or instance.is_finished != instance._signal_old_is_finished:
_refresh_session_org_resource_statistics_cache(instance)
@receiver(post_delete, sender=Session)
def on_session_deleted_refresh_org_resource_statistics_cache(sender, instance, **kwargs):
_refresh_session_org_resource_statistics_cache(instance)