diff --git a/apps/assets/migrations/0100_auto_20220711_1413.py b/apps/assets/migrations/0100_auto_20220711_1413.py index 8bb3309e4..dd7e34394 100644 --- a/apps/assets/migrations/0100_auto_20220711_1413.py +++ b/apps/assets/migrations/0100_auto_20220711_1413.py @@ -1,6 +1,7 @@ # Generated by Django 3.2.12 on 2022-07-11 06:13 import time +import math from django.utils import timezone from itertools import groupby from django.db import migrations @@ -40,9 +41,13 @@ def migrate_asset_accounts(apps, schema_editor): if system_user: # 更新一次系统用户的认证属性 account_values.update({attr: getattr(system_user, attr, '') for attr in all_attrs}) - account_values['created_by'] = str(system_user.id) account_values['privileged'] = system_user.type == 'admin' \ or system_user.username in ['root', 'Administrator'] + if system_user.su_enabled and system_user.su_from: + created_by = f'{str(system_user.id)}::{str(system_user.su_from.username)}' + else: + created_by = str(system_user.id) + account_values['created_by'] = created_by auth_book_auth = {attr: getattr(auth_book, attr, '') for attr in all_attrs if getattr(auth_book, attr, '')} # 最终优先使用 auth_book 的认证属性 @@ -117,6 +122,70 @@ def migrate_asset_accounts(apps, schema_editor): print("\t - histories: {}".format(len(accounts_to_history))) +def update_asset_accounts_su_from(apps, schema_editor): + # Update accounts su_from + print("\n\tStart update asset accounts su_from field") + account_model = apps.get_model('accounts', 'Account') + platform_model = apps.get_model('assets', 'Platform') + asset_model = apps.get_model('assets', 'Asset') + platform_ids = list(platform_model.objects.filter(su_enabled=True).values_list('id', flat=True)) + + count = 0 + step_size = 1000 + count_account = 0 + while True: + start = time.time() + asset_ids = asset_model.objects \ + .filter(platform_id__in=platform_ids) \ + .values_list('id', flat=True)[count:count + step_size] + asset_ids = list(asset_ids) + if not asset_ids: + break + count += len(asset_ids) + + accounts = list(account_model.objects.filter(asset_id__in=asset_ids)) + + # {asset_id_account_username: account.id}} + asset_accounts_mapper = {} + for a in accounts: + try: + k = f'{a.asset_id}_{a.username}' + asset_accounts_mapper[k] = str(a.id) + except Exception as e: + pass + + update_accounts = [] + for a in accounts: + try: + if not a.created_by: + continue + created_by_list = a.created_by.split('::') + if len(created_by_list) != 2: + continue + su_from_username = created_by_list[1] + if not su_from_username: + continue + k = f'{a.asset_id}_{su_from_username}' + su_from_id = asset_accounts_mapper.get(k) + if not su_from_id: + continue + a.su_from_id = su_from_id + update_accounts.append(a) + except Exception as e: + pass + + count_account += len(update_accounts) + + log_msg = "\t - [{}]: Update accounts su_from: {}-{} {:.2f}s" + try: + account_model.objects.bulk_update(update_accounts, ['su_from_id']) + except Exception as e: + status = 'Failed' + else: + status = 'Success' + print(log_msg.format(status, count_account - len(update_accounts), count_account, time.time() - start)) + + def migrate_db_accounts(apps, schema_editor): app_perm_model = apps.get_model('perms', 'ApplicationPermission') account_model = apps.get_model('accounts', 'Account') @@ -196,5 +265,6 @@ class Migration(migrations.Migration): operations = [ migrations.RunPython(migrate_asset_accounts), + migrations.RunPython(update_asset_accounts_su_from), migrations.RunPython(migrate_db_accounts), ]