Merge pull request #1671 from IFRCGo/fix/auto-internal-plan-load

Country plan import: Only download pdf files

Author: szabozoltan69

country_plan/management/commands/ingest_country_plan_file.py

@@ -1,6 +1,6 @@
 import logging
 import requests
-from typing import Union
+from typing import Union, Optional, Tuple
 from datetime import datetime
 from django.utils.timezone import make_aware
 from django.core.management.base import BaseCommand
@@ -20,38 +20,57 @@
 INTERNAL_SOURCE = 'https://go-api.ifrc.org/Api/FedNetAppeals?AppealsTypeId=1844&Hidden=false'
 
 
+def parse_date(text: Optional[str]) -> Optional[datetime]:
+    """
+    Convert Appeal API datetime into django datetime
+    Parameters
+    ----------
+      text : str
+          Datetime eg: 2022-11-29T11:24:00
+    """
+    if text:
+        return make_aware(
+            # NOTE: Format is assumed by looking at the data from Appeal API
+            datetime.strptime(text, '%Y-%m-%dT%H:%M:%S')
+        )
+
+
+def get_meta_from_url(url) -> Tuple[Optional[str], str]:
+    """
+    Fetch url headers and return content-type and filename
+    """
+    def _get_filename_from_headers(resp):
+        try:
+            # Eg: Content-Disposition: 'attachment;filename=UP_Botswana_2023.pdf'
+            return resp.headers.get('Content-Disposition').split(';')[1].split('=')[1]
+        except Exception:
+            return 'document.pdf'
+
+    # Check if it is a pdf file
+    resp = requests.head(url)
+    return resp.headers.get('Content-Type'), _get_filename_from_headers(resp)
+
+
 class Command(BaseCommand):
     @staticmethod
-    def parse_date(text: str) -> Union[datetime, None]:
+    def load_file_to_country_plan(country_plan: CountryPlan, url: str, filename: str, field_name: str):
         """
-        Convert Appeal API datetime into django datetime
-        Parameters
-        ----------
-          text : str
-              Datetime eg: 2022-11-29T11:24:00
+        Fetch file using url and save to country_plan
         """
-        if text:
-            return make_aware(
-                # NOTE: Format is assumed by looking at the data from Appeal API
-                datetime.strptime(text, '%Y-%m-%dT%H:%M:%S')
-            )
-
-    @staticmethod
-    def load_file_to_country_plan(country_plan: CountryPlan, url: str, filename: str, field_name: str):
         with DownloadFileManager(url, suffix='.pdf') as f:
             getattr(country_plan, field_name).save(
                 filename,
                 File(f),
             )
 
-    def load_for_country(self, country_data, file_field, field_inserted_date_field):
+    def load_for_country(self, country_data: dict, file_field: str, field_inserted_date_field: str):
         country_iso2 = country_data.get('LocationCountryCode')
         country_name = country_data.get('LocationCountryName')
-        public_plan_url = country_data.get('BaseDirectory') or '' + country_data.get('BaseFileName') or ''
-        inserted_date = self.parse_date(country_data.get('Inserted'))
+        plan_url = country_data['BaseDirectory'] + country_data['BaseFileName']
+        inserted_date = parse_date(country_data.get('Inserted'))
         if (
             (country_iso2 is None and country_name is None) or
-            public_plan_url is None or
+            plan_url is None or
             inserted_date is None
         ):
             return
@@ -72,14 +91,18 @@ def load_for_country(self, country_data, file_field, field_inserted_date_field):
         if existing_inserted_date and existing_inserted_date >= inserted_date:
             # No need to do anything here
             return
-        self.stdout.write(f'- Saving data for country:: {country_plan.country.name}')
-        public_plan_url = country_data['BaseDirectory'] + country_data['BaseFileName']
+        content_type, content_name = get_meta_from_url(plan_url)
+        self.stdout.write(f'- Checking plan file for country:: {country_plan.country.name}')
+        if content_type != 'application/pdf':
+            # Only looking for PDFs
+            return
+        self.stdout.write('  - Saving data')
         setattr(country_plan, field_inserted_date_field, inserted_date)
         self.load_file_to_country_plan(
             country_plan,
-            public_plan_url,
+            plan_url,
             # NOTE: File provided are PDF,
-            f"{file_field.replace('_', '-').replace('file', '')}-{country_data['BaseFileName']}.pdf",
+            f"{file_field.replace('_', '-').replace('-file', '')}-{content_name}",
             file_field,
         )
         country_plan.is_publish = True
@@ -90,7 +113,6 @@ def load_for_country(self, country_data, file_field, field_inserted_date_field):
                 'is_publish',
             )
         )
-        return True
 
     def load(self, url: str, file_field: str, field_inserted_date_field: str):
         auth = (settings.APPEALS_USER, settings.APPEALS_PASS)

country_plan/tests/test_commands.py

@@ -45,7 +45,14 @@
         'Inserted': '2022-11-29T11:24:00',
         'LocationCountryCode': 'GR',
         'LocationCountryName': 'Greece'
-    }
+    },
+    {
+        'BaseDirectory': FILE_BASE_DIRECTORY,
+        'BaseFileName': 'NOOP',
+        'Inserted': '2022-11-29T11:24:00',
+        'LocationCountryCode': 'XY',
+        'LocationCountryName': 'Myanmar'
+    },
 ]
 
 INTERNAL_APPEAL_COUNTRY_PLAN_MOCK_RESPONSE = [
@@ -85,13 +92,18 @@ def raise_for_status(self):
         def iter_content(self, **_):
             return self.stream
 
-    def __init__(self, json=None, stream=None):
+    def __init__(self, json=None, stream=None, headers=None):
+        self._headers = headers
         self._json = json
         self.stream = stream
 
     def json(self):
         return self._json
 
+    @property
+    def headers(self):
+        return self._headers
+
     def __enter__(self):
         return MockResponse.FileStream(self.stream)
 
@@ -108,6 +120,18 @@ def mock_request(url, *_, **kwargs):
         if url.startswith(FILE_BASE_DIRECTORY):
             return MockResponse(stream=[b''])
 
+    def mock_request_head(url, *_, **kwargs):
+        headers = {
+            'Content-Type': 'application/pdf',
+            'Content-Disposition': 'attachment;filename=Sample_document_2023.pdf',
+        }
+        if url.endswith('NOOP'):
+            headers['Content-Type'] = 'html/text'
+        elif url.endswith('000004'):
+            headers['Content-Disposition'] = ''
+        return MockResponse(headers=headers)
+
+    @mock.patch('country_plan.management.commands.ingest_country_plan_file.requests.head', side_effect=mock_request_head)
     @mock.patch('country_plan.management.commands.ingest_country_plan_file.requests.get', side_effect=mock_request)
     @mock.patch('main.utils.requests.get', side_effect=mock_request)
     def test_country_plan_ingest(self, *_):
@@ -139,3 +163,7 @@ def test_country_plan_ingest(self, *_):
         assert CountryPlan.objects.filter(is_publish=True).count() == 5
         assert CountryPlan.objects.exclude(public_plan_file='').count() == 4
         assert CountryPlan.objects.exclude(internal_plan_file='').count() == 2
+        # First downloaded document. Others will have Sample_document_2023{random-chars}.pdf
+        assert CountryPlan.objects.filter(country__iso='SY').first().public_plan_file.name.endswith('Sample_document_2023.pdf')
+        # Without attachment filename
+        assert CountryPlan.objects.filter(country__iso='GR').first().public_plan_file.name.endswith('document.pdf')