jumpserver/apps/orgs/migrations/0010_auto_20210219_1241.py

# Generated by Django 3.1 on 2021-02-19 04:41

import time
import sys

from django.db import migrations


default_id = '00000000-0000-0000-0000-000000000001'


def add_default_org(apps, schema_editor):
    org_cls = apps.get_model('orgs', 'Organization')
    defaults = {'name': 'DEFAULT', 'id': default_id}
    org_cls.objects.get_or_create(defaults=defaults, id=default_id)


def migrate_default_org_id(apps, schema_editor):
    org_app_models = [
        ('applications', ['Application']),
        ('assets', [
            'AdminUser', 'Asset', 'AuthBook', 'CommandFilter',
            'CommandFilterRule', 'Domain', 'Gateway', 'GatheredUser',
            'Label', 'Node', 'SystemUser'
        ]),
        ('audits', ['FTPLog', 'OperateLog']),
        ('ops', ['AdHoc', 'AdHocExecution', 'CommandExecution', 'Task']),
        ('perms', ['ApplicationPermission', 'AssetPermission', 'UserAssetGrantedTreeNodeRelation']),
        ('terminal', ['Session', 'Command']),
        ('tickets', ['Ticket']),
        ('users', ['UserGroup']),
    ]
    print("")
    for app, models_name in org_app_models:
        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 = [
        ('orgs', '0009_auto_20201023_1628'),
        ('perms', '0018_auto_20210208_1515'),
    ]

    operations = [
        migrations.RunPython(add_default_org),
        migrations.RunPython(migrate_default_org_id),
        migrations.RunPython(add_all_user_to_default_org)
    ]