perf: 优化生成假数据 (#4759)

* perf: 优化生成假数据

apps/assets/models/asset.py

@@ -355,36 +355,3 @@ class Asset(ProtocolsMixin, NodesRelationMixin, OrgModelMixin):
     class Meta:
         unique_together = [('org_id', 'hostname')]
         verbose_name = _("Asset")
-
-    @classmethod
-    def generate_fake(cls, count=100):
-        from .user import AdminUser, SystemUser
-        from random import seed, choice
-        from django.db import IntegrityError
-        from .node import Node
-        from orgs.utils import get_current_org
-        from orgs.models import Organization
-        org = get_current_org()
-        if not org or not org.is_real():
-            Organization.default().change_to()
-
-        nodes = list(Node.objects.all())
-        seed()
-        for i in range(count):
-            ip = [str(i) for i in random.sample(range(255), 4)]
-            asset = cls(ip='.'.join(ip),
-                        hostname='.'.join(ip),
-                        admin_user=choice(AdminUser.objects.all()),
-                        created_by='Fake')
-            try:
-                asset.save()
-                asset.protocols = 'ssh/22'
-                if nodes and len(nodes) > 3:
-                    _nodes = random.sample(nodes, 3)
-                else:
-                    _nodes = [Node.default_node()]
-                asset.nodes.set(_nodes)
-                logger.debug('Generate fake asset : %s' % asset.ip)
-            except IntegrityError:
-                print('Error continue')
-                continue

apps/assets/models/cluster.py

@@ -38,27 +38,3 @@ class Cluster(models.Model):
     class Meta:
         ordering = ['name']
         verbose_name = _("Cluster")
-
-    @classmethod
-    def generate_fake(cls, count=5):
-        from random import seed, choice
-        import forgery_py
-        from django.db import IntegrityError
-
-        seed()
-        for i in range(count):
-            cluster = cls(name=forgery_py.name.full_name(),
-                          bandwidth='200M',
-                          contact=forgery_py.name.full_name(),
-                          phone=forgery_py.address.phone(),
-                          address=forgery_py.address.city() + forgery_py.address.street_address(),
-                          # operator=choice(['北京联通', '北京电信', 'BGP全网通']),
-                          operator=choice([_('Beijing unicom'), _('Beijing telecom'), _('BGP full netcom')]),
-                          comment=forgery_py.lorem_ipsum.sentence(),
-                          created_by='Fake')
-            try:
-                cluster.save()
-                logger.debug('Generate fake asset group: %s' % cluster.name)
-            except IntegrityError:
-                print('Error continue')
-                continue

apps/assets/models/group.py

@@ -33,21 +33,3 @@ class AssetGroup(models.Model):
     def initial(cls):
         asset_group = cls(name=_('Default'), comment=_('Default asset group'))
         asset_group.save()
-
-    @classmethod
-    def generate_fake(cls, count=100):
-        from random import seed
-        import forgery_py
-        from django.db import IntegrityError
-
-        seed()
-        for i in range(count):
-            group = cls(name=forgery_py.name.full_name(),
-                        comment=forgery_py.lorem_ipsum.sentence(),
-                        created_by='Fake')
-            try:
-                group.save()
-                logger.debug('Generate fake asset group: %s' % group.name)
-            except IntegrityError:
-                print('Error continue')
-                continue

apps/assets/models/node.py

@@ -97,9 +97,11 @@ class FamilyMixin:
     def all_children(self):
         return self.get_all_children(with_self=False)
 
-    def create_child(self, value, _id=None):
+    def create_child(self, value=None, _id=None):
         with atomic(savepoint=False):
             child_key = self.get_next_child_key()
+            if value is None:
+                value = child_key
             child = self.__class__.objects.create(
                 id=_id, key=child_key, value=value, parent_key=self.key,
             )
@@ -456,20 +458,3 @@ class Node(OrgModelMixin, SomeNodesMixin, FamilyMixin, NodeAssetsMixin):
         if self.has_children_or_has_assets():
             return
         return super().delete(using=using, keep_parents=keep_parents)
-
-    @classmethod
-    def generate_fake(cls, count=100):
-        import random
-        org = get_current_org()
-        if not org or not org.is_real():
-            Organization.default().change_to()
-        nodes = list(cls.objects.all())
-        if count > 100:
-            length = 100
-        else:
-            length = count
-
-        for i in range(length):
-            node = random.choice(nodes)
-            child = node.create_child('Node {}'.format(i))
-            print("{}. {}".format(i, child))

apps/assets/models/user.py

@@ -65,26 +65,6 @@ class AdminUser(BaseUser):
         unique_together = [('name', 'org_id')]
         verbose_name = _("Admin user")
 
-    @classmethod
-    def generate_fake(cls, count=10):
-        from random import seed
-        import forgery_py
-        from django.db import IntegrityError
-
-        seed()
-        for i in range(count):
-            obj = cls(name=forgery_py.name.full_name(),
-                      username=forgery_py.internet.user_name(),
-                      password=forgery_py.lorem_ipsum.word(),
-                      comment=forgery_py.lorem_ipsum.sentence(),
-                      created_by='Fake')
-            try:
-                obj.save()
-                logger.debug('Generate fake asset group: %s' % obj.name)
-            except IntegrityError:
-                print('Error continue')
-                continue
-
 
 class SystemUser(BaseUser):
     PROTOCOL_SSH = 'ssh'
@@ -199,23 +179,3 @@ class SystemUser(BaseUser):
         ordering = ['name']
         unique_together = [('name', 'org_id')]
         verbose_name = _("System user")
-
-    @classmethod
-    def generate_fake(cls, count=10):
-        from random import seed
-        import forgery_py
-        from django.db import IntegrityError
-
-        seed()
-        for i in range(count):
-            obj = cls(name=forgery_py.name.full_name(),
-                      username=forgery_py.internet.user_name(),
-                      password=forgery_py.lorem_ipsum.word(),
-                      comment=forgery_py.lorem_ipsum.sentence(),
-                      created_by='Fake')
-            try:
-                obj.save()
-                logger.debug('Generate fake asset group: %s' % obj.name)
-            except IntegrityError:
-                print('Error continue')
-                continue

apps/perms/models/asset_permission.py

@@ -136,53 +136,6 @@ class AssetPermission(BasePermission):
         assets = Asset.objects.filter(id__in=assets_ids)
         return assets
 
-    @classmethod
-    def generate_fake(cls, count=100):
-        from ..hands import User, Node, SystemUser
-        import random
-
-        org = get_current_org()
-        if not org or not org.is_real():
-            Organization.default().change_to()
-
-        nodes = list(Node.objects.all())
-        assets = list(Asset.objects.all())
-        system_users = list(SystemUser.objects.all())
-        users = User.objects.filter(username='admin')
-
-        for i in range(count):
-            name = "fake_perm_to_admin_{}".format(str(uuid.uuid4())[:6])
-            perm = cls(name=name)
-            try:
-                perm.save()
-                perm.users.set(users)
-                if system_users and len(system_users) > 3:
-                    _system_users = random.sample(system_users, 3)
-                elif system_users:
-                    _system_users = [system_users[0]]
-                else:
-                    _system_users = []
-                perm.system_users.set(_system_users)
-
-                if nodes and len(nodes) > 3:
-                    _nodes = random.sample(nodes, 3)
-                else:
-                    _nodes = [Node.default_node()]
-                perm.nodes.set(_nodes)
-
-                if assets and len(assets) > 3:
-                    _assets = random.sample(assets, 3)
-                elif assets:
-                    _assets = [assets[0]]
-                else:
-                    _assets = []
-                perm.assets.set(_assets)
-
-                logger.debug('Generate fake perm: %s' % perm.name)
-
-            except Exception as e:
-                print('Error continue')
-                continue
 
 
 class UserGrantedMappingNode(FamilyMixin, models.JMSBaseModel):

apps/perms/utils/user_asset_permission.py

@@ -256,6 +256,13 @@ def rebuild_user_mapping_nodes(user):
     logger.info(f'>>> {dt_formater(now())} end rebuild {user} mapping nodes')
 
 
+def rebuild_all_user_mapping_nodes():
+    from users.models import User
+    users = User.objects.all()
+    for user in users:
+        rebuild_user_mapping_nodes(user)
+
+
 def get_user_granted_nodes_list_via_mapping_node(user):
     """
     这里的 granted nodes, 是整棵树需要的node，推算出来的也算

apps/users/models/group.py

@@ -39,20 +39,3 @@ class UserGroup(OrgModelMixin):
         else:
             group = default_group[0]
         return group
-
-    @classmethod
-    def generate_fake(cls, count=100):
-        from random import seed, choice
-        import forgery_py
-        from . import User
-
-        seed()
-        for i in range(count):
-            group = cls(name=forgery_py.name.full_name(),
-                        comment=forgery_py.lorem_ipsum.sentence(),
-                        created_by=choice(User.objects.all()).username)
-            try:
-                group.save()
-            except IntegrityError:
-                print('Error continue')
-                continue

apps/users/models/user.py

@@ -670,28 +670,3 @@ class User(AuthMixin, TokenMixin, RoleMixin, MFAMixin, AbstractUser):
         if self.email and self.source == self.SOURCE_LOCAL:
             return True
         return False
-
-    @classmethod
-    def generate_fake(cls, count=100):
-        from random import seed, choice
-        import forgery_py
-        from django.db import IntegrityError
-        from .group import UserGroup
-
-        seed()
-        for i in range(count):
-            user = cls(username=forgery_py.internet.user_name(True),
-                       email=forgery_py.internet.email_address(),
-                       name=forgery_py.name.full_name(),
-                       password=make_password(forgery_py.lorem_ipsum.word()),
-                       role=choice(list(dict(User.ROLE.choices).keys())),
-                       wechat=forgery_py.internet.user_name(True),
-                       comment=forgery_py.lorem_ipsum.sentence(),
-                       created_by=choice(cls.objects.all()).username)
-            try:
-                user.save()
-            except IntegrityError:
-                print('Duplicate Error, continue ...')
-                continue
-            user.groups.add(choice(UserGroup.objects.all()))
-            user.save()

requirements/requirements.txt

@@ -33,7 +33,7 @@ enum-compat==0.0.2
 ephem==3.7.6.0
 eventlet==0.24.1
 future==0.16.0
-ForgeryPy==0.1
+ForgeryPy3==0.3.2
 greenlet==0.4.14
 gunicorn==19.9.0
 idna==2.6

utils/generate_fake_data/generate.py

+#!/usr/bin/env python
+#
+import os
+import sys
+import django
+import argparse
+
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+APPS_DIR = os.path.join(BASE_DIR, 'apps')
+sys.path.insert(0, APPS_DIR)
+
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "jumpserver.settings")
+django.setup()
+
+from resources.assets import AssetsGenerator, NodesGenerator, SystemUsersGenerator, AdminUsersGenerator
+from resources.users import UserGroupGenerator, UserGenerator
+from resources.perms import AssetPermissionGenerator
+
+
+resource_generator_mapper = {
+    'asset': AssetsGenerator,
+    'node': NodesGenerator,
+    'system_user': SystemUsersGenerator,
+    'admin_user': AdminUsersGenerator,
+    'user': UserGenerator,
+    'user_group': UserGroupGenerator,
+    'asset_permission': AssetPermissionGenerator
+}
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Generate fake data')
+    parser.add_argument(
+        'resource', type=str,
+        choices=resource_generator_mapper.keys(),
+        help="resource to generate"
+    )
+    parser.add_argument('-c', '--count', type=int, default=100)
+    parser.add_argument('-b', '--batch_size', type=int, default=100)
+    parser.add_argument('-o', '--org', type=str, default='')
+    args = parser.parse_args()
+    resource, count, batch_size, org_id = args.resource, args.count, args.batch_size, args.org
+    generator_cls = resource_generator_mapper[resource]
+    generator = generator_cls(org_id=org_id, batch_size=batch_size)
+    generator.generate(count)
+
+
+if __name__ == '__main__':
+    main()

utils/generate_fake_data/resources/assets.py

+from random import choice
+import random
+import forgery_py
+
+from .base import FakeDataGenerator
+
+from assets.models import *
+from assets.utils import check_node_assets_amount
+
+
+class AdminUsersGenerator(FakeDataGenerator):
+    resource = 'admin_user'
+
+    def do_generate(self, batch, batch_size):
+        admin_users = []
+        for i in batch:
+            username = forgery_py.internet.user_name(True)
+            password = forgery_py.basic.password()
+            admin_users.append(AdminUser(
+                name=username.title(),
+                username=username,
+                password=password,
+                org_id=self.org.id,
+                created_by='Fake',
+            ))
+        AdminUser.objects.bulk_create(admin_users, ignore_conflicts=True)
+
+
+class SystemUsersGenerator(FakeDataGenerator):
+    def do_generate(self, batch, batch_size):
+        system_users = []
+        protocols = list(dict(SystemUser.PROTOCOL_CHOICES).keys())
+        for i in batch:
+            username = forgery_py.internet.user_name(True)
+            protocol = random.choice(protocols)
+            name = username.title()
+            name = f'{name}-{protocol}'
+            system_users.append(SystemUser(
+                name=name,
+                username=username,
+                password=forgery_py.basic.password(),
+                protocol=protocol,
+                org_id=self.org.id,
+                created_by='Fake',
+            ))
+        SystemUser.objects.bulk_create(system_users, ignore_conflicts=True)
+
+
+class NodesGenerator(FakeDataGenerator):
+    resource = 'node'
+
+    def do_generate(self, batch, batch_size):
+        nodes_to_generate_children = list(Node.objects.all())
+        for i in batch:
+            parent = random.choice(nodes_to_generate_children)
+            parent.create_child()
+
+
+class AssetsGenerator(FakeDataGenerator):
+    resource = 'asset'
+    admin_users_id: list
+    nodes_id: list
+
+    def pre_generate(self):
+        self.admin_users_id = list(AdminUser.objects.all().values_list('id', flat=True))
+        self.nodes_id = list(Node.objects.all().values_list('id', flat=True))
+
+    def set_assets_nodes(self, assets):
+        assets_id = [asset.id for asset in assets]
+        objs = []
+        for asset_id in assets_id:
+            nodes_id_add_to = random.sample(self.nodes_id, 3)
+            objs_add = [Asset.nodes.through(asset_id=asset_id, node_id=nid) for nid in nodes_id_add_to]
+            objs.extend(objs_add)
+        Asset.nodes.through.objects.bulk_create(objs, ignore_conflicts=True)
+
+    def do_generate(self, batch, batch_size):
+        assets = []
+
+        for i in batch:
+            ip = forgery_py.internet.ip_v4()
+            hostname = forgery_py.email.address().replace('@', '.')
+            hostname = f'{hostname}-{ip}'
+            data = dict(
+                ip=ip,
+                hostname=hostname,
+                admin_user_id=choice(self.admin_users_id),
+                created_by='Fake',
+                org_id=self.org.id
+            )
+            assets.append(Asset(**data))
+        creates = Asset.objects.bulk_create(assets, ignore_conflicts=True)
+        self.set_assets_nodes(creates)
+
+    def after_generate(self):
+        check_node_assets_amount()

utils/generate_fake_data/resources/base.py

+#!/usr/bin/python
+
+from random import seed
+from itertools import islice
+
+from orgs.models import Organization
+
+
+class FakeDataGenerator:
+    resource = 'Fake'
+
+    def __init__(self, batch_size=100, org_id=None):
+        self.batch_size = batch_size
+        self.org = self.switch_org(org_id)
+        seed()
+
+    def switch_org(self, org_id):
+        o = Organization.get_instance(org_id, default=True)
+        if o:
+            o.change_to()
+        print('Current org is: {}'.format(o))
+        return o
+
+    def do_generate(self, batch, batch_size):
+        raise NotImplementedError
+
+    def pre_generate(self):
+        pass
+
+    def after_generate(self):
+        pass
+
+    def generate(self, count=100):
+        self.pre_generate()
+        counter = iter(range(count))
+        created = 0
+        while True:
+            batch = list(islice(counter, self.batch_size))
+            if not batch:
+                break
+            self.do_generate(batch, self.batch_size)
+            from_size = created
+            created += len(batch)
+            print('Generate %s: %s-%s' % (self.resource, from_size, created))
+        self.after_generate()
\ No newline at end of file

utils/generate_fake_data/resources/perms.py

+from random import choice, sample
+import forgery_py
+
+from .base import FakeDataGenerator
+
+from users.models import *
+from assets.models import *
+from perms.models import *
+
+
+class AssetPermissionGenerator(FakeDataGenerator):
+    resource = 'asset_permission'
+    users_id: list
+    user_groups_id: list
+    assets_id: list
+    nodes_id: list
+    system_users_id: list
+
+    def pre_generate(self):
+        self.nodes_id = list(Node.objects.all().values_list('id', flat=True))
+        self.assets_id = list(Asset.objects.all().values_list('id', flat=True))
+        self.system_users_id = list(SystemUser.objects.all().values_list('id', flat=True))
+        self.users_id = list(User.objects.all().values_list('id', flat=True))
+        self.user_groups_id = list(UserGroup.objects.all().values_list('id', flat=True))
+
+    def set_users(self, perms):
+        through = AssetPermission.users.through
+        choices = self.users_id
+        relation_name = 'user_id'
+        self.set_relations(perms, through, relation_name, choices)
+
+    def set_user_groups(self, perms):
+        through = AssetPermission.user_groups.through
+        choices = self.user_groups_id
+        relation_name = 'usergroup_id'
+        self.set_relations(perms, through, relation_name, choices)
+
+    def set_assets(self, perms):
+        through = AssetPermission.assets.through
+        choices = self.assets_id
+        relation_name = 'asset_id'
+        self.set_relations(perms, through, relation_name, choices)
+
+    def set_nodes(self, perms):
+        through = AssetPermission.nodes.through
+        choices = self.nodes_id
+        relation_name = 'node_id'
+        self.set_relations(perms, through, relation_name, choices)
+
+    def set_system_users(self, perms):
+        through = AssetPermission.system_users.through
+        choices = self.system_users_id
+        relation_name = 'systemuser_id'
+        self.set_relations(perms, through, relation_name, choices)
+
+    def set_relations(self, perms, through, relation_name, choices, choice_count=None):
+        relations = []
+
+        for perm in perms:
+            if choice_count is None:
+                choice_count = choice(range(8))
+            resources_id = sample(choices, choice_count)
+            for rid in resources_id:
+                data = {'assetpermission_id': perm.id}
+                data[relation_name] = rid
+                relations.append(through(**data))
+        through.objects.bulk_create(relations, ignore_conflicts=True)
+
+    def do_generate(self, batch, batch_size):
+        perms = []
+
+        for i in batch:
+            name = forgery_py.basic.text()
+            name = f'AssetPermission: {name}'
+            perm = AssetPermission(name=name, org_id=self.org.id)
+            perms.append(perm)
+        created = AssetPermission.objects.bulk_create(perms, ignore_conflicts=True)
+        self.set_users(created)
+        self.set_user_groups(created)
+        self.set_assets(created)
+        self.set_nodes(created)
+        self.set_system_users(created)

utils/generate_fake_data/resources/users.py

+from random import choice, sample
+import forgery_py
+
+from .base import FakeDataGenerator
+
+from users.models import *
+from orgs.models import OrganizationMember
+
+
+class UserGroupGenerator(FakeDataGenerator):
+    resource = 'usergroup'
+
+    def do_generate(self, batch, batch_size):
+        groups = []
+        for i in batch:
+            group_name = forgery_py.name.job_title()
+            groups.append(UserGroup(name=group_name, org_id=self.org.id))
+        UserGroup.objects.bulk_create(groups, ignore_conflicts=True)
+
+
+class UserGenerator(FakeDataGenerator):
+    resource = 'user'
+    roles: list
+    groups_id: list
+
+    def pre_generate(self):
+        self.roles = list(dict(User.ROLE.choices).keys())
+        self.groups_id = list(UserGroup.objects.all().values_list('id', flat=True))
+
+    def set_org(self, users):
+        relations = []
+        for u in users:
+            relations.append(OrganizationMember(
+                org_id=self.org.id,
+                user_id=u.id,
+            ))
+        OrganizationMember.objects.bulk_create(relations, ignore_conflicts=True)
+
+    def set_groups(self, users):
+        relations = []
+        for i in users:
+            groups_to_join = sample(self.groups_id, 3)
+            _relations = [User.groups.through(user_id=i.id, usergroup_id=gid) for gid in groups_to_join]
+            relations.extend(_relations)
+        User.groups.through.objects.bulk_create(relations, ignore_conflicts=True)
+
+    def do_generate(self, batch, batch_size):
+        users = []
+        for i in batch:
+            username = forgery_py.internet.user_name(True)
+            email = forgery_py.internet.email_address()
+            u = User(
+                username=username,
+                email=email,
+                name=username.title(),
+                role=choice(self.roles),
+                created_by='Faker'
+            )
+            users.append(u)
+        users = User.objects.bulk_create(users, ignore_conflicts=True)
+        self.set_org(users)
+        self.set_groups(users)