From 3b8aab8c25e4438729e600f24f466d3b93b7c6b9 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Tue, 2 Apr 2024 15:55:08 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BF=AE=E6=94=B9=20export=20=E4=BD=BF?=
 =?UTF-8?q?=E7=94=A8=E7=9A=84=20serializer?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/common/api/mixin.py      | 12 ++++++++++++
 apps/common/api/serializer.py |  2 ++
 2 files changed, 14 insertions(+)

diff --git a/apps/common/api/mixin.py b/apps/common/api/mixin.py
index 7db2056ac..40ea6ee27 100644
--- a/apps/common/api/mixin.py
+++ b/apps/common/api/mixin.py
@@ -6,6 +6,7 @@ from typing import Callable
 
 from django.db import models
 from django.db.models.signals import m2m_changed
+from rest_framework.request import Request
 from rest_framework.response import Response
 from rest_framework.settings import api_settings
 
@@ -89,6 +90,7 @@ class RelationMixin:
 
 class QuerySetMixin:
     action: str
+    request: Request
     get_serializer_class: Callable
     get_queryset: Callable
 
@@ -98,8 +100,18 @@ class QuerySetMixin:
             return queryset
         if self.action == 'metadata':
             queryset = queryset.none()
+        queryset = self.setup_eager_loading(queryset)
         return queryset
 
+    # Todo: 未来考虑自定义 pagination
+    def setup_eager_loading(self, queryset):
+        if self.request.query_params.get('format') not in ['csv', 'xlsx']:
+            return queryset
+        serializer_class = self.get_serializer_class()
+        if not serializer_class or not hasattr(serializer_class, 'setup_eager_loading'):
+            return queryset
+        return serializer_class.setup_eager_loading(queryset)
+
     def paginate_queryset(self, queryset):
         page = super().paginate_queryset(queryset)
         serializer_class = self.get_serializer_class()
diff --git a/apps/common/api/serializer.py b/apps/common/api/serializer.py
index 84cad19ff..539ff2868 100644
--- a/apps/common/api/serializer.py
+++ b/apps/common/api/serializer.py
@@ -27,6 +27,8 @@ class SerializerMixin:
             return None
         serializer_classes = dict(serializer_classes)
         view_action = self.request.query_params.get('action') or self.action or 'list'
+        if self.request.query_params.get('format'):
+            view_action = 'retrieve'
         serializer_class = serializer_classes.get(view_action)
 
         if serializer_class is None: