jumpserver/apps/assets/migrations/0102_auto_20220711_1413.py

# Generated by Django 3.2.12 on 2022-07-11 06:13

import time
from django.db import migrations


def migrate_accounts(apps, schema_editor):
    auth_book_model = apps.get_model('assets', 'AuthBook')
    account_model = apps.get_model('assets', 'Account')

    count = 0
    bulk_size = 1000
    print("\nStart migrate accounts")
    while True:
        start = time.time()
        auth_books = auth_book_model.objects \
            .prefetch_related('systemuser') \
            .all()[count:count+bulk_size]
        count += len(auth_books)
        if not auth_books:
            break

        accounts = []
        # auth book 和 account 相同的属性
        same_attrs = [
            'id', 'comment', 'date_created', 'date_updated',
            'created_by', 'asset_id', 'org_id',
        ]
        # 认证的属性，可能是 authbook 的，可能是 systemuser 的
        auth_attrs = ['username', 'password', 'private_key', 'public_key']

        for auth_book in auth_books:
            values = {attr: getattr(auth_book, attr) for attr in same_attrs}
            values['protocol'] = 'ssh'
            values['version'] = 1

            system_user = auth_book.systemuser
            if auth_book.systemuser:
                values.update({attr: getattr(system_user, attr) for attr in auth_attrs})
                values['protocol'] = system_user.protocol
                values['created_by'] = str(system_user.id)
                values['type'] = system_user.type

            auth_book_auth = {attr: getattr(auth_book, attr) for attr in auth_attrs}
            auth_book_auth = {attr: value for attr, value in auth_book_auth.items() if value}
            values.update(auth_book_auth)

            account = account_model(**values)
            accounts.append(account)

        account_model.objects.bulk_create(accounts, ignore_conflicts=True)
        print("Create accounts: {}-{} using: {:.2f}s".format(
            count - bulk_size, count, time.time()-start
        ))


class Migration(migrations.Migration):

    dependencies = [
        ('assets', '0101_auto_20220711_1409'),
    ]

    operations = [
        migrations.RunPython(migrate_accounts)
    ]