refactor(exports): rearrange export tasks

Author: Hicham

cohort/services/cohort_result.py

@@ -8,7 +8,7 @@
 from cohort.serializers import WSJobStatus, JobName
 from cohort.services.base_service import CommonService
 from cohort.services.dated_measure import dm_service
-from cohort.services.utils import get_authorization_header, ServerError
+from cohort.services.utils import ServerError
 from admin_cohort.services.ws_event_manager import WebsocketManager, WebSocketMessageType
 from cohort.tasks import create_cohort
 
@@ -36,7 +36,12 @@ def build_query(cohort_source_id: str, fhir_filter_id: str = None) -> str:
                  }
         return json.dumps(query)
 
-    def create_cohort_subset(self, request, owner_id: str, table_name: str, source_cohort: CohortResult, fhir_filter_id: str) -> CohortResult:
+    def create_cohort_subset(self,
+                             auth_headers: dict,
+                             owner_id: str,
+                             table_name: str,
+                             source_cohort: CohortResult,
+                             fhir_filter_id: str) -> CohortResult:
         def copy_query_snapshot(snapshot: RequestQuerySnapshot) -> RequestQuerySnapshot:
             return RequestQuerySnapshot.objects.create(owner=snapshot.owner,
                                                        request=snapshot.request,
@@ -62,7 +67,7 @@ def copy_dated_measure(dm: DatedMeasure) -> DatedMeasure:
                                                     dated_measure=new_dm,
                                                     request_query_snapshot=new_rqs)
         with transaction.atomic():
-            self.handle_cohort_creation(cohort_subset, request, False)
+            self.handle_cohort_creation(cohort_subset, auth_headers)
         return cohort_subset
 
     @staticmethod
@@ -74,9 +79,9 @@ def count_active_jobs():
         return CohortResult.objects.filter(request_job_status__in=active_statuses) \
             .count()
 
-    def handle_cohort_creation(self, cohort: CohortResult, request, global_estimate: bool) -> None:
+    def handle_cohort_creation(self, cohort: CohortResult, auth_headers: dict, global_estimate: bool=False) -> None:
         if global_estimate:
-            dm_service.handle_global_count(cohort, request)
+            dm_service.handle_global_count(cohort, auth_headers)
         try:
             if cohort.parent_cohort and cohort.sampling_ratio:
                 json_query = self.build_query(cohort_source_id=cohort.parent_cohort.group_id)
@@ -91,7 +96,7 @@ def handle_cohort_creation(self, cohort: CohortResult, request, global_estimate:
 
             create_cohort.s(cohort_id=cohort.pk,
                             json_query=json_query,
-                            auth_headers=get_authorization_header(request),
+                            auth_headers=auth_headers,
                             cohort_creator_cls=self.operator_cls,
                             sampling_ratio=cohort.sampling_ratio) \
                          .apply_async()

cohort/services/dated_measure.py

@@ -3,21 +3,20 @@
 from cohort.models.dated_measure import GLOBAL_DM_MODE
 from cohort.serializers import WSJobStatus, JobName
 from cohort.services.base_service import CommonService
-from cohort.services.utils import get_authorization_header, ServerError
+from cohort.services.utils import ServerError
 from admin_cohort.services.ws_event_manager import WebsocketManager, WebSocketMessageType
 from cohort.tasks import cancel_previous_count_jobs, count_cohort
 
 
 class DatedMeasureService(CommonService):
     job_type = "count"
 
-    def handle_count(self, dm: DatedMeasure, request) -> None:
-        stage_details = request.data.get("stageDetails", None)
+    def handle_count(self, dm: DatedMeasure, auth_headers: dict, stage_details) -> None:
         cancel_previous_count_jobs.s(dm_id=dm.uuid, cohort_counter_cls=self.operator_cls).apply_async()
         try:
             count_cohort.s(dm_id=dm.uuid,
                            json_query=dm.request_query_snapshot.serialized_query,
-                           auth_headers=get_authorization_header(request),
+                           auth_headers=auth_headers,
                            cohort_counter_cls=self.operator_cls,
                            stage_details=stage_details
                            ) \
@@ -26,14 +25,14 @@ def handle_count(self, dm: DatedMeasure, request) -> None:
             dm.delete()
             raise ServerError("Could not launch count request") from e
 
-    def handle_global_count(self, cohort: CohortResult, request) -> None:
+    def handle_global_count(self, cohort: CohortResult, auth_headers: dict) -> None:
         dm_global = DatedMeasure.objects.create(mode=GLOBAL_DM_MODE,
-                                                owner=request.user,
-                                                request_query_snapshot_id=request.data.get("request_query_snapshot"))
+                                                owner=cohort.owner,
+                                                request_query_snapshot=cohort.request_query_snapshot)
         try:
             count_cohort.s(dm_id=dm_global.uuid,
                            json_query=dm_global.request_query_snapshot.serialized_query,
-                           auth_headers=get_authorization_header(request),
+                           auth_headers=auth_headers,
                            cohort_counter_cls=self.operator_cls,
                            global_estimate=True) \
                 .apply_async()

