add write option to s3 manaagement commands

uktrade · Dec 14, 2024 · 0a468e4 · 0a468e4
1 parent 4b0bcb9
commit 0a468e4
Show file tree

Hide file tree

Showing 10 changed files with 97 additions and 78 deletions.
diff --git a/dataservices/management/commands/helpers.py b/dataservices/management/commands/helpers.py
@@ -93,28 +93,29 @@ def load_data(self):
         """
         raise NotImplementedError('subclasses of MarketGuidesDataIngestionCommand must provide a load_data() method')
 
+    def save_import_data(self, data):
+        """
+        The procedure for saving the data. Subclasses must implement this method.
+        """
+        raise NotImplementedError('subclasses of MarketGuidesDataIngestionCommand must provide a load_data() method')
+
     def handle(self, *args, **options):
-        data, self.save_func = self.load_data(save_data=False)
-        prefix = 'Would create'
+
+        if not options['write']:
+            data = self.load_data(save_data=False)
+            prefix = 'Would create'
+        else:
+            prefix = 'Created'
+            data = self.load_data(save_data=True)
+            self.save_import_data(data)
+
         if isinstance(data, list):
             count = len(data)
+        elif isinstance(data, io.TextIOWrapper):
+            count = len(data.readlines())
         else:
             count = None
 
-        if options['write'] and self.save_func:
-            if isinstance(data, list):
-                count = len(data)
-                prefix = 'Created'
-            else:
-                count = None
-            breakpoint()
-            self.save_func(data)
-        else:
-            if isinstance(data, list):
-                count = len(data)
-            else:
-                count = len(data.readlines())
-
         if count:
             self.stdout.write(self.style.SUCCESS(f'{prefix} {count} records.'))
 

diff --git a/dataservices/management/commands/import_dbt_investment_opportunities.py b/dataservices/management/commands/import_dbt_investment_opportunities.py
@@ -68,10 +68,9 @@ class Command(BaseS3IngestionCommand, S3DownloadMixin):
 
     def load_data(self, save_data=True, *args, **options):
         data = self.do_handle(prefix=settings.INVESTMENT_OPPORTUNITIES_S3_PREFIX)
-        return data, Command.save_investment_opportunities_data
-
-    def save_investment_opportunities_data(self, data):
+        return data
 
+    def save_import_data(self, data):
         engine = sa.create_engine(settings.DATABASE_URL, future=True)
 
         metadata = sa.MetaData()

diff --git a/dataservices/management/commands/import_dbt_sectors.py b/dataservices/management/commands/import_dbt_sectors.py
@@ -59,10 +59,9 @@ def load_data(self, save_data=True, *args, **options):
         data = self.do_handle(
             prefix=settings.DBT_SECTOR_S3_PREFIX,
         )
-        return data, self.save_dbt_sectors_data
+        return data
 
-    @staticmethod
-    def save_dbt_sectors_data(data):
+    def save_import_data(self, data):
 
         engine = sa.create_engine(settings.DATABASE_URL, future=True)
 

diff --git a/dataservices/management/commands/import_eyb_business_cluster_information.py b/dataservices/management/commands/import_eyb_business_cluster_information.py
@@ -288,20 +288,19 @@ def save_tmp_table_data(self, save_data):
         if save_data:
             return
 
-        data = self.save_uk_business_employee_counts_data(save_data=save_data)
+        data = self.save_import_data(save_data=save_data)
         return data
 
     def load_data(self, save_data=True, *args, **options):
         try:
             data = self.save_tmp_table_data(save_data)
-            return data, self.save_uk_business_employee_counts_data
+            return data
         except Exception:
             logger.exception("import_eyb_business_cluster_information failed to ingest data from s3")
         finally:
             self.delete_temp_tables(TEMP_TABLES)
 
-
-    def save_uk_business_employee_counts_data(self, data=[], save_data=True):
+    def save_import_data(self, data=[], save_data=True):
         if save_data:
             self.save_tmp_table_data(save_data)
 

diff --git a/dataservices/management/commands/import_eyb_rent_data.py b/dataservices/management/commands/import_eyb_rent_data.py
@@ -74,9 +74,9 @@ def load_data(self, save_data=True, *args, **options):
         data = self.do_handle(
             prefix=settings.EYB_RENT_S3_PREFIX,
         )
-        return data, self.save_eyb_rent_data
+        return data
 
-    def save_eyb_rent_data(self, data):
+    def save_import_data(self, data):
 
         engine = sa.create_engine(settings.DATABASE_URL, future=True)
 

diff --git a/dataservices/management/commands/import_eyb_salary_data.py b/dataservices/management/commands/import_eyb_salary_data.py
@@ -70,9 +70,9 @@ def load_data(self, save_data=True, *args, **options):
         data = self.do_handle(
             prefix=settings.EYB_SALARY_S3_PREFIX,
         )
-        return data, self.save_eyb_salary_data
+        return data
 
-    def save_eyb_salary_data(self, data):
+    def save_import_data(self, data):
 
         engine = sa.create_engine(settings.DATABASE_URL, future=True)
 

diff --git a/dataservices/management/commands/import_postcodes_from_s3.py b/dataservices/management/commands/import_postcodes_from_s3.py
@@ -80,9 +80,9 @@ def load_data(self, save_data=True, *args, **options):
         data = self.do_handle(
             prefix=settings.POSTCODE_FROM_S3_PREFIX,
         )
-        return data, self.save_postcode_data
+        return data
 
-    def save_postcode_data(self, data):
+    def save_import_data(self, data):
         engine = sa.create_engine(settings.DATABASE_URL, future=True)
 
         metadata = sa.MetaData()

diff --git a/dataservices/management/commands/import_sectors_gva_value_bands.py b/dataservices/management/commands/import_sectors_gva_value_bands.py
@@ -66,9 +66,9 @@ def load_data(self, save_data=True, *args, **options):
         data = self.do_handle(
             prefix=settings.DBT_SECTORS_GVA_VALUE_BANDS_DATA_S3_PREFIX,
         )
-        return data, self.save_sectors_gva_value_bands_data
+        return data
 
-    def save_sectors_gva_value_bands_data(self, data):
+    def save_import_data(self, data):
 
         engine = sa.create_engine(settings.DATABASE_URL, future=True)
 

diff --git a/dataservices/tasks.py b/dataservices/tasks.py
@@ -38,7 +38,7 @@ def run_comtrade_data_ingest(period):
 
 @app.task()
 def run_postcode_data_ingest():
-    call_command('import_postcodes_from_s3')
+    call_command('import_postcodes_from_s3', '--write')
 
 
 @app.task()
@@ -48,27 +48,27 @@ def run_import_countries_territories_regions_dw():
 
 @app.task()
 def run_import_dbt_investment_opportunities():
-    call_command('import_dbt_investment_opportunities')
+    call_command('import_dbt_investment_opportunities', '--write')
 
 
 @app.task()
 def run_import_dbt_sectors():
-    call_command('import_dbt_sectors')
+    call_command('import_dbt_sectors', '--write')
 
 
 @app.task()
 def run_import_eyb_business_cluster_information():
-    call_command('import_eyb_business_cluster_information')
+    call_command('import_eyb_business_cluster_information', '--write')
 
 
 @app.task()
 def run_import_eyb_rent_data():
-    call_command('import_eyb_rent_data')
+    call_command('import_eyb_rent_data', '--write')
 
 
 @app.task()
 def run_import_eyb_salary_data():
-    call_command('import_eyb_salary_data')
+    call_command('import_eyb_salary_data', '--write')
 
 
 @app.task()