Merge pull request #14806 from jumpserver/pr@pam@risk_account

perf: Risk account

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"].id}_{d["risk"]}'
+            account = d["account"]
+            key = f'{account.id}_{d["risk"]}'
             origin_risk = ori_risk_map.get(key)
 
             if origin_risk and origin_risk.status != ConfirmOrIgnore.pending:
@@ -209,9 +210,9 @@ class CheckAccountManager(BaseManager):
                 update_risk(origin_risk)
             else:
                 create_risk({
-                    "account": d["account"],
-                    "asset": d["account"].asset,
-                    "username": d["account"].username,
+                    "account": account,
+                    "asset": account.asset,
+                    "username": account.username,
                     "risk": d["risk"],
                     "details": [{"datetime": now, 'type': 'init'}],
                 })

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

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():

apps/accounts/models/automations/check_account.py

@@ -58,7 +58,9 @@ 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')
+    )
     username = models.CharField(max_length=32, verbose_name=_('Username'))
     account = models.ForeignKey(
         'accounts.Account', on_delete=models.CASCADE, related_name='risks',