fix: Fixed issue of v2 to v3 Account missing `su_from`

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),
     ]