From e7f8086f23b23b6c64b0e164b562cd9e36f74e1a Mon Sep 17 00:00:00 2001
From: caoqianming <caoqianming@foxmail.com>
Date: Mon, 26 Feb 2024 08:08:47 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20base=20=E6=8F=90=E4=BE=9B=E6=9F=A5?=
 =?UTF-8?q?=E8=AF=A2=E5=AD=90=E9=9B=86=E6=95=B0=E6=8D=AE=E6=96=B9=E6=B3=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/utils/queryset.py | 25 +++++++++++++++----------
 apps/utils/viewsets.py | 11 ++++++++++-
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/apps/utils/queryset.py b/apps/utils/queryset.py
index 9f7b5d32..e346091f 100755
--- a/apps/utils/queryset.py
+++ b/apps/utils/queryset.py
@@ -2,22 +2,27 @@ from django.apps import apps
 from django.db.models import Q
 
 
-def get_child_queryset_u(checkQueryset, obj, hasParent=True):
+def get_child_queryset_u(checkQueryset, hasParent=True):
     '''
     获取所有子集
     查的范围checkQueryset
     父obj
     是否包含父默认True
     '''
-    cls = type(obj)
-    queryset = cls.objects.none()
-    if hasParent:
-        queryset = cls.objects.filter(pk=obj.id)
-    child_queryset = checkQueryset.filter(parent=obj)
-    while child_queryset.exists():
-        queryset = queryset | child_queryset
-        child_queryset = checkQueryset.filter(parent__in=child_queryset)
-    return queryset
+    cls = type(checkQueryset.model)
+    if hasattr(cls, 'parent'):
+        queryset = cls.objects.none()
+        if hasParent:
+            queryset = checkQueryset
+        child_queryset = checkQueryset.filter(parent__in=queryset)
+        while child_queryset.exists():
+            queryset = queryset | child_queryset
+            child_queryset = checkQueryset.filter(parent__in=child_queryset)
+        return queryset
+    elif hasParent:
+        return checkQueryset
+    else:
+        return checkQueryset.none()
 
 
 def get_child_queryset(name, pk, hasParent=True):
diff --git a/apps/utils/viewsets.py b/apps/utils/viewsets.py
index 869187c4..6870db08 100755
--- a/apps/utils/viewsets.py
+++ b/apps/utils/viewsets.py
@@ -12,7 +12,7 @@ from apps.system.models import DataFilter, Dept, User
 from apps.utils.errors import PKS_ERROR
 from apps.utils.mixins import MyLoggingMixin, BulkCreateModelMixin, BulkUpdateModelMixin, BulkDestroyModelMixin
 from apps.utils.permission import ALL_PERMS, RbacPermission, get_user_perms_map
-from apps.utils.queryset import get_child_queryset2
+from apps.utils.queryset import get_child_queryset2, get_child_queryset_u
 from apps.utils.serializers import PkSerializer, ComplexSerializer
 from rest_framework.throttling import UserRateThrottle
 from drf_yasg.utils import swagger_auto_schema
@@ -100,6 +100,15 @@ class CustomGenericViewSet(MyLoggingMixin, GenericViewSet):
             queryset = queryset.select_related(*self.select_related_fields)
         if self.prefetch_related_fields:
             queryset = queryset.prefetch_related(*self.prefetch_related_fields)
+
+        # 查询出子集数据,请谨慎使用
+        has_children = self.request.query_params.get('has_children', 'no')
+        if has_children == 'hasparent':
+            queryset = get_child_queryset_u(queryset, True)
+        elif has_children == 'noparent':
+            queryset = get_child_queryset_u(queryset, False)
+        else:
+            pass
         return queryset
 
     def get_queryset(self):