mirror of https://github.com/jumpserver/jumpserver
61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
|
# Generated by Django 3.1.13 on 2021-12-01 11:01
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import django
|
||
|
|
import time
|
||
|
|
|
||
|
|
app_path = '***** Change me *******'
|
||
|
|
sys.path.insert(0, app_path)
|
||
|
|
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "jumpserver.settings")
|
||
|
|
django.setup()
|
||
|
|
|
||
|
|
from django.apps import apps
|
||
|
|
from django.db import connection
|
||
|
|
|
||
|
|
# ========================== 添加到需要测试的 migrations 上方 ==========================
|
||
|
|
|
||
|
|
|
||
|
|
from django.db import migrations
|
||
|
|
|
||
|
|
from rbac.builtin import BuiltinRole
|
||
|
|
|
||
|
|
|
||
|
|
def migrate_some_binding(apps, schema_editor):
|
||
|
|
db_alias = schema_editor.connection.alias
|
||
|
|
user_model = apps.get_model('users', 'User')
|
||
|
|
role_binding_model = apps.get_model('rbac', 'SystemRoleBinding')
|
||
|
|
users = user_model.objects.using(db_alias).all()
|
||
|
|
|
||
|
|
grouped_users = [users[i:i+1000] for i in range(0, len(users), 1000)]
|
||
|
|
for i, group in enumerate(grouped_users):
|
||
|
|
role_bindings = []
|
||
|
|
start = time.time()
|
||
|
|
for user in group:
|
||
|
|
role = BuiltinRole.get_system_role_by_old_name(user.role)
|
||
|
|
role_binding = role_binding_model(scope='system', user_id=user.id, role_id=role.id)
|
||
|
|
role_bindings.append(role_binding)
|
||
|
|
role_binding_model.objects.bulk_create(role_bindings, ignore_conflicts=True)
|
||
|
|
print("Create role binding: {}-{} using: {:.2f}s".format(
|
||
|
|
i*1000, i*1000 + len(group), time.time()-start
|
||
|
|
))
|
||
|
|
|
||
|
|
|
||
|
|
class Migration(migrations.Migration):
|
||
|
|
|
||
|
|
dependencies = [
|
||
|
|
('rbac', '0003_auto_20211130_1037'),
|
||
|
|
]
|
||
|
|
|
||
|
|
operations = [
|
||
|
|
migrations.RunPython(migrate_some_binding),
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
# ================== 添加到下方 ======================
|
||
|
|
def main():
|
||
|
|
schema_editor = connection.schema_editor()
|
||
|
|
migrate_some_binding(apps, schema_editor)
|
||
|
|
|
||
|
|
|
||
|
|
# main()
|