From c73b49fe30649d36bd853a0ee6a0161d276749c8 Mon Sep 17 00:00:00 2001 From: fit2bot <68588906+fit2bot@users.noreply.github.com> Date: Mon, 12 Oct 2020 12:44:30 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E7=94=9F=E6=88=90?= =?UTF-8?q?=E5=81=87=E6=95=B0=E6=8D=AE=20(#4759)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf: 优化生成假数据 --- apps/assets/models/asset.py | 33 ------- apps/assets/models/cluster.py | 24 ----- apps/assets/models/group.py | 18 ---- apps/assets/models/node.py | 21 +--- apps/assets/models/user.py | 40 -------- apps/perms/models/asset_permission.py | 47 --------- apps/perms/utils/user_asset_permission.py | 7 ++ apps/users/models/group.py | 17 ---- apps/users/models/user.py | 25 ----- requirements/requirements.txt | 2 +- utils/generate_fake_data/__init__.py | 0 utils/generate_fake_data/generate.py | 49 ++++++++++ .../generate_fake_data/resources/__init__.py | 0 utils/generate_fake_data/resources/assets.py | 96 +++++++++++++++++++ utils/generate_fake_data/resources/base.py | 45 +++++++++ utils/generate_fake_data/resources/perms.py | 82 ++++++++++++++++ utils/generate_fake_data/resources/users.py | 62 ++++++++++++ 17 files changed, 345 insertions(+), 223 deletions(-) create mode 100644 utils/generate_fake_data/__init__.py create mode 100644 utils/generate_fake_data/generate.py create mode 100644 utils/generate_fake_data/resources/__init__.py create mode 100644 utils/generate_fake_data/resources/assets.py create mode 100644 utils/generate_fake_data/resources/base.py create mode 100644 utils/generate_fake_data/resources/perms.py create mode 100644 utils/generate_fake_data/resources/users.py diff --git a/apps/assets/models/asset.py b/apps/assets/models/asset.py index 65414e86d..900a87060 100644 --- a/apps/assets/models/asset.py +++ b/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 diff --git a/apps/assets/models/cluster.py b/apps/assets/models/cluster.py index 84f8f2374..6c0692ab9 100644 --- a/apps/assets/models/cluster.py +++ b/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 diff --git a/apps/assets/models/group.py b/apps/assets/models/group.py index b9bf16f18..ee5472622 100644 --- a/apps/assets/models/group.py +++ b/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 diff --git a/apps/assets/models/node.py b/apps/assets/models/node.py index 5b4dd7259..7c1c81cde 100644 --- a/apps/assets/models/node.py +++ b/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)) diff --git a/apps/assets/models/user.py b/apps/assets/models/user.py index 3fb862760..087a05bef 100644 --- a/apps/assets/models/user.py +++ b/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 diff --git a/apps/perms/models/asset_permission.py b/apps/perms/models/asset_permission.py index 5611f6ad5..9fc00f4a7 100644 --- a/apps/perms/models/asset_permission.py +++ b/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): diff --git a/apps/perms/utils/user_asset_permission.py b/apps/perms/utils/user_asset_permission.py index 01a558ae2..3fb1ed9d0 100644 --- a/apps/perms/utils/user_asset_permission.py +++ b/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，推算出来的也算 diff --git a/apps/users/models/group.py b/apps/users/models/group.py index 0ae636f76..d06fce989 100644 --- a/apps/users/models/group.py +++ b/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 diff --git a/apps/users/models/user.py b/apps/users/models/user.py index 22adab238..81f0f593b 100644 --- a/apps/users/models/user.py +++ b/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() diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 018632ed3..0b71e5394 100644 --- a/requirements/requirements.txt +++ b/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 diff --git a/utils/generate_fake_data/__init__.py b/utils/generate_fake_data/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/utils/generate_fake_data/generate.py b/utils/generate_fake_data/generate.py new file mode 100644 index 000000000..f504c8e55 --- /dev/null +++ b/utils/generate_fake_data/generate.py @@ -0,0 +1,49 @@ +#!/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() diff --git a/utils/generate_fake_data/resources/__init__.py b/utils/generate_fake_data/resources/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/utils/generate_fake_data/resources/assets.py b/utils/generate_fake_data/resources/assets.py new file mode 100644 index 000000000..a0edc9f08 --- /dev/null +++ b/utils/generate_fake_data/resources/assets.py @@ -0,0 +1,96 @@ +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() diff --git a/utils/generate_fake_data/resources/base.py b/utils/generate_fake_data/resources/base.py new file mode 100644 index 000000000..ba45f602b --- /dev/null +++ b/utils/generate_fake_data/resources/base.py @@ -0,0 +1,45 @@ +#!/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 diff --git a/utils/generate_fake_data/resources/perms.py b/utils/generate_fake_data/resources/perms.py new file mode 100644 index 000000000..e76e3618f --- /dev/null +++ b/utils/generate_fake_data/resources/perms.py @@ -0,0 +1,82 @@ +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) diff --git a/utils/generate_fake_data/resources/users.py b/utils/generate_fake_data/resources/users.py new file mode 100644 index 000000000..05332e4e2 --- /dev/null +++ b/utils/generate_fake_data/resources/users.py @@ -0,0 +1,62 @@ +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)