Greatuk 101 automate IMF data

dataservices/management/commands/import_market_guides_data.py

@@ -24,6 +24,10 @@ class Command(BaseCommand):
             'table_name': 'trade__uk_services_nsa',
             'view_name': 'TopFiveServicesExportsByCountryView',
         },
+        'import_world_economic_outlook_data': {
+            'table_name': 'world_economic_outlook__by_country',
+            'view_name': 'EconomicHighlightsView',
+        },
     }
 
     def add_arguments(self, parser):

dataservices/management/commands/import_world_economic_outlook_data.py

@@ -1,15 +1,14 @@
 import pandas as pd
 import sqlalchemy as sa
-from django.conf import settings
-from django.core.management import BaseCommand
 
 from dataservices.models import Country, WorldEconomicOutlookByCountry
 
+from .helpers import MarketGuidesDataIngestionCommand
 
-class Command(BaseCommand):
+
+class Command(MarketGuidesDataIngestionCommand):
     help = 'Import IMF world economic outlook data by country from Data Workspace'
 
-    engine = sa.create_engine(settings.DATA_WORKSPACE_DATASETS_URL, execution_options={'stream_results': True})
     sql = '''
         SELECT
             iso AS ons_iso_alpha_3_code,
@@ -38,7 +37,14 @@ class Command(BaseCommand):
             AND NULLIF(TRIM(x.value), '') IS NOT NULL;
     '''
 
-    def handle(self, *args, **options):
+    def add_arguments(self, parser):
+        parser.add_argument(
+            '--write',
+            action='store_true',
+            help='Store dataset records',
+        )
+
+    def load_data(self):
         data = []
         chunks = pd.read_sql(sa.text(self.sql), self.engine, chunksize=10000)
 
@@ -64,7 +70,4 @@ def handle(self, *args, **options):
                     )
                 )
 
-        WorldEconomicOutlookByCountry.objects.all().delete()
-        WorldEconomicOutlookByCountry.objects.bulk_create(data)
-
-        self.stdout.write(self.style.SUCCESS('All done, bye!'))
+        return data

dataservices/management/commands/tests/test_import_data.py

@@ -388,8 +388,15 @@ def test_import_world_economic_outlook_data(read_sql_mock, world_economic_outloo
     assert len(models.WorldEconomicOutlookByCountry.objects.all()) == 0
 
     management.call_command('import_countries')
+
+    # Dry run
     management.call_command('import_world_economic_outlook_data')
 
+    assert len(models.WorldEconomicOutlookByCountry.objects.all()) == 0
+
+    # Write option
+    management.call_command('import_world_economic_outlook_data', '--write')
+
     assert len(models.WorldEconomicOutlookByCountry.objects.all()) == 10
 
 
@@ -426,7 +433,7 @@ def test_import_metadata_source_data_filter_tables():
 @pytest.mark.django_db
 @pytest.mark.parametrize(
     'env, review_requested_x_times',
-    [('dev', 0), ('staging', 3), ('uat', 0), ('production', 0)],
+    [('dev', 0), ('staging', 4), ('uat', 0), ('production', 0)],
 )
 @mock.patch('dataservices.management.commands.import_market_guides_data.call_command')
 @mock.patch('dataservices.management.commands.helpers.MarketGuidesDataIngestionCommand.should_ingestion_run')
@@ -439,14 +446,15 @@ def test_import_market_guides_data(
             'import_uk_total_trade_data',
             'import_uk_trade_in_goods_data',
             'import_uk_trade_in_services_data',
+            'import_world_economic_outlook_data',
         ]
         mock_should_run.return_value = False
         management.call_command('import_market_guides_data', '--write')
         assert mock_call_command.call_count == 0
 
         mock_should_run.return_value = True
         management.call_command('import_market_guides_data', '--write')
-        assert mock_call_command.call_count == 6
+        assert mock_call_command.call_count == 8
         assert mock_send_review_request.call_count == review_requested_x_times
 
         for command in command_list:
@@ -462,12 +470,13 @@ def test_import_market_guides_data_dry_run(mock_call_command, mock_should_run):
         'import_uk_total_trade_data',
         'import_uk_trade_in_goods_data',
         'import_uk_trade_in_services_data',
+        'import_world_economic_outlook_data',
     ]
     mock_should_run.return_value = True
 
     management.call_command('import_market_guides_data')
 
-    assert mock_call_command.call_count == 6
+    assert mock_call_command.call_count == 8
 
     for command in command_list:
         assert command in str(mock_call_command.call_args_list)
@@ -486,9 +495,9 @@ def test_import_market_guides_data_error(mock_call_command, mock_should_run, moc
     management.call_command('import_market_guides_data')
 
     mock_email_string = str(mock_error_email.mock_calls)
-    assert mock_call_command.call_count == 3
+    assert mock_call_command.call_count == 4
     assert 'oops' in str(mock_call_command.side_effect)
-    assert mock_error_email.call_count == 3
+    assert mock_error_email.call_count == 4
     assert 'area_of_error' in mock_email_string
     assert 'error_type' in mock_email_string
     assert 'error_details' in mock_email_string