Code change s3 ingestion add write option

dataservices/core/mixins.py

@@ -2,8 +2,10 @@
 import zlib
 
 import boto3
+import sqlalchemy as sa
 from django.conf import settings
 from pg_bulk_ingest import to_file_like_obj
+from sqlalchemy.ext.declarative import declarative_base
 
 
 def unzip_s3_gzip_file(file_body, max_bytes):
@@ -52,16 +54,25 @@ def get_s3_file(key):
 
 class S3DownloadMixin:
 
-    def do_handle(self, prefix, save_func):
+    def delete_temp_tables(self, table_names):
+        Base = declarative_base()
+        metadata = sa.MetaData()
+        engine = sa.create_engine(settings.DATABASE_URL, future=True)
+        metadata.reflect(bind=engine)
+        for name in table_names:
+            table = metadata.tables.get(name, None)
+            if table is not None:
+                Base.metadata.drop_all(engine, [table], checkfirst=True)
+
+    def do_handle(self, prefix):
         """
         Download latest data file from s3
         unzip downloaded data file
         store latest data in the database
         params:
             prefix: str - Bucket Path on the Dataservices s3 bucket.
-            save_func: method - Method that saves the <data> param to the database.
         """
-        assert None not in [prefix, save_func]
+        assert None not in [prefix]
 
         page_iterator = get_s3_paginator(prefix)
         files = []
@@ -81,4 +92,6 @@ def do_handle(self, prefix, save_func):
                     chunks = unzip_s3_gzip_file(body, (32 + zlib.MAX_WBITS))
                     text_lines = io.TextIOWrapper(to_file_like_obj(chunks), encoding="utf-8", newline="")
                     if text_lines:
-                        save_func(text_lines)
+                        return text_lines
+                    else:
+                        return None

dataservices/management/commands/helpers.py

@@ -76,6 +76,50 @@ def handle(self, *args, **options):
         self.stdout.write(self.style.SUCCESS(f'{prefix} {count} records.'))
 
 
+class BaseS3IngestionCommand(BaseCommand):
+
+    save_func = None
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            '--write',
+            action='store_true',
+            help='Store dataset records',
+        )
+
+    def load_data(self):
+        """
+        The procedure for fetching the data. Subclasses must implement this method.
+        """
+        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):
+
+        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 count:
+            self.stdout.write(self.style.SUCCESS(f'{prefix} {count} records.'))
+
+
 class MarketGuidesDataIngestionCommand(BaseDataWorkspaceIngestionCommand):
     def should_ingestion_run(self, view_name, table_name):
         dataflow_metadata = self.get_dataflow_metadata(table_name)

dataservices/management/commands/import_dbt_investment_opportunities.py

@@ -2,10 +2,9 @@
 
 import sqlalchemy as sa
 from django.conf import settings
-from django.core.management.base import BaseCommand
 
 from dataservices.core.mixins import S3DownloadMixin
-from dataservices.management.commands.helpers import ingest_data
+from dataservices.management.commands.helpers import BaseS3IngestionCommand, ingest_data
 
 
 def get_investment_opportunities_data_table(metadata):
@@ -63,29 +62,25 @@ def get_table_data():
     )
 
 
-def save_investment_opportunities_data(data):
+class Command(BaseS3IngestionCommand, S3DownloadMixin):
 
-    engine = sa.create_engine(settings.DATABASE_URL, future=True)
-
-    metadata = sa.MetaData()
-
-    data_table = get_investment_opportunities_data_table(metadata)
+    help = 'Import DBT investment opportunities data from s3'
 
-    def on_before_visible(conn, ingest_table, batch_metadata):
-        pass
+    def load_data(self, save_data=True, *args, **options):
+        data = self.do_handle(prefix=settings.INVESTMENT_OPPORTUNITIES_S3_PREFIX)
+        return data
 
-    def batches(_):
-        yield get_investment_opportunities_batch(data, data_table)
+    def save_import_data(self, data):
+        engine = sa.create_engine(settings.DATABASE_URL, future=True)
 
-    ingest_data(engine, metadata, on_before_visible, batches)
+        metadata = sa.MetaData()
 
+        data_table = get_investment_opportunities_data_table(metadata)
 
-class Command(BaseCommand, S3DownloadMixin):
+        def on_before_visible(conn, ingest_table, batch_metadata):
+            pass
 
-    help = 'Import DBT investment opportunities data from s3'
+        def batches(_):
+            yield get_investment_opportunities_batch(data, data_table)
 
-    def handle(self, *args, **options):
-        self.do_handle(
-            prefix=settings.INVESTMENT_OPPORTUNITIES_S3_PREFIX,
-            save_func=save_investment_opportunities_data,
-        )
+        ingest_data(engine, metadata, on_before_visible, batches)

dataservices/management/commands/import_dbt_sectors.py

@@ -2,10 +2,9 @@
 
 import sqlalchemy as sa
 from django.conf import settings
-from django.core.management.base import BaseCommand
 
 from dataservices.core.mixins import S3DownloadMixin
-from dataservices.management.commands.helpers import ingest_data
+from dataservices.management.commands.helpers import BaseS3IngestionCommand, ingest_data
 
 
 def get_dbtsector_table_batch(data, data_table):
@@ -52,29 +51,28 @@ def get_dbtsector_postgres_table(metadata):
     )
 
 
-def save_dbt_sectors_data(data):
+class Command(BaseS3IngestionCommand, S3DownloadMixin):
 
-    engine = sa.create_engine(settings.DATABASE_URL, future=True)
+    help = 'Import DBT Sector list data from s3'
 
-    metadata = sa.MetaData()
+    def load_data(self, save_data=True, *args, **options):
+        data = self.do_handle(
+            prefix=settings.DBT_SECTOR_S3_PREFIX,
+        )
+        return data
 
-    data_table = get_dbtsector_postgres_table(metadata)
+    def save_import_data(self, data):
 
-    def on_before_visible(conn, ingest_table, batch_metadata):
-        pass
+        engine = sa.create_engine(settings.DATABASE_URL, future=True)
 
-    def batches(_):
-        yield get_dbtsector_table_batch(data, data_table)
+        metadata = sa.MetaData()
 
-    ingest_data(engine, metadata, on_before_visible, batches)
+        data_table = get_dbtsector_postgres_table(metadata)
 
+        def on_before_visible(conn, ingest_table, batch_metadata):
+            pass
 
-class Command(BaseCommand, S3DownloadMixin):
+        def batches(_):
+            yield get_dbtsector_table_batch(data, data_table)
 
-    help = 'Import DBT Sector list data from s3'
-
-    def handle(self, *args, **options):
-        self.do_handle(
-            prefix=settings.DBT_SECTOR_S3_PREFIX,
-            save_func=save_dbt_sectors_data,
-        )
+        ingest_data(engine, metadata, on_before_visible, batches)