Merge pull request #2182 from IFRCGo/feature/country-plan-import-clea…

…n-up

Clean-up StrategicPriority & MembershipCoordination on import

country_plan/tasks.py

@@ -1,6 +1,7 @@
 import logging
 import operator
 import re
+import typing
 from functools import reduce
 
 import pandas as pd
@@ -14,6 +15,10 @@
 logger = logging.getLogger(__name__)
 
 
+class RollBackException(Exception):
+    pass
+
+
 class CountryPlanImporter:
     CSV_COLUMN = [
         "ISO3",
@@ -124,15 +129,11 @@ def _save_country_plan(cls, row):
         )
 
     @classmethod
-    def process(cls, file):
-        data = pd.read_excel(file, skiprows=4, na_filter=False, usecols=cls.CSV_COLUMN)
-        data = data[:-1]
-
+    def _process(cls, data) -> typing.List[typing.Dict]:
         errors = []
         for row in data.itertuples():
             try:
-                with transaction.atomic():
-                    cls._save_country_plan(row)
+                cls._save_country_plan(row)
             except Exception as e:
                 logger.error(f"Failed to import Country-Plan: iso3:{row.ISO3}", exc_info=True)
                 errors.append(
@@ -143,6 +144,32 @@ def process(cls, file):
                 )
         return errors
 
+    @classmethod
+    def process(cls, file) -> typing.List[typing.Dict]:
+        data = pd.read_excel(file, skiprows=4, na_filter=False, usecols=cls.CSV_COLUMN)
+        data = data[:-1]
+        errors = []
+
+        try:
+            with transaction.atomic():
+                # Clean-up existing dataset StrategicPriority,MembershipCoordination
+                StrategicPriority.objects.all().delete()
+                MembershipCoordination.objects.all().delete()
+
+                # Add new data
+                errors = cls._process(data)
+                if errors:
+                    raise RollBackException()
+        except Exception as e:
+            # handle custom rollback error manually
+            if not isinstance(e, RollBackException):
+                logger.error("Failed to process Country-Plan", exc_info=True)
+                errors.append(dict(error=str(e)))
+            pass
+
+        # Return errors
+        return errors
+
 
 @shared_task
 def process_data_import(pk):