Updated the query to no use union - failing a lot of tests

MIT-LCP · Apr 12, 2024 · d836c0b · d836c0b
1 parent f43a814
commit d836c0b
Showing 1 changed file with 21 additions and 34 deletions.
diff --git a/physionet-django/user/managers.py b/physionet-django/user/managers.py
@@ -1,4 +1,5 @@
-from django.db.models import DateTimeField, ExpressionWrapper, QuerySet, F, Q, OuterRef, Subquery
+from django.db.models import (DateTimeField, ExpressionWrapper, QuerySet, F, Q,
+                              OuterRef, Subquery, Case, When)
 from django.utils import timezone
 
 from user.enums import TrainingStatus, RequiredField
@@ -9,55 +10,41 @@ def get_review(self):
         return self.filter(status=TrainingStatus.REVIEW)
 
     def get_valid(self):
-        non_platfrom_trainings = self.filter(
-            Q(status=TrainingStatus.ACCEPTED),
-            ~Q(training_type__required_field=RequiredField.PLATFORM),
-            Q(training_type__valid_duration__isnull=True)
-            | Q(process_datetime__gte=timezone.now() - F('training_type__valid_duration')),
-        ).annotate(
-            valid_datetime=ExpressionWrapper(
-                F('process_datetime') + F('training_type__valid_duration'), output_field=DateTimeField()
-            )
-        )
-
         course_duration = Course.objects.filter(trainings=OuterRef('pk')).values('valid_duration')[:1]
 
-        platform_trainings = self.filter(
+        return self.filter(
             Q(status=TrainingStatus.ACCEPTED),
-            Q(training_type__required_field=RequiredField.PLATFORM),
-            Q(process_datetime__gte=timezone.now() -Subquery(course_duration)),
+            Q(training_type__valid_duration__isnull=True)
+            | Q(process_datetime__gte=timezone.now() - Case(
+                When(training_type__required_field=RequiredField.PLATFORM, then=Subquery(course_duration)),
+                default=F('training_type__valid_duration')
+            )),
         ).annotate(
             valid_datetime=ExpressionWrapper(
-                F('process_datetime') + Subquery(course_duration), output_field=DateTimeField()
+                F('process_datetime') + Case(
+                    When(training_type__required_field=RequiredField.PLATFORM, then=Subquery(course_duration)),
+                    default=F('training_type__valid_duration')
+                ), output_field=DateTimeField()
             )
         )
 
-        return non_platfrom_trainings.union(platform_trainings)
-
     def get_expired(self):
-        non_platform_trainings = self.filter(
-            Q(status=TrainingStatus.ACCEPTED),
-            ~Q(training_type__required_field=RequiredField.PLATFORM),
-            Q(process_datetime__lt=timezone.now() - F('training_type__valid_duration')),
-        ).annotate(
-            valid_datetime=ExpressionWrapper(
-                F('process_datetime') + F('training_type__valid_duration'), output_field=DateTimeField()
-            )
-        )
-
         course_duration = Course.objects.filter(trainings=OuterRef('pk')).values('valid_duration')[:1]
 
-        platform_trainings = self.filter(
+        return self.filter(
             Q(status=TrainingStatus.ACCEPTED),
-            Q(training_type__required_field=RequiredField.PLATFORM),
-            Q(process_datetime__lt=timezone.now() - Subquery(course_duration)),
+            Q(process_datetime__lt=timezone.now() - Case(
+                When(training_type__required_field=RequiredField.PLATFORM, then=Subquery(course_duration)),
+                default=F('training_type__valid_duration')
+            )),
         ).annotate(
             valid_datetime=ExpressionWrapper(
-                F('process_datetime') + Subquery(course_duration), output_field=DateTimeField()
+                F('process_datetime') + Case(
+                    When(training_type__required_field=RequiredField.PLATFORM, then=Subquery(course_duration)),
+                    default=F('training_type__valid_duration')
+                ), output_field=DateTimeField()
             )
         )
 
-        return non_platform_trainings.union(platform_trainings)
-
     def get_rejected(self):
         return self.filter(status=TrainingStatus.REJECTED)