Merge pull request #10746 from jumpserver/pr@dev@perf_category_api_sql

perf: 修改 category 引起的 sql 查询过多

apps/assets/const/custom.py

@@ -1,3 +1,6 @@
+from collections import defaultdict
+
+from common.decorators import cached_method
 from .base import BaseType
 
 
@@ -9,7 +12,8 @@ class CustomTypes(BaseType):
         except Exception:
             return []
         types = set([p.type for p in platforms])
-        return [(t, t) for t in types]
+        choices = [(t, t) for t in types]
+        return choices
 
     @classmethod
     def _get_base_constrains(cls) -> dict:
@@ -37,13 +41,20 @@ class CustomTypes(BaseType):
         return constrains
 
     @classmethod
+    @cached_method(5)
     def _get_protocol_constrains(cls) -> dict:
-        constrains = {}
-        for platform in cls.get_custom_platforms():
-            choices = list(platform.protocols.values_list('name', flat=True))
-            if platform.type in constrains:
-                choices = constrains[platform.type]['choices'] + choices
-            constrains[platform.type] = {'choices': choices}
+        from assets.models import PlatformProtocol
+        _constrains = defaultdict(set)
+        protocols = PlatformProtocol.objects \
+            .filter(platform__category='custom') \
+            .values_list('name', 'platform__type')
+        for name, tp in protocols:
+            _constrains[tp].add(name)
+
+        constrains = {
+            tp: {'choices': list(choices)}
+            for tp, choices in _constrains.items()
+        }
         return constrains
 
     @classmethod
@@ -51,6 +62,8 @@ class CustomTypes(BaseType):
         return {}
 
     @classmethod
+    @cached_method(5)
     def get_custom_platforms(cls):
         from assets.models import Platform
-        return Platform.objects.filter(category='custom')
+        platforms = Platform.objects.filter(category='custom')
+        return platforms

apps/assets/const/types.py

@@ -193,7 +193,6 @@ class AllTypes(ChoicesMixin):
         }
         return node
 
-
     @classmethod
     def asset_to_node(cls, asset, pid):
         node = {

apps/authentication/views/login.py

@@ -2,15 +2,19 @@
 #
 
 from __future__ import unicode_literals
-import os
+
 import datetime
+import os
 from typing import Callable
 
-from django.db import IntegrityError
-from django.templatetags.static import static
+from django.conf import settings
+from django.contrib.auth import BACKEND_SESSION_KEY
 from django.contrib.auth import login as auth_login, logout as auth_logout
-from django.http import HttpResponse, HttpRequest
+from django.db import IntegrityError
+from django.http import HttpRequest, HttpResponse
 from django.shortcuts import reverse, redirect
+from django.templatetags.static import static
+from django.urls import reverse_lazy
 from django.utils.decorators import method_decorator
 from django.utils.translation import ugettext as _, get_language
 from django.views.decorators.cache import never_cache
@@ -18,16 +22,13 @@ from django.views.decorators.csrf import csrf_protect
 from django.views.decorators.debug import sensitive_post_parameters
 from django.views.generic.base import TemplateView, RedirectView
 from django.views.generic.edit import FormView
-from django.conf import settings
-from django.urls import reverse_lazy
-from django.contrib.auth import BACKEND_SESSION_KEY
 
 from common.utils import FlashMessageUtil, static_or_direct
 from users.utils import (
     redirect_user_first_login_or_index
 )
-from ..const import RSA_PRIVATE_KEY, RSA_PUBLIC_KEY
 from .. import mixins, errors
+from ..const import RSA_PRIVATE_KEY, RSA_PUBLIC_KEY
 from ..forms import get_user_login_form_cls
 
 __all__ = [

apps/common/decorators.py

@@ -6,6 +6,7 @@ import inspect
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
+from functools import wraps
 
 from django.db import transaction
 
@@ -217,3 +218,24 @@ def do_test():
     end = time.time()
     using = end - s
     print("end : %s, using: %s" % (end, using))
+
+
+def cached_method(ttl=20):
+    _cache = {}
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            key = (func, args, tuple(sorted(kwargs.items())))
+            # 检查缓存是否存在且未过期
+            if key in _cache and time.time() - _cache[key]['timestamp'] < ttl:
+                return _cache[key]['result']
+
+            # 缓存过期或不存在，执行方法并缓存结果
+            result = func(*args, **kwargs)
+            _cache[key] = {'result': result, 'timestamp': time.time()}
+            return result
+
+        return wrapper
+
+    return decorator