Merge pull request #64 from nineaiyu/dev

Dev

Author: nineaiyu

common/base/magic.py

@@ -11,7 +11,7 @@
 from importlib import import_module
 
 from django.core.cache import cache
-from django.db import close_old_connections
+from django.db import close_old_connections, connection
 from django.http.response import HttpResponse
 
 from common.utils import get_logger
@@ -341,3 +341,39 @@ def wrapper(*_args, **_kwargs):
         return wrapper
 
     return decorator
+
+
+import functools
+
+
+def timeit(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        logger.info(f"{func.__name__} run time:{end_time - start_time}")
+        return result
+
+    return wrapper
+
+
+class SQLCounter:
+    def __init__(self):
+        self.count = 0
+
+    def __call__(self, execute, sql, params, many, context):
+        self.count += 1
+        return execute(sql, params, many, context)
+
+
+def count_sql_queries(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        sql_counter = SQLCounter()
+        with connection.execute_wrapper(sql_counter):
+            result = func(*args, **kwargs)
+        logger.info(f"{func.__name__} sql queries count: {sql_counter.count}")
+        return result
+
+    return wrapper

common/cache/redis.py

@@ -30,10 +30,11 @@ def format_input(data):
 
 class CacheList(object):
 
-    def __init__(self, key, max_size=1024):
+    def __init__(self, key, max_size=1024, timeout=None):
         self.connect = get_redis_connection("default")
         self.key = key
         self.max_size = max_size
+        self.timeout = timeout
 
     def auto_ltrim(self):
         stop = self.connect.llen(self.key)
@@ -44,6 +45,8 @@ def auto_ltrim(self):
     def push(self, json_data, *args):
         self.connect.lpush(self.key, json.dumps(json_data), *[json.dumps(x) for x in args])
         self.auto_ltrim()
+        if self.timeout is not None:
+            self.connect.expire(self.key, self.timeout)
 
     def pop(self):
         try:
@@ -56,10 +59,11 @@ def pop(self):
     def delete(self):
         self.connect.delete(self.key)
 
+    def len(self):
+        return self.connect.llen(self.key)
 
-class RobotMsgCache(CacheList):
-    def __init__(self, key='', max_size=1024):
-        super().__init__(f'ai_robot_chat_msg_{key}', max_size)
+    def get_all(self):
+        return [format_return(k) for k in self.connect.lrange(self.key, 0, -1)]
 
 
 class CacheSet(object):

common/core/filter.py

@@ -19,6 +19,7 @@
 from rest_framework.exceptions import NotAuthenticated
 from rest_framework.filters import BaseFilterBackend
 
+from common.base.magic import timeit, count_sql_queries
 from common.cache.storage import CommonResourceIDsCache
 from common.core.db.utils import RelatedManager
 from common.utils import get_logger
@@ -140,6 +141,8 @@ def get_filter_q_base(model, permission, user_obj=None, dept_obj=None):
     return q1
 
 
+@timeit
+@count_sql_queries
 def get_filter_queryset(queryset: QuerySet, user_obj: UserInfo):
     """
     1.获取所有数据权限规则

common/core/modelset.py

@@ -4,10 +4,13 @@
 # filename : modelset
 # author : ly_13
 # date : 6/2/2023
+import itertools
 import json
+import uuid
 from hashlib import md5
 from typing import Callable
 
+import math
 from django.conf import settings
 from django.db import transaction
 from django.forms.widgets import SelectMultiple, DateTimeInput
@@ -29,15 +32,32 @@
 from common.core.config import SysConfig
 from common.core.response import ApiResponse
 from common.core.serializers import BasePrimaryKeyRelatedField
-from common.core.utils import get_query_post_pks
 from common.drf.renders.csv import CSVFileRenderer
 from common.drf.renders.excel import ExcelFileRenderer
 from common.swagger.utils import get_default_response_schema
+from common.tasks import background_task_view_set_job
 from common.utils import get_logger
 
 logger = get_logger(__name__)
 
 
+def run_view_by_celery_task(view, request, kwargs, list_data, batch_length=100):
+    task = kwargs.get("task", request.query_params.get('task', 'true').lower() in ['true', '1', 'yes'])  # 默认为任务异步导入
+    if task:
+        view_str = f"{view.__class__.__module__}.{view.__class__.__name__}"
+        meta = request.META
+        task_id = uuid.uuid4()
+        meta["task_count"] = math.ceil(len(list_data) / batch_length)
+        meta["action"] = view.action
+        for index, batch in enumerate(itertools.batched(list_data, batch_length)):
+            meta["task_id"] = f"{task_id}_{index}"
+            meta["task_index"] = index
+            res = background_task_view_set_job.apply_async(args=(view_str, meta, json.dumps(batch), view.action_map),
+                                                           task_id=meta["task_id"])
+            logger.info(f"add {view_str} task success. {res}")
+        return ApiResponse(detail=_("Task add success"))
+
+
 class CacheDetailResponseMixin(object):
     def get_cache_key(self, view_instance, view_method, request, args, kwargs):
         func_name = f'{view_instance.__class__.__name__}_{view_method.__name__}'
@@ -106,20 +126,14 @@ class RankAction(object):
     get_queryset: Callable
 
     @extend_schema(
-        request=OpenApiRequest(
-            build_object_type(
-                properties={'pks': build_array_type(build_basic_type(OpenApiTypes.STR))},
-                required=['pks'],
-                description="主键列表"
-            )
-        ),
+        request=OpenApiRequest(build_array_type(build_basic_type(OpenApiTypes.STR))),
         responses=get_default_response_schema()
     )
     @action(methods=['post'], detail=False, url_path='rank')
     def rank(self, request, *args, **kwargs):
         """{cls}排序"""
         rank = 1
-        for pk in get_query_post_pks(request):
+        for pk in request.data:
             self.filter_queryset(self.get_queryset()).filter(pk=pk).update(rank=rank)
             rank += 1
         return ApiResponse(detail=_("Sorting saved successfully"))
@@ -336,7 +350,7 @@ def get_input_type(value, info):
         return ApiResponse(data=results)
 
 
-class BaseViewSet(GenericViewSet):
+class BaseViewSet(object):
     action: Callable
     extra_filter_class = []
 
@@ -373,30 +387,26 @@ class BatchDestroyAction(object):
     perform_destroy: Callable
 
     @extend_schema(
-        request=OpenApiRequest(
-            build_object_type(
-                properties={'pks': build_array_type(build_basic_type(OpenApiTypes.STR))},
-                required=['pks'],
-                description="主键列表"
-            )
-        ),
+        request=OpenApiRequest(build_array_type(build_basic_type(OpenApiTypes.STR))),
         responses=get_default_response_schema()
     )
     @action(methods=['post'], detail=False, url_path='batch-destroy')
     def batch_destroy(self, request, *args, **kwargs):
         """批量删除{cls}"""
-        pks = get_query_post_pks(request)
-        if not pks:
-            return ApiResponse(code=1003, detail=_("Operation failed. Primary key list does not exist"))
+
+        # response = run_view_by_celery_task(self, request, kwargs, request.data, batch_length=30)
+        # if response:
+        #     return response
+
         # queryset  delete() 方法进行批量删除，并不调用模型上的任何 delete() 方法,需要通过循环对象进行删除
         count = 0
-        for instance in self.filter_queryset(self.get_queryset()).filter(pk__in=pks):
+        for instance in self.filter_queryset(self.get_queryset()).filter(pk__in=request.data):
             try:
                 deleted, _rows_count = self.perform_destroy(instance)
                 if deleted:
                     count += 1
-            except Exception:
-                pass
+            except Exception as e:
+                logger.error(f"failed to destroy instance {instance} with error {e}")
         return ApiResponse(detail=_("Operation successful. Batch deleted {} data").format(count))
 
 
@@ -481,6 +491,11 @@ class ImportExportDataAction(CreateAction, UpdateAction, OnlyExportDataAction):
     @transaction.atomic
     def import_data(self, request, *args, **kwargs):
         """导入{cls}数据"""
+
+        response = run_view_by_celery_task(self, request, kwargs, request.data)
+        if response:
+            return response
+
         act = request.query_params.get('action')
         ignore_error = request.query_params.get('ignore_error', 'false') == 'true'
         if act and request.data:
@@ -509,25 +524,25 @@ def import_data(self, request, *args, **kwargs):
         return ApiResponse(detail=_("Operation failed. Abnormal data"), code=1001)
 
 
-class DetailUpdateModelSet(UpdateAction, DetailAction, BaseViewSet):
+class DetailUpdateModelSet(BaseViewSet, UpdateAction, DetailAction, GenericViewSet):
     pass
 
 
-class OnlyListModelSet(ListAction, SearchFieldsAction, SearchColumnsAction, BaseViewSet):
+class OnlyListModelSet(BaseViewSet, ListAction, SearchFieldsAction, SearchColumnsAction, GenericViewSet):
     pass
 
 
 # 全部 ViewSet 包含增删改查 
-class BaseModelSet(CreateAction, DestroyAction, UpdateAction, ListAction, DetailAction, SearchFieldsAction,
-                   SearchColumnsAction, BatchDestroyAction, BaseViewSet):
+class BaseModelSet(BaseViewSet, CreateAction, DestroyAction, UpdateAction, ListAction, DetailAction, SearchFieldsAction,
+                   SearchColumnsAction, BatchDestroyAction, GenericViewSet):
     pass
 
 
 # 只允许读和删除，不允许创建和修改
-class ListDeleteModelSet(DestroyAction, ListAction, DetailAction, SearchFieldsAction, SearchColumnsAction,
-                         BatchDestroyAction, BaseViewSet):
+class ListDeleteModelSet(BaseViewSet, DestroyAction, ListAction, DetailAction, SearchFieldsAction, SearchColumnsAction,
+                         BatchDestroyAction, GenericViewSet):
     pass
 
 
-class NoDetailModelSet(UpdateAction, DetailAction, SearchColumnsAction, BaseViewSet):
+class NoDetailModelSet(BaseViewSet, UpdateAction, DetailAction, SearchColumnsAction, GenericViewSet):
     pass

common/notifications.py

@@ -6,7 +6,7 @@
 from common.models import Monitor
 from notifications.backends import BACKEND
 from notifications.models import SystemMsgSubscription
-from notifications.notifications import SystemMessage
+from notifications.notifications import SystemMessage, UserMessage
 from system.models import UserInfo
 
 
@@ -128,3 +128,39 @@ def get_monitor_latest_average_value(num=3):
 
     def initial_terminals(self):
         self._terminals = [self.get_monitor_latest_average_value()]
+
+
+class TaskMessage(object):
+    def get_html_msg(self) -> dict:
+        context = dict(
+            subject=self.subject,
+            name=self.user.nickname,
+            **self.task,
+        )
+        message = render_to_string('notify/msg_task.html', context)
+        return {
+            'subject': self.subject,
+            'message': message
+        }
+
+
+class ImportDataMessage(TaskMessage, UserMessage):
+    category = 'Task Message'
+    category_label = _('Task Message')
+    message_type_label = _('Import data message')
+
+    def __init__(self, user, task):
+        super().__init__(user)
+        self.task = task
+        self.subject = _('Import {} data {} message').format(self.task.get("view_doc"), self.task.get("status"))
+
+
+class BatchDeleteDataMessage(TaskMessage, UserMessage):
+    category = 'Task Message'
+    category_label = _('Task Message')
+    message_type_label = _('Batch delete data message')
+
+    def __init__(self, user, task):
+        super().__init__(user)
+        self.task = task
+        self.subject = _('Batch delete {} data {} message').format(self.task.get("view_doc"), self.task.get("status"))

common/tasks.py

@@ -6,20 +6,25 @@
 # date : 7/30/2024
 import datetime
 import os
+from io import BytesIO
 
 from celery import shared_task
 from celery.utils.log import get_task_logger
 from django.conf import settings
+from django.core.handlers.wsgi import WSGIRequest
 from django.core.mail import send_mail, EmailMultiAlternatives, get_connection
 from django.utils import timezone
+from django.utils.module_loading import import_string
 from django.utils.translation import gettext_lazy as _
 from django_celery_beat.models import PeriodicTask
 
+from common.cache.redis import CacheList
 from common.celery.decorator import register_as_period_task, after_app_ready_start
 from common.celery.utils import delete_celery_periodic_task, disable_celery_periodic_task, get_celery_periodic_task, \
     create_or_update_celery_periodic_tasks
 from common.models import Monitor
-from common.notifications import ServerPerformanceCheckUtil
+from common.notifications import ServerPerformanceCheckUtil, ImportDataMessage, BatchDeleteDataMessage
+from common.utils.timezone import local_now_display
 from server.celery import app
 
 logger = get_task_logger(__name__)
@@ -130,3 +135,42 @@ def create_or_update_registered_periodic_tasks():
 @register_as_period_task(interval=60)
 def check_server_performance_period():
     ServerPerformanceCheckUtil().check_and_publish()
+
+
+@shared_task(verbose_name=_("Run background task view set"))
+def background_task_view_set_job(view: str, meta: dict, data: str, action_map: dict):
+    cache = CacheList(f"view_task_{meta.get("task_id").split("_")[0]}", timeout=3600 * 24)
+    task_info = {
+        "start_time": local_now_display(),
+        "task_id": meta.get("task_id"),
+        "task_index": meta.get("task_index")
+    }
+    view_func = import_string(view)
+    b_data = data.encode("utf-8")
+    meta["wsgi.input"] = BytesIO(b_data)
+    meta["CONTENT_TYPE"] = "application/json"
+    meta["CONTENT_LENGTH"] = len(b_data)
+    request = WSGIRequest(meta)
+    result = view_func.as_view(action_map)(request, task=False)
+    task_info["result"] = result.data.get("detail", result.data)
+    task_info["end_time"] = local_now_display()
+    task_info["status"] = result.data.get("code") == 1000
+    cache.push(task_info)
+    if cache.len() == meta["task_count"]:
+        task_results = cache.get_all()
+        cache.delete()
+        state = all([task["status"] for task in task_results])
+        task_info = {
+            "task_name": view,
+            "view_doc": view_func.__doc__,
+            "state": state,
+            "status": _("Operation successful") if state else _("Operation failed"),
+            "tasks": sorted(task_results, key=lambda task: task["task_index"])
+        }
+        match meta["action"]:
+            case "import_data":
+                ImportDataMessage(getattr(request, "user"), task_info).publish()
+            case "batch_destroy":
+                BatchDeleteDataMessage(getattr(request, "user"), task_info).publish()
+
+    return task_info