From 3aa52e8f4213acac521f8522e833d80b1711c406 Mon Sep 17 00:00:00 2001 From: rithviknishad Date: Tue, 15 Oct 2024 13:01:02 +0530 Subject: [PATCH] make asset, bed, assetbed get_queryset reusable based on user --- care/facility/api/viewsets/asset.py | 17 +------ care/facility/api/viewsets/bed.py | 48 +++----------------- care/facility/api/viewsets/camera_preset.py | 49 ++++----------------- care/utils/queryset/asset_bed.py | 47 ++++++++++++++++++++ 4 files changed, 65 insertions(+), 96 deletions(-) create mode 100644 care/utils/queryset/asset_bed.py diff --git a/care/facility/api/viewsets/asset.py b/care/facility/api/viewsets/asset.py index 15dd00e2aa..fc66eff4bf 100644 --- a/care/facility/api/viewsets/asset.py +++ b/care/facility/api/viewsets/asset.py @@ -62,6 +62,7 @@ from care.utils.assetintegration.asset_classes import AssetClasses from care.utils.cache.cache_allowed_facilities import get_accessible_facilities from care.utils.filters.choicefilter import CareChoiceFilter, inverse_choices +from care.utils.queryset.asset_bed import get_asset_queryset from care.utils.queryset.asset_location import get_asset_location_queryset from care.utils.queryset.facility import get_facility_queryset from config.authentication import MiddlewareAuthentication @@ -290,21 +291,7 @@ class AssetViewSet( filterset_class = AssetFilter def get_queryset(self): - user = self.request.user - queryset = self.queryset - if user.is_superuser: - pass - elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]: - queryset = queryset.filter(current_location__facility__state=user.state) - elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]: - queryset = queryset.filter( - current_location__facility__district=user.district - ) - else: - allowed_facilities = get_accessible_facilities(user) - queryset = queryset.filter( - current_location__facility__id__in=allowed_facilities - ) + queryset = get_asset_queryset(user=self.request.user, queryset=self.queryset) return queryset.annotate( latest_status=Subquery( AvailabilityRecord.objects.filter( diff --git a/care/facility/api/viewsets/bed.py b/care/facility/api/viewsets/bed.py index 336b5f83c2..db9dd6652f 100644 --- a/care/facility/api/viewsets/bed.py +++ b/care/facility/api/viewsets/bed.py @@ -30,6 +30,7 @@ from care.users.models import User from care.utils.cache.cache_allowed_facilities import get_accessible_facilities from care.utils.filters.choicefilter import CareChoiceFilter, inverse_choices +from care.utils.queryset.asset_bed import get_asset_bed_queryset, get_bed_queryset inverse_bed_type = inverse_choices(BedTypeChoices) @@ -76,27 +77,14 @@ class BedViewSet( filterset_class = BedFilter def get_queryset(self): - user = self.request.user - queryset = self.queryset - - queryset = queryset.annotate( + queryset = self.queryset.annotate( is_occupied=Exists( ConsultationBed.objects.filter( bed__id=OuterRef("id"), end_date__isnull=True ) ) ) - - if user.is_superuser: - pass - elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]: - queryset = queryset.filter(facility__state=user.state) - elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]: - queryset = queryset.filter(facility__district=user.district) - else: - allowed_facilities = get_accessible_facilities(user) - queryset = queryset.filter(facility__id__in=allowed_facilities) - return queryset + return get_bed_queryset(user=self.request.user, queryset=queryset) @transaction.atomic def create(self, request, *args, **kwargs): @@ -168,18 +156,7 @@ class AssetBedViewSet( lookup_field = "external_id" def get_queryset(self): - user = self.request.user - queryset = self.queryset - if user.is_superuser: - pass - elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]: - queryset = queryset.filter(bed__facility__state=user.state) - elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]: - queryset = queryset.filter(bed__facility__district=user.district) - else: - allowed_facilities = get_accessible_facilities(user) - queryset = queryset.filter(bed__facility__id__in=allowed_facilities) - return queryset + return get_asset_bed_queryset(user=self.request.user, queryset=self.queryset) class PatientAssetBedFilter(filters.FilterSet): @@ -212,20 +189,9 @@ class PatientAssetBedViewSet(ListModelMixin, GenericViewSet): ] def get_queryset(self): - user = self.request.user - queryset = self.queryset - if user.is_superuser: - pass - elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]: - queryset = queryset.filter(bed__facility__state=user.state) - elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]: - queryset = queryset.filter(bed__facility__district=user.district) - else: - allowed_facilities = get_accessible_facilities(user) - queryset = queryset.filter(bed__facility__id__in=allowed_facilities) - return queryset.filter( - bed__facility__external_id=self.kwargs["facility_external_id"] - ) + return get_asset_bed_queryset( + user=self.request.user, queryset=self.queryset + ).filter(bed__facility__external_id=self.kwargs["facility_external_id"]) class ConsultationBedFilter(filters.FilterSet): diff --git a/care/facility/api/viewsets/camera_preset.py b/care/facility/api/viewsets/camera_preset.py index 17c9a30f9e..bfb168834b 100644 --- a/care/facility/api/viewsets/camera_preset.py +++ b/care/facility/api/viewsets/camera_preset.py @@ -5,9 +5,12 @@ from rest_framework.viewsets import GenericViewSet, ModelViewSet from care.facility.api.serializers.camera_preset import CameraPresetSerializer -from care.facility.models import Asset, AssetBed, Bed, CameraPreset -from care.users.models import User -from care.utils.cache.cache_allowed_facilities import get_accessible_facilities +from care.facility.models import CameraPreset +from care.utils.queryset.asset_bed import ( + get_asset_bed_queryset, + get_asset_queryset, + get_bed_queryset, +) class AssetBedCameraPresetViewSet(ModelViewSet): @@ -19,19 +22,9 @@ class AssetBedCameraPresetViewSet(ModelViewSet): permission_classes = (IsAuthenticated,) def get_asset_bed_obj(self): - user = self.request.user - queryset = AssetBed.objects.filter( + queryset = get_asset_bed_queryset(self.request.user).filter( external_id=self.kwargs["assetbed_external_id"] ) - if user.is_superuser: - pass - elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]: - queryset = queryset.filter(bed__facility__state=user.state) - elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]: - queryset = queryset.filter(bed__facility__district=user.district) - else: - allowed_facilities = get_accessible_facilities(user) - queryset = queryset.filter(bed__facility__id__in=allowed_facilities) return get_object_or_404(queryset) def get_queryset(self): @@ -52,35 +45,11 @@ class CameraPresetViewSet(GenericViewSet, ListModelMixin): permission_classes = (IsAuthenticated,) def get_bed_obj(self, external_id: str): - user = self.request.user - queryset = Bed.objects.filter(external_id=external_id) - if user.is_superuser: - pass - elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]: - queryset = queryset.filter(facility__state=user.state) - elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]: - queryset = queryset.filter(facility__district=user.district) - else: - allowed_facilities = get_accessible_facilities(user) - queryset = queryset.filter(facility__id__in=allowed_facilities) + queryset = get_bed_queryset(self.request.user).filter(external_id=external_id) return get_object_or_404(queryset) def get_asset_obj(self, external_id: str): - user = self.request.user - queryset = Asset.objects.filter(external_id=external_id) - if user.is_superuser: - pass - elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]: - queryset = queryset.filter(current_location__facility__state=user.state) - elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]: - queryset = queryset.filter( - current_location__facility__district=user.district - ) - else: - allowed_facilities = get_accessible_facilities(user) - queryset = queryset.filter( - current_location__facility__id__in=allowed_facilities - ) + queryset = get_asset_queryset(self.request.user).filter(external_id=external_id) return get_object_or_404(queryset) def get_queryset(self): diff --git a/care/utils/queryset/asset_bed.py b/care/utils/queryset/asset_bed.py new file mode 100644 index 0000000000..472605edb5 --- /dev/null +++ b/care/utils/queryset/asset_bed.py @@ -0,0 +1,47 @@ +from care.facility.models import Asset, AssetBed, Bed +from care.users.models import User +from care.utils.cache.cache_allowed_facilities import get_accessible_facilities + + +def get_asset_bed_queryset(user, queryset=None): + queryset = queryset or AssetBed.objects.all() + if user.is_superuser: + pass + elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]: + queryset = queryset.filter(bed__facility__state=user.state) + elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]: + queryset = queryset.filter(bed__facility__district=user.district) + else: + allowed_facilities = get_accessible_facilities(user) + queryset = queryset.filter(bed__facility__id__in=allowed_facilities) + return queryset + + +def get_bed_queryset(user, queryset=None): + queryset = queryset or Bed.objects.all() + if user.is_superuser: + pass + elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]: + queryset = queryset.filter(facility__state=user.state) + elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]: + queryset = queryset.filter(facility__district=user.district) + else: + allowed_facilities = get_accessible_facilities(user) + queryset = queryset.filter(facility__id__in=allowed_facilities) + return queryset + + +def get_asset_queryset(user, queryset=None): + queryset = queryset or Asset.objects.all() + if user.is_superuser: + pass + elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]: + queryset = queryset.filter(current_location__facility__state=user.state) + elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]: + queryset = queryset.filter(current_location__facility__district=user.district) + else: + allowed_facilities = get_accessible_facilities(user) + queryset = queryset.filter( + current_location__facility__id__in=allowed_facilities + ) + return queryset