jumpserver/utils/test_run_migrations.py

69 lines
1.8 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_system_role_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')
count = 0
bulk_size = 1000
while True:
users = user_model.objects.using(db_alias) \
.only('role', 'id') \
.all()[count:count+bulk_size]
if not users:
break
role_bindings = []
start = time.time()
for user in users:
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("\tCreate role binding: {}-{} using: {:.2f}s".format(
count, count + len(users), time.time()-start
))
count += len(users)
class Migration(migrations.Migration):
dependencies = [
('rbac', '0003_auto_20211130_1037'),
]
operations = [
migrations.RunPython(migrate_system_role_binding),
]
# ================== 添加到下方 ======================
def main():
schema_editor = connection.schema_editor()
migrate_system_role_binding(apps, schema_editor)
# main()