refactor: 适配新的 default 组织

apps/orgs/migrations/0010_auto_20210219_1241.py

@@ -1,6 +1,8 @@
 # Generated by Django 3.1 on 2021-02-19 04:41
 
 import time
+import sys
+
 from django.db import migrations
 
 
@@ -38,12 +40,33 @@ def migrate_default_org_id(apps, schema_editor):
         for model_name in models_name:
             t_start = time.time()
             print("Migrate model org id: {}".format(model_name), end='')
+            sys.stdout.flush()
+
             model_cls = apps.get_model(app, model_name)
             model_cls.objects.filter(org_id='').update(org_id=default_id)
             interval = round((time.time() - t_start) * 1000, 2)
             print(" done, use {} ms".format(interval))
 
 
+def add_all_user_to_default_org(apps, schema_editor):
+    User = apps.get_model('users', 'User')
+    Organization = apps.get_model('orgs', 'Organization')
+
+    users = User.objects.all()
+    default_org = Organization.objects.get(id=default_id)
+
+    t_start = time.time()
+    count = users.count()
+    print(f'{count} users to add')
+
+    batch_size = 1000
+    for i in range(0, count, batch_size):
+        default_org.members.add(*users[i:i+batch_size])
+        print(f'Add {i+1}-{i+batch_size} users')
+    interval = round((time.time() - t_start) * 1000, 2)
+    print(f'done, use {interval} ms')
+
+
 class Migration(migrations.Migration):
 
     dependencies = [
@@ -52,5 +75,6 @@ class Migration(migrations.Migration):
 
     operations = [
         migrations.RunPython(add_default_org),
-        migrations.RunPython(migrate_default_org_id)
+        migrations.RunPython(migrate_default_org_id),
+        migrations.RunPython(add_all_user_to_default_org)
     ]

apps/orgs/signals_handler/__init__.py

@@ -0,0 +1,2 @@
+from . import common
+from . import cache

apps/orgs/signals_handler/cache.py

@@ -0,0 +1,78 @@
+from django.db.models.signals import m2m_changed
+from django.db.models.signals import post_save, pre_delete
+from django.dispatch import receiver
+
+from orgs.models import Organization, OrganizationMember
+from assets.models import Node
+from perms.models import (AssetPermission, ApplicationPermission)
+from users.models import UserGroup, User
+from applications.models import Application
+from assets.models import Asset, AdminUser, SystemUser, Domain, Gateway
+from common.const.signals import POST_PREFIX
+from orgs.caches import OrgResourceStatisticsCache
+
+
+def refresh_user_amount_on_user_create_or_delete(user_id):
+    orgs = Organization.objects.filter(m2m_org_members__user_id=user_id).distinct()
+    for org in orgs:
+        org_cache = OrgResourceStatisticsCache(org)
+        org_cache.expire('users_amount')
+
+
+@receiver(post_save, sender=User)
+def on_user_create_refresh_cache(sender, instance, created, **kwargs):
+    if created:
+        refresh_user_amount_on_user_create_or_delete(instance.id)
+
+
+@receiver(pre_delete, sender=User)
+def on_user_delete_refresh_cache(sender, instance, **kwargs):
+    refresh_user_amount_on_user_create_or_delete(instance.id)
+
+
+@receiver(m2m_changed, sender=OrganizationMember)
+def on_org_user_changed_refresh_cache(sender, action, instance, reverse, pk_set, **kwargs):
+    if not action.startswith(POST_PREFIX):
+        return
+
+    if reverse:
+        orgs = Organization.objects.filter(id__in=pk_set)
+    else:
+        orgs = [instance]
+
+    for org in orgs:
+        org_cache = OrgResourceStatisticsCache(org)
+        org_cache.expire('users_amount')
+
+
+class OrgResourceStatisticsRefreshUtil:
+    model_cache_field_mapper = {
+        ApplicationPermission: 'app_perms_amount',
+        AssetPermission: 'asset_perms_amount',
+        Application: 'applications_amount',
+        Gateway: 'gateways_amount',
+        Domain: 'domains_amount',
+        SystemUser: 'system_users_amount',
+        AdminUser: 'admin_users_amount',
+        Node: 'nodes_amount',
+        Asset: 'assets_amount',
+        UserGroup: 'groups_amount',
+    }
+
+    @classmethod
+    def refresh_if_need(cls, instance):
+        cache_field_name = cls.model_cache_field_mapper.get(type(instance))
+        if cache_field_name:
+            org_cache = OrgResourceStatisticsCache(instance.org)
+            org_cache.expire(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(pre_delete)
+def on_pre_delete_refresh_org_resource_statistics_cache(sender, instance, **kwargs):
+    OrgResourceStatisticsRefreshUtil.refresh_if_need(instance)

apps/orgs/signals_handler/common.py

@@ -4,18 +4,15 @@ from collections import defaultdict
 from functools import partial
 
 from django.db.models.signals import m2m_changed
-from django.db.models.signals import post_save, pre_delete
+from django.db.models.signals import post_save
 from django.dispatch import receiver
 
 from orgs.utils import tmp_to_org
-from .models import Organization, OrganizationMember
-from .hands import set_current_org, Node, get_current_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 applications.models import Application
-from assets.models import Asset, AdminUser, SystemUser, Domain, Gateway
-from common.const.signals import PRE_REMOVE, POST_REMOVE, POST_PREFIX
-from .caches import OrgResourceStatisticsCache
+from common.const.signals import PRE_REMOVE, POST_REMOVE
 
 
 @receiver(post_save, sender=Organization)
@@ -111,70 +108,8 @@ def on_org_user_changed(action, instance, reverse, pk_set, **kwargs):
             _clear_users_from_org(org, leaved_users)
 
 
-# 缓存相关
-# -----------------------------------------------------
-
-def refresh_user_amount_on_user_create_or_delete(user_id):
-    orgs = Organization.objects.filter(m2m_org_members__user_id=user_id).distinct()
-    for org in orgs:
-        org_cache = OrgResourceStatisticsCache(org)
-        org_cache.expire('users_amount')
-
-
 @receiver(post_save, sender=User)
 def on_user_create_refresh_cache(sender, instance, created, **kwargs):
     if created:
-        refresh_user_amount_on_user_create_or_delete(instance.id)
-
-
-@receiver(pre_delete, sender=User)
-def on_user_delete_refresh_cache(sender, instance, **kwargs):
-    refresh_user_amount_on_user_create_or_delete(instance.id)
-
-
-@receiver(m2m_changed, sender=OrganizationMember)
-def on_org_user_changed_refresh_cache(sender, action, instance, reverse, pk_set, **kwargs):
-    if not action.startswith(POST_PREFIX):
-        return
-
-    if reverse:
-        orgs = Organization.objects.filter(id__in=pk_set)
-    else:
-        orgs = [instance]
-
-    for org in orgs:
-        org_cache = OrgResourceStatisticsCache(org)
-        org_cache.expire('users_amount')
-
-
-class OrgResourceStatisticsRefreshUtil:
-    model_cache_field_mapper = {
-        ApplicationPermission: 'app_perms_amount',
-        AssetPermission: 'asset_perms_amount',
-        Application: 'applications_amount',
-        Gateway: 'gateways_amount',
-        Domain: 'domains_amount',
-        SystemUser: 'system_users_amount',
-        AdminUser: 'admin_users_amount',
-        Node: 'nodes_amount',
-        Asset: 'assets_amount',
-        UserGroup: 'groups_amount',
-    }
-
-    @classmethod
-    def refresh_if_need(cls, instance):
-        cache_field_name = cls.model_cache_field_mapper.get(type(instance))
-        if cache_field_name:
-            org_cache = OrgResourceStatisticsCache(instance.org)
-            org_cache.expire(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(pre_delete)
-def on_pre_delete_refresh_org_resource_statistics_cache(sender, instance, **kwargs):
-    OrgResourceStatisticsRefreshUtil.refresh_if_need(instance)
+        default_org = Organization.default()
+        default_org.members.add(instance)

apps/perms/utils/asset/user_permission.py

@@ -198,9 +198,6 @@ class UserGrantedTreeRefreshController:
             builded_orgs_id = {org_id.decode() for org_id in ret[0]}
             ids = orgs_id - builded_orgs_id
             orgs = set()
-            if Organization.DEFAULT_ID in ids:
-                ids.remove(Organization.DEFAULT_ID)
-                orgs.add(Organization.default())
             orgs.update(Organization.objects.filter(id__in=ids))
             logger.info(f'Need rebuild orgs are {orgs}, builed orgs are {ret[0]}, all orgs are {orgs_id}')
             return orgs
@@ -293,7 +290,7 @@ class UserGrantedTreeRefreshController:
 
     @lazyproperty
     def orgs(self):
-        orgs = {*self.user.orgs.all().distinct(), Organization.default()}
+        orgs = {*self.user.orgs.all().distinct()}
         return orgs
 
     @timeit
@@ -302,11 +299,7 @@ class UserGrantedTreeRefreshController:
 
         with UserGrantedTreeRebuildLock(user_id=user.id):
             with tmp_to_root_org():
-                orgids = self.orgs_id.copy()
-                orgids.remove(Organization.default().id)
-                orgids.add('')  # 添加 default
-
-                UserAssetGrantedTreeNodeRelation.objects.filter(user=user).exclude(org_id__in=orgids).delete()
+                UserAssetGrantedTreeNodeRelation.objects.filter(user=user).exclude(org_id__in=self.orgs_id).delete()
                 exists = UserAssetGrantedTreeNodeRelation.objects.filter(user=user).exists()
 
             if force or not exists: