Merge pull request #3 from airflow-plugins/feature/add-keys

Feature/add keys

Author: Ben

Diff for: operators/s3_to_redshift.py

@@ -60,11 +60,33 @@ class S3ToRedshiftOperator(BaseOperator):
                                     with. Only required if using a load_type of
                                     "upsert".
     :type incremental_key:          string
+    :param foreign_key:             *(optional)* This specifies any foreign_keys
+                                    in the table and which corresponding table
+                                    and key they reference. This may be either
+                                    a dictionary or list of dictionaries (for
+                                    multiple foreign keys). The fields that are
+                                    required in each dictionary are:
+                                        - column_name
+                                        - reftable
+                                        - ref_column
+    :type foreign_key:              dictionary
+    :param distkey:                 *(optional)* The distribution key for the
+                                    table. Only one key may be specified.
+    :type distkey:                  string
+    :param sortkey:                 *(optional)* The sort keys for the table.
+                                    If more than one key is specified, set this
+                                    as a list.
+    :type sortkey:                  string
+    :param sort_type:               *(optional)* The style of distribution
+                                    to sort the table. Possible values include:
+                                        - compound
+                                        - interleaved
+                                    Defaults to "compound".
+    :type sort_type:                string
     """
 
     template_fields = ('s3_key',
-                       'origin_schema',
-                       'com')
+                       'origin_schema')
 
     @apply_defaults
     def __init__(self,
@@ -81,7 +103,10 @@ def __init__(self,
                  load_type='append',
                  primary_key=None,
                  incremental_key=None,
-                 timeformat='auto',
+                 foreign_key={},
+                 distkey=None,
+                 sortkey='',
+                 sort_type='COMPOUND',
                  *args,
                  **kwargs):
         super().__init__(*args, **kwargs)
@@ -98,14 +123,29 @@ def __init__(self,
         self.load_type = load_type
         self.primary_key = primary_key
         self.incremental_key = incremental_key
-        self.timeformat = timeformat
+        self.foreign_key = foreign_key
+        self.distkey = distkey
+        self.sortkey = sortkey
+        self.sort_type = sort_type
 
         if self.load_type.lower() not in ["append", "rebuild", "upsert"]:
             raise Exception('Please choose "append", "rebuild", or "upsert".')
 
         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)):
+            raise Exception('Sort Keys must be specified as either a string or list.')
+
+        if not (isinstance(self.foreign_key, dict) or isinstance(self.foreign_key, list)):
+            raise Exception('Foreign Keys must be specified as either a dictionary or a list of dictionaries.')
+
+        if ((',' in self.distkey) or not isinstance(self.distkey, str)):
+            raise Exception('Only one distribution key may be specified.')
+
+        if self.sort_type.lower() not in ('compound', 'interleaved'):
+            raise Exception('Please choose "compound" or "interleaved" for sort type.')
+
     def execute(self, context):
         # Append a random string to the end of the staging table to ensure
         # no conflicts if multiple processes running concurrently.
@@ -337,6 +377,8 @@ def create_if_not_exists(self, schema, pg_hook, temp=False):
         for item in schema:
             k = "{quote}{key}{quote}".format(quote='"', key=item['name'])
             field = ' '.join([k, item['type']])
+            if isinstance(self.sortkey, str) and self.sortkey == item['name']:
+                field += ' sortkey'
             output += field
             output += ', '
         # Remove last comma and space after schema items loop ends
@@ -346,12 +388,50 @@ def create_if_not_exists(self, schema, pg_hook, temp=False):
         else:
             copy_table = self.table
         create_schema_query = \
-            '''CREATE SCHEMA IF NOT EXISTS "{0}";'''.format(
-                                                          self.redshift_schema)
+            '''
+            CREATE SCHEMA IF NOT EXISTS "{0}";
+            '''.format(self.redshift_schema)
+
+        pk = ''
+        fk = ''
+        dk = ''
+        sk = ''
+
+        if self.primary_key:
+            pk = ', primary key("{0}")'.format(self.primary_key)
+
+        if self.foreign_key:
+            if isinstance(self.foreign_key, list):
+                fk = ', '
+                for i, e in enumerate(self.foreign_key):
+                    fk += 'foreign key("{0}") references {1}("{2}")'.format(e['column_name'],
+                                                                        e['reftable'],
+                                                                        e['ref_column'])
+                    if i != (len(self.foreign_key) - 1):
+                        fk += ', '""
+            elif isinstance(self.foreign_key, dict):
+                fk += ', '
+                fk += 'foreign key("{0}") references {1}("{2}")'.format(self.foreign_key['column_name'],
+                                                                   self.foreign_key['reftable'],
+                                                                   self.foreign_key['ref_column'])
+        if self.distkey:
+            dk = 'distkey({})'.format(self.distkey)
+
+        if self.sortkey:
+            if isinstance(self.sortkey, list):
+                sk += '{0} sortkey({1})'.format(self.sort_type, ', '.join(["{}".format(e) for e in self.sortkey]))
+
         create_table_query = \
-            '''CREATE TABLE IF NOT EXISTS "{0}"."{1}" ({2})'''.format(
-                                                          self.redshift_schema,
-                                                          copy_table,
-                                                          output)
+            '''
+            CREATE TABLE IF NOT EXISTS "{schema}"."{table}"
+            ({fields}{primary_key}{foreign_key}) {distkey} {sortkey}
+            '''.format(schema=self.redshift_schema,
+                       table=copy_table,
+                       fields=output,
+                       primary_key=pk,
+                       foreign_key=fk,
+                       distkey=dk,
+                       sortkey=sk)
+
         pg_hook.run(create_schema_query)
         pg_hook.run(create_table_query)