# -*- coding: utf-8 -*- # import threading from collections import defaultdict from functools import partial from django.dispatch import receiver from django.utils.functional import LazyObject from common.db.utils import close_old_connections from django.db.models.signals import m2m_changed from django.db.models.signals import post_save, post_delete, pre_delete from orgs.utils import tmp_to_org from orgs.models import Organization, OrganizationMember from orgs.hands import set_current_org, Node, get_current_org from perms.models import (AssetPermission, ApplicationPermission) from users.models import UserGroup, User from common.const.signals import PRE_REMOVE, POST_REMOVE from common.decorator import on_transaction_commit from common.signals import django_ready from common.utils import get_logger from common.utils.connection import RedisPubSub from assets.models import CommandFilterRule logger = get_logger(__file__) def get_orgs_mapping_for_memory_pub_sub(): return RedisPubSub('fm.orgs_mapping') class OrgsMappingForMemoryPubSub(LazyObject): def _setup(self): self._wrapped = get_orgs_mapping_for_memory_pub_sub() orgs_mapping_for_memory_pub_sub = OrgsMappingForMemoryPubSub() def expire_orgs_mapping_for_memory(org_id): orgs_mapping_for_memory_pub_sub.publish(str(org_id)) @receiver(django_ready) def subscribe_orgs_mapping_expire(sender, **kwargs): logger.debug("Start subscribe for expire orgs mapping from memory") def keep_subscribe_org_mapping(): while True: try: subscribe = orgs_mapping_for_memory_pub_sub.subscribe() msgs = subscribe.listen() # 开始之前关闭连接,因为server端可能关闭了连接,而 client 还在 CONN_MAX_AGE 中 close_old_connections() for message in msgs: if message['type'] != 'message': continue if message['data'] == b'error': raise ValueError Organization.expire_orgs_mapping() logger.debug('Expire orgs mapping: ' + str(message['data'])) except Exception as e: logger.exception(f'subscribe_orgs_mapping_expire: {e}') Organization.expire_orgs_mapping() finally: # 结束收关闭连接 close_old_connections() t = threading.Thread(target=keep_subscribe_org_mapping) t.daemon = True t.start() @receiver(post_save, sender=Organization) def on_org_create_or_update(sender, instance, created=False, **kwargs): # 必须放到最开始, 因为下面调用Node.save方法时会获取当前组织的org_id(即instance.org_id), 如果不过期会找不到 expire_orgs_mapping_for_memory(instance.id) old_org = get_current_org() set_current_org(instance) node_root = Node.org_root() if node_root.value != instance.name: node_root.value = instance.name node_root.save() set_current_org(old_org) @receiver(pre_delete, sender=Organization) def on_org_delete(sender, instance, **kwargs): expire_orgs_mapping_for_memory(instance.id) @receiver(pre_delete, sender=Organization) def on_org_delete(sender, instance, **kwargs): # 删除该组织下所有 节点 with tmp_to_org(instance): root_node = Node.org_root() if root_node: root_node.delete() def _remove_users(model, users, org, user_field_name='users'): with tmp_to_org(org): if not isinstance(users, (tuple, list, set)): users = (users, ) user_field = getattr(model, user_field_name) m2m_model = user_field.through reverse = user_field.reverse if reverse: m2m_field_name = user_field.field.m2m_reverse_field_name() else: m2m_field_name = user_field.field.m2m_field_name() relations = m2m_model.objects.filter(**{ 'user__in': users, f'{m2m_field_name}__org_id': org.id }) object_id_user_ids_map = defaultdict(set) m2m_field_attr_name = f'{m2m_field_name}_id' for relation in relations: object_id = getattr(relation, m2m_field_attr_name) object_id_user_ids_map[object_id].add(relation.user_id) objects = model.objects.filter(id__in=object_id_user_ids_map.keys()) send_m2m_change_signal = partial( m2m_changed.send, sender=m2m_model, reverse=reverse, model=User, using=model.objects.db ) for obj in objects: send_m2m_change_signal( instance=obj, pk_set=object_id_user_ids_map[obj.id], action=PRE_REMOVE ) relations.delete() for obj in objects: send_m2m_change_signal( instance=obj, pk_set=object_id_user_ids_map[obj.id], action=POST_REMOVE ) def _clear_users_from_org(org, users): """ 清理用户在该组织下的相关数据 """ if not users: return models = (AssetPermission, ApplicationPermission, UserGroup) for m in models: _remove_users(m, users, org) _remove_users(CommandFilterRule, users, org, user_field_name='reviewers') @receiver(m2m_changed, sender=OrganizationMember) def on_org_user_changed(action, instance, reverse, pk_set, **kwargs): if action == 'post_remove': if reverse: user = instance org_pk_set = pk_set orgs = Organization.objects.filter(id__in=org_pk_set) for org in orgs: if not org.members.filter(id=user.id).exists(): _clear_users_from_org(org, user) else: org = instance user_pk_set = pk_set leaved_users = set(pk_set) - set(org.members.filter(id__in=user_pk_set).values_list('id', flat=True)) _clear_users_from_org(org, leaved_users) @receiver(post_save, sender=User) @on_transaction_commit def on_user_created_set_default_org(sender, instance, created, **kwargs): if not created: return if instance.orgs.count() > 0: return Organization.default().members.add(instance)