add truncate option separate from rebuild option

Author: benjamingregory

Diff for: operators/s3_to_redshift.py

@@ -2,12 +2,14 @@
 import random
 import string
 import logging
+
+from airflow.utils.db import provide_session
+from airflow.models import Connection
 from airflow.utils.decorators import apply_defaults
+
 from airflow.models import BaseOperator
 from airflow.hooks.S3_hook import S3Hook
 from airflow.hooks.postgres_hook import PostgresHook
-from airflow.utils.db import provide_session
-from airflow.models import Connection
 
 
 class S3ToRedshiftOperator(BaseOperator):
@@ -46,9 +48,12 @@ class S3ToRedshiftOperator(BaseOperator):
                                     possible values include "mysql".
     :type origin_datatype:          string
     :param load_type:               The method of loading into Redshift that
-                                    should occur. Options are "append",
-                                    "rebuild", and "upsert". Defaults to
-                                    "append."
+                                    should occur. Options:
+                                        - "append"
+                                        - "rebuild"
+                                        - "truncate"
+                                        - "upsert"
+                                    Defaults to "append."
     :type load_type:                string
     :param primary_key:             *(optional)* The primary key for the
                                     destination table. Not enforced by redshift
@@ -128,10 +133,10 @@ def __init__(self,
         self.sortkey = sortkey
         self.sort_type = sort_type
 
-        if self.load_type.lower() not in ["append", "rebuild", "upsert"]:
+        if self.load_type.lower() not in ("append", "rebuild", "truncate", "upsert"):
             raise Exception('Please choose "append", "rebuild", or "upsert".')
 
-        if self.schema_location.lower() not in ['s3', 'local']:
+        if self.schema_location.lower() not in ('s3', 'local'):
             raise Exception('Valid Schema Locations are "s3" or "local".')
 
         if not (isinstance(self.sortkey, str) or isinstance(self.sortkey, list)):
@@ -152,9 +157,12 @@ def execute(self, context):
         letters = string.ascii_lowercase
         random_string = ''.join(random.choice(letters) for _ in range(7))
         self.temp_suffix = '_tmp_{0}'.format(random_string)
+
         if self.origin_schema:
             schema = self.read_and_format()
+
         pg_hook = PostgresHook(self.redshift_conn_id)
+
         self.create_if_not_exists(schema, pg_hook)
         self.reconcile_schemas(schema, pg_hook)
         self.copy_data(pg_hook, schema)
@@ -221,7 +229,6 @@ def read_and_format(self):
                             if i['type'] == e['avro']:
                                 i['type'] = e['redshift']
 
-        print(schema)
         return schema
 
     def reconcile_schemas(self, schema, pg_hook):
@@ -277,7 +284,7 @@ def getS3Conn():
             elif aws_role_arn:
                 creds = ("aws_iam_role={0}"
                     .format(aws_role_arn))
-            
+
             return creds
 
         # Delete records from the destination table where the incremental_key
@@ -331,6 +338,11 @@ def getS3Conn():
             FILLTARGET
             '''.format(self.redshift_schema, self.table, self.temp_suffix)
 
+        drop_sql = \
+            '''
+            DROP TABLE IF EXISTS "{0}"."{1}"
+            '''.format(self.redshift_schema, self.table)
+
         drop_temp_sql = \
             '''
             DROP TABLE IF EXISTS "{0}"."{1}{2}"
@@ -366,6 +378,13 @@ def getS3Conn():
                                                      base_sql)
         if self.load_type == 'append':
             pg_hook.run(load_sql)
+        elif self.load_type == 'rebuild':
+            pg_hook.run(drop_sql)
+            self.create_if_not_exists(schema, pg_hook)
+            pg_hook.run(load_sql)
+        elif self.load_type == 'truncate':
+            pg_hook.run(truncate_sql)
+            pg_hook.run(load_sql)
         elif self.load_type == 'upsert':
             self.create_if_not_exists(schema, pg_hook, temp=True)
             load_temp_sql = \
@@ -378,9 +397,6 @@ def getS3Conn():
             pg_hook.run(delete_confirm_sql)
             pg_hook.run(append_sql, autocommit=True)
             pg_hook.run(drop_temp_sql)
-        elif self.load_type == 'rebuild':
-            pg_hook.run(truncate_sql)
-            pg_hook.run(load_sql)
 
     def create_if_not_exists(self, schema, pg_hook, temp=False):
         output = ''