cohort/services/feasibility_study.py

@@ -15,7 +15,7 @@
 from cohort.models import FeasibilityStudy
 from cohort.services.base_service import CommonService
 
-from cohort.services.utils import get_authorization_header, ServerError
+from cohort.services.utils import ServerError
 from cohort.tasks import feasibility_study_count, send_feasibility_study_notification, send_email_feasibility_report_error, \
                          send_email_feasibility_report_ready
 
@@ -49,11 +49,11 @@ def bound_number(n: int) -> str:
 class FeasibilityStudyService(CommonService):
     job_type = "count"
 
-    def handle_feasibility_study_count(self, fs: FeasibilityStudy, request) -> None:
+    def handle_feasibility_study_count(self, fs: FeasibilityStudy, auth_headers: dict) -> None:
         try:
             chain(*(feasibility_study_count.s(fs_id=fs.uuid,
                                               json_query=fs.request_query_snapshot.serialized_query,
-                                              auth_headers=get_authorization_header(request),
+                                              auth_headers=auth_headers,
                                               cohort_counter_cls=self.operator_cls),
                     send_feasibility_study_notification.s(fs.uuid)))()
         except Exception as e:

cohort/tests/test_service_cohort_result.py

@@ -1,7 +1,5 @@
 from unittest import mock
 
-from requests import Request as HttpRequest
-
 from admin_cohort.tests.tests_tools import new_random_user, TestCaseWithDBs
 from cohort.models import Folder, Request, RequestQuerySnapshot, CohortResult, DatedMeasure, FhirFilter
 from cohort.services.cohort_result import CohortResultService
@@ -31,12 +29,10 @@ def setUp(self):
                                                      filter="param=value")
         self.cohort_result_service = CohortResultService()
 
-    @mock.patch('cohort.services.cohort_result.get_authorization_header')
     @mock.patch('cohort.services.cohort_result.create_cohort.apply_async')
-    def test_create_cohort_subset(self, mock_get_auth_headers, mock_celery_task):
-        mock_get_auth_headers.return_value = {"authorization": "token"}
+    def test_create_cohort_subset(self, mock_celery_task):
         mock_celery_task.return_value = None
-        cohort_subset = self.cohort_result_service.create_cohort_subset(request=HttpRequest(data={}),
+        cohort_subset = self.cohort_result_service.create_cohort_subset(auth_headers={},
                                                                         owner_id=self.user1.pk,
                                                                         table_name="Table_01",
                                                                         source_cohort=self.source_cohort,

cohort/tests/tests_view_cohort_result.py

@@ -231,12 +231,10 @@ def setUp(self):
             status=status.HTTP_400_BAD_REQUEST,
         )
 
-    @mock.patch('cohort.services.cohort_result.get_authorization_header')
     @mock.patch('cohort.services.cohort_result.create_cohort.apply_async')
     @mock.patch('cohort.services.dated_measure.count_cohort.apply_async')
     def check_create_case_with_mock(self, case: CohortCreateCase, mock_count_task: MagicMock, mock_create_task: MagicMock,
-                                    mock_header: MagicMock, other_view: any, view_kwargs: dict):
-        mock_header.return_value = None
+                                    other_view: any, view_kwargs: dict):
         mock_create_task.return_value = None
         mock_count_task.return_value = None
 
