From 3926d9ff46f26b1cfa9643d367997aefd3760786 Mon Sep 17 00:00:00 2001
From: feng <1304903146@qq.com>
Date: Tue, 14 Jan 2025 10:24:48 +0800
Subject: [PATCH] perf: Risk account

---
 .../automations/check_account/manager.py        |  8 +++++---
 apps/accounts/filters.py                        |  7 ++-----
 apps/accounts/models/account.py                 | 15 ++++-----------
 .../models/automations/check_account.py         | 17 +++++++++++------
 4 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/apps/accounts/automations/check_account/manager.py b/apps/accounts/automations/check_account/manager.py
index c212d20ff..a52a36226 100644
--- a/apps/accounts/automations/check_account/manager.py
+++ b/apps/accounts/automations/check_account/manager.py
@@ -194,7 +194,8 @@ class CheckAccountManager(BaseManager):
 
         now = timezone.now().isoformat()
         for d in self.batch_risks:
-            key = f'{d["account"].asset_id}_{d["account"].username}_{d["risk"]}'
+            account = d["account"]
+            key = f'{account.asset_id}_{account.username}_{d["risk"]}'
             origin_risk = ori_risk_map.get(key)
 
             if origin_risk and origin_risk.status != ConfirmOrIgnore.pending:
@@ -209,8 +210,9 @@ class CheckAccountManager(BaseManager):
                 update_risk(origin_risk)
             else:
                 create_risk({
-                    "asset": d["account"].asset,
-                    "username": d["account"].username,
+                    "account": account,
+                    "asset": account.asset,
+                    "username": account.username,
                     "risk": d["risk"],
                     "details": [{"datetime": now, 'type': 'init'}],
                 })
diff --git a/apps/accounts/filters.py b/apps/accounts/filters.py
index 84dd82452..a4eda5ff6 100644
--- a/apps/accounts/filters.py
+++ b/apps/accounts/filters.py
@@ -7,7 +7,7 @@ from django_filters import rest_framework as drf_filters
 from assets.models import Node
 from common.drf.filters import BaseFilterSet
 from common.utils.timezone import local_zero_hour, local_now
-from .models import Account, GatheredAccount, ChangeSecretRecord, AccountRisk
+from .models import Account, GatheredAccount, ChangeSecretRecord
 
 
 class AccountFilterSet(BaseFilterSet):
@@ -70,10 +70,7 @@ class AccountFilterSet(BaseFilterSet):
         if not value:
             return queryset
 
-        risks = AccountRisk.objects.filter(risk=value)
-        usernames = risks.values_list('username', flat=True)
-        assets = risks.values_list('asset', flat=True)
-        queryset = queryset.filter(username__in=usernames, asset__in=assets)
+        queryset = queryset.filter(risks__risk=value)
         return queryset
 
     @staticmethod
diff --git a/apps/accounts/models/account.py b/apps/accounts/models/account.py
index d7ee6efd4..f7185d005 100644
--- a/apps/accounts/models/account.py
+++ b/apps/accounts/models/account.py
@@ -1,5 +1,4 @@
 from django.db import models
-from django.db.models import Exists, OuterRef
 from django.utils.translation import gettext_lazy as _
 from simple_history.models import HistoricalRecords
 
@@ -169,20 +168,14 @@ class Account(AbsConnectivity, LabeledMixin, BaseAccount):
 
     @classmethod
     def get_risks(cls, queryset=None, risk_type=None):
-        # TODO 数据量大时，子查询性能不佳，考虑用原生sql或者在模型层面做出改动
-        from accounts.models import AccountRisk
-        subquery = AccountRisk.objects.filter(
-            asset_id=OuterRef('asset_id'),
-            username=OuterRef('username')
-        )
-
-        if risk_type:
-            subquery = subquery.filter(risk=risk_type)
+        query = {
+            'risks__risk': risk_type
+        }
 
         if queryset is None:
             queryset = cls.objects.all()
 
-        return queryset.filter(Exists(subquery))
+        return queryset.filter(**query)
 
 
 def replace_history_model_with_mixin():
diff --git a/apps/accounts/models/automations/check_account.py b/apps/accounts/models/automations/check_account.py
index 9964aa030..0bae122bd 100644
--- a/apps/accounts/models/automations/check_account.py
+++ b/apps/accounts/models/automations/check_account.py
@@ -58,14 +58,19 @@ class RiskChoice(TextChoices):
 
 
 class AccountRisk(JMSOrgBaseModel):
-    asset = models.ForeignKey('assets.Asset', on_delete=models.CASCADE, related_name='risks', verbose_name=_('Asset'))
+    asset = models.ForeignKey(
+        'assets.Asset', on_delete=models.CASCADE, related_name='risks', verbose_name=_('Asset')
+    )
+    account = models.ForeignKey(
+        'accounts.Account', on_delete=models.CASCADE, related_name='risks', verbose_name=_('Account'), null=True
+    )
+    status = models.CharField(
+        max_length=32, choices=ConfirmOrIgnore.choices,
+        default=ConfirmOrIgnore.pending, blank=True, verbose_name=_('Status')
+    )
     username = models.CharField(max_length=32, verbose_name=_('Username'))
-    account = models.ForeignKey('accounts.Account', on_delete=models.CASCADE, related_name='risks',
-                                verbose_name=_('Account'), null=True)
-    risk = models.CharField(max_length=128, verbose_name=_('Risk'), choices=RiskChoice.choices)
-    status = models.CharField(max_length=32, choices=ConfirmOrIgnore.choices, default=ConfirmOrIgnore.pending,
-                              blank=True, verbose_name=_('Status'))
     details = models.JSONField(default=list, verbose_name=_('Details'))
+    risk = models.CharField(max_length=128, verbose_name=_('Risk'), choices=RiskChoice.choices)
 
     class Meta:
         verbose_name = _('Account risk')