@@ -253,7 +251,6 @@ def check_create_case_with_mock(self, case: CohortCreateCase, mock_count_task: M
                 mock_create_task.assert_called() if case.mock_create_task_called else mock_create_task.assert_not_called()
 
         mock_create_task.assert_called() if case.mock_create_task_called else mock_create_task.assert_not_called()
-        mock_header.assert_called() if case.mock_create_task_called else mock_header.assert_not_called()
 
     def check_create_case(self, case: CohortCreateCase, other_view: Any = None, **view_kwargs):
         return self.check_create_case_with_mock(case, other_view=other_view or None, view_kwargs=view_kwargs)

cohort/tests/tests_view_dated_measure.py

@@ -166,19 +166,16 @@ def setUp(self):
                                                       request_job_status=JobStatus.pending,
                                                       owner=self.user1)
 
-    @mock.patch('cohort.services.dated_measure.get_authorization_header')
     @mock.patch('cohort.services.dated_measure.cancel_previous_count_jobs.apply_async')
     @mock.patch('cohort.services.dated_measure.count_cohort.apply_async')
-    def check_create_case_with_mock(self, case: DMCreateCase, mock_count_task: MagicMock, mock_cancel_task: MagicMock, mock_header: MagicMock,
+    def check_create_case_with_mock(self, case: DMCreateCase, mock_count_task: MagicMock, mock_cancel_task: MagicMock,
                                     other_view: any, view_kwargs: dict):
-        mock_header.return_value = None
         mock_cancel_task.return_value = None
         mock_count_task.return_value = None
 
         with self.captureOnCommitCallbacks(execute=True):
             super(DatedMeasuresCreateTests, self).check_create_case(case, other_view, **(view_kwargs or {}))
 
-        mock_header.assert_called() if case.mock_header_called else mock_header.assert_not_called()
         mock_cancel_task.assert_called() if case.mock_cancel_task_called else mock_cancel_task.assert_not_called()
         mock_count_task.assert_called() if case.mock_count_task_called else mock_count_task.assert_not_called()
 

cohort/views/cohort_result.py

@@ -16,8 +16,8 @@
 from cohort.serializers import CohortResultSerializer, CohortRightsSerializer, CohortResultPatchSerializer, CohortResultCreateSerializer, \
     SampledCohortResultCreateSerializer
 from cohort.services.cohort_rights import cohort_rights_service
+from cohort.services.utils import get_authorization_header
 from cohort.views.shared import UserObjectsRestrictedViewSet
-from exports.services.export import export_service
 
 
 class CohortFilter(filters.FilterSet):
@@ -116,7 +116,8 @@ def list(self, request, *args, **kwargs):
     def create(self, request, *args, **kwargs):
         global_estimate = request.data.pop("global_estimate", False)
         response = super().create(request, *args, **kwargs)
-        transaction.on_commit(lambda: cohort_service.handle_cohort_creation(request=request,
+        auth_headers = get_authorization_header(request)
+        transaction.on_commit(lambda: cohort_service.handle_cohort_creation(auth_headers=auth_headers,
                                                                             cohort=response.data.serializer.instance,
                                                                             global_estimate=global_estimate))
         return response
@@ -136,8 +137,6 @@ def partial_update(self, request, *args, **kwargs):
         else:
             response = super().partial_update(request, *args, **kwargs)
             cohort_service.handle_cohort_post_update(cohort=cohort, caller=request.user.username)
-            if cohort.export_table.exists():
-                export_service.check_all_cohort_subsets_created(export=cohort.export_table.first().export)
         cohort_service.ws_send_to_client(cohort=cohort)
         return response
 

cohort/views/dated_measure.py

@@ -10,7 +10,7 @@
 from cohort.serializers import DatedMeasureSerializer, DatedMeasureCreateSerializer, DatedMeasurePatchSerializer
 from cohort.services.dated_measure import dm_service
 from cohort.services.request_refresh_schedule import requests_refresher_service
-from cohort.services.utils import await_celery_task
+from cohort.services.utils import await_celery_task, get_authorization_header
 from cohort.views.shared import UserObjectsRestrictedViewSet
 
 _logger = logging.getLogger('info')
@@ -44,8 +44,11 @@ def get_queryset(self):
     @transaction.atomic
     def create(self, request, *args, **kwargs):
         response = super().create(request, *args, **kwargs)
-        transaction.on_commit(lambda: dm_service.handle_count(request=request,
-                                                              dm=response.data.serializer.instance))
+        auth_headers = get_authorization_header(request)
+        stage_details = request.data.get("stageDetails", None)
+        transaction.on_commit(lambda: dm_service.handle_count(dm=response.data.serializer.instance,
+                                                              auth_headers=auth_headers,
+                                                              stage_details=stage_details))
         return response
 
     @await_celery_task

cohort/views/feasibility_study.py

@@ -12,6 +12,7 @@
 from cohort.models import FeasibilityStudy
 from cohort.serializers import FeasibilityStudySerializer, FeasibilityStudyCreateSerializer, FeasibilityStudyPatchSerializer
 from cohort.services.feasibility_study import feasibility_study_service
+from cohort.services.utils import get_authorization_header
 from cohort.views.shared import UserObjectsRestrictedViewSet
 
 _logger = logging.getLogger('info')
@@ -48,7 +49,8 @@ def retrieve(self, request, *args, **kwargs):
     @transaction.atomic
     def create(self, request, *args, **kwargs):
         response = super().create(request, *args, **kwargs)
-        transaction.on_commit(lambda: feasibility_study_service.handle_feasibility_study_count(request=request,
+        auth_headers = get_authorization_header(request)
+        transaction.on_commit(lambda: feasibility_study_service.handle_feasibility_study_count(auth_headers=auth_headers,
                                                                                                fs=response.data.serializer.instance))
         return response
 

exporters/enums.py

@@ -1,16 +1,16 @@
-from enum import Enum, StrEnum
+from enum import StrEnum
 
 from admin_cohort.types import JobStatus
 
 
-class ExportTypes(Enum):
+class ExportTypes(StrEnum):
     CSV = "csv"
     XLSX = "xlsx"
     HIVE = "hive"
 
     @staticmethod
     def default() -> str:
-        return ExportTypes.CSV.value
+        return ExportTypes.CSV
 
     @property
     def allow_download(self) -> bool:
@@ -41,15 +41,15 @@ class APIJobStatus(StrEnum):
     flowerNotAccessible = 'flowerNotAccessible'
 
 
-status_mapper = {APIJobStatus.Received.value: JobStatus.new,
-                 APIJobStatus.Pending.value: JobStatus.pending,
-                 APIJobStatus.Retry.value: JobStatus.pending,
-                 APIJobStatus.Running.value: JobStatus.started,
-                 APIJobStatus.FinishedSuccessfully.value: JobStatus.finished,
-                 APIJobStatus.FinishedWithError.value: JobStatus.failed,
-                 APIJobStatus.FinishedWithTimeout.value: JobStatus.failed,
-                 APIJobStatus.flowerNotAccessible.value: JobStatus.failed,
-                 APIJobStatus.Failure.value: JobStatus.failed,
-                 APIJobStatus.NotFound.value: JobStatus.failed,
-                 APIJobStatus.Revoked.value: JobStatus.cancelled
+status_mapper = {APIJobStatus.Received: JobStatus.new,
+                 APIJobStatus.Pending: JobStatus.pending,
+                 APIJobStatus.Retry: JobStatus.pending,
+                 APIJobStatus.Running: JobStatus.started,
+                 APIJobStatus.FinishedSuccessfully: JobStatus.finished,
+                 APIJobStatus.FinishedWithError: JobStatus.failed,
+                 APIJobStatus.FinishedWithTimeout: JobStatus.failed,
+                 APIJobStatus.flowerNotAccessible: JobStatus.failed,
+                 APIJobStatus.Failure: JobStatus.failed,
+                 APIJobStatus.NotFound: JobStatus.failed,
+                 APIJobStatus.Revoked: JobStatus.cancelled
                  }