Merge pull request #1 from robertengel/master

Extend redshift_sqlalchemy by UNLOAD and COPY support

Author: hearsaydev

redshift_sqlalchemy/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.4'
+__version__ = '0.5.0a'
 
 from sqlalchemy.dialects import registry
 

redshift_sqlalchemy/dialect.py

@@ -108,6 +108,7 @@ def _fetch_redshift_column_attributes(self, column):
             text += " SORTKEY"
         return text
 
+
 class RedshiftDialect(PGDialect_psycopg2):
     name = 'redshift'
     ddl_compiler = RedShiftDDLCompiler
@@ -179,37 +180,113 @@ class UnloadFromSelect(Executable, ClauseElement):
     ''' Prepares a RedShift unload statement to drop a query to Amazon S3
     http://docs.aws.amazon.com/redshift/latest/dg/r_UNLOAD_command_examples.html
     '''
-    def __init__(self, select, bucket, access_key, secret_key, parallel='on'):
+    def __init__(self, select, unload_location, access_key, secret_key, session_token='', options={}):
         ''' Initializes an UnloadFromSelect instance
 
         Args:
             self: An instance of UnloadFromSelect
             select: The select statement to be unloaded
-            bucket: The Amazon S3 bucket where the result will be stored
-            access_key: The Amazon Access Key ID
-            secret_key: The Amazon Secret Access Key
-            parallel: If 'ON' the result will be written to multiple files. If
-                'OFF' the result will write to one (1) file up to 6.2GB before
-                splitting
+            unload_location: The Amazon S3 bucket where the result will be stored
+            access_key - AWS Access Key (required)
+            secret_key - AWS Secret Key (required)
+            session_token - AWS STS Session Token (optional)
+            options - Set of optional parameters to modify the UNLOAD sql
+                parallel: If 'ON' the result will be written to multiple files. If
+                    'OFF' the result will write to one (1) file up to 6.2GB before
+                    splitting
+                add_quotes: Boolean value for ADDQUOTES; defaults to True
+                delimiter - File delimiter. Defaults to ','
         '''
         self.select = select
-        self.bucket = bucket
+        self.unload_location = unload_location
         self.access_key = access_key
         self.secret_key = secret_key
-        self.parallel = parallel
+        self.session_token = session_token
+        self.options = options
 
 
 @compiles(UnloadFromSelect)
 def visit_unload_from_select(element, compiler, **kw):
     ''' Returns the actual sql query for the UnloadFromSelect class
+
+    '''
+    return """
+           UNLOAD ('%(query)s') TO '%(unload_location)s'
+           CREDENTIALS 'aws_access_key_id=%(access_key)s;aws_secret_access_key=%(secret_key)s%(session_token)s'
+           DELIMITER '%(delimiter)s'
+           %(add_quotes)s
+           ALLOWOVERWRITE
+           PARALLEL %(parallel)s;
+           """ % \
+           {'query': compiler.process(element.select, unload_select=True, literal_binds=True),
+            'unload_location': element.unload_location,
+            'access_key': element.access_key,
+            'secret_key': element.secret_key,
+            'session_token': ';token=%s' % element.session_token if element.session_token else '',
+            'add_quotes': 'ADDQUOTES' if bool(element.options.get('add_quotes', True)) else '',
+            'delimiter': element.options.get('delimiter', ','),
+            'parallel': element.options.get('parallel', 'ON')}
+
+
+class CopyCommand(Executable, ClauseElement):
+    ''' Prepares a RedShift COPY statement
+    '''
+    def __init__(self, schema_name, table_name, data_location, access_key, secret_key, session_token='', options={}):
+        ''' Initializes a CopyCommand instance
+
+        Args:
+            self: An instance of CopyCommand
+            schema_name - Schema associated with the table_name
+            table_name: The table to copy the data into
+            data_location The Amazon S3 location from where to copy
+            access_key - AWS Access Key (required)
+            secret_key - AWS Secret Key (required)
+            session_token - AWS STS Session Token (optional)
+            options - Set of optional parameters to modify the COPY sql
+                delimiter - File delimiter; defaults to ','
+                ignore_header - Integer value of number of lines to skip at the start of each file
+                null - String value denoting what to interpret as a NULL value from the file; defaults to '---'
+                empty_as_null - Boolean value denoting whether to load VARCHAR fields with
+                                empty values as NULL instead of empty string; defaults to True
+                blanks_as_null - Boolean value denoting whether to load VARCHAR fields with
+                                 whitespace only values as NULL instead of whitespace; defaults to True
+        '''
+        self.schema_name = schema_name
+        self.table_name = table_name
+        self.data_location = data_location
+        self.access_key = access_key
+        self.secret_key = secret_key
+        self.session_token = session_token
+        self.options = options
+
+
+@compiles(CopyCommand)
+def visit_copy_command(element, compiler, **kw):
+    ''' Returns the actual sql query for the CopyCommand class
     '''
-    return "unload ('%(query)s') to '%(bucket)s' credentials 'aws_access_key_id=%(access_key)s;aws_secret_access_key=%(secret_key)s' delimiter ',' addquotes allowoverwrite parallel %(parallel)s" % {
-        'query': compiler.process(element.select, unload_select=True, literal_binds=True),
-        'bucket': element.bucket,
-        'access_key': element.access_key,
-        'secret_key': element.secret_key,
-        'parallel': element.parallel,
-    }
+    return """
+           COPY %(schema_name)s.%(table_name)s FROM '%(data_location)s'
+           CREDENTIALS 'aws_access_key_id=%(access_key)s;aws_secret_access_key=%(secret_key)s%(session_token)s'
+           CSV
+           TRUNCATECOLUMNS
+           DELIMITER '%(delimiter)s'
+           IGNOREHEADER %(ignore_header)s
+           NULL '%(null)s'
+           %(empty_as_null)s
+           %(blanks_as_null)s;
+           """ % \
+           {'schema_name': element.schema_name,
+            'table_name': element.table_name,
+            'data_location': element.data_location,
+            'access_key': element.access_key,
+            'secret_key': element.secret_key,
+            'session_token': ';token=%s' % element.session_token if element.session_token else '',
+            'null': element.options.get('null', '---'),
+            'delimiter': element.options.get('delimiter', ','),
+            'ignore_header': element.options.get('ignore_header', 0),
+            'empty_as_null': 'EMPTYASNULL' if bool(element.options.get('empty_as_null', True)) else '',
+            'blanks_as_null': 'BLANKSASNULL' if bool(element.options.get('blanks_as_null', True)) else ''}
+
 
 @compiles(BindParameter)
 def visit_bindparam(bindparam, compiler, **kw):

setup.py

@@ -10,7 +10,7 @@ def run_tests(self):
 
 setup(
     name='redshift-sqlalchemy',
-    version='0.4.1',
+    version='0.5.0a',
     description='Amazon Redshift Dialect for sqlalchemy',
     long_description=open("README.rst").read(),
     author='Matt George',

tests/test_copy_command.py

@@ -0,0 +1,18 @@
+from unittest import TestCase
+from redshift_sqlalchemy.dialect import CopyCommand
+
+
+class TestCopyCommand(TestCase):
+    def setUp(self):
+        pass
+
+    def test_basic_copy_case(self):
+        ''' Tests that the simplest type of CopyCommand works
+        '''
+        expected_result = "COPY schema1.t1 FROM 's3://mybucket/data/listing/'\n" \
+                          "           CREDENTIALS 'aws_access_key_id=cookies;aws_secret_access_key=cookies'\n" \
+                          "           CSV\n           TRUNCATECOLUMNS\n           DELIMITER ','\n" \
+                          "           IGNOREHEADER 0\n           NULL '---'\n           EMPTYASNULL\n" \
+                          "           BLANKSASNULL;"
+        insert = CopyCommand('schema1', 't1', 's3://mybucket/data/listing/', 'cookies', 'cookies')
+        self.assertEqual(expected_result, str(insert).strip())

tests/test_unload_from_select.py

@@ -14,13 +14,21 @@ def setUp(self):
     def test_basic_unload_case(self):
         ''' Tests that the simplest type of UnloadFromSelect works
         '''
-        expected_result = "unload ('SELECT count(t1.id) AS count_1 \nFROM t1') to 'cookies' credentials 'aws_access_key_id=cookies;aws_secret_access_key=cookies' delimiter ',' addquotes allowoverwrite parallel on"
-        insert = UnloadFromSelect(select([func.count(self.t1.c.id)]), 'cookies', 'cookies', 'cookies')
-        self.assertEqual(expected_result, str(insert))
+        expected_result = "UNLOAD ('SELECT count(t1.id) AS count_1 \nFROM t1') TO 's3://bucket/key'\n           " \
+                          "CREDENTIALS 'aws_access_key_id=cookies;aws_secret_access_key=cookies;token=cookies'\n" \
+                          "           DELIMITER ','\n           ADDQUOTES\n           ALLOWOVERWRITE\n          " \
+                          " PARALLEL ON;"
+        insert = UnloadFromSelect(select([func.count(self.t1.c.id)]), 's3://bucket/key', 'cookies', 'cookies',
+                                  'cookies')
+        self.assertEqual(expected_result, str(insert).strip())
 
     def test_parallel_off_unload_case(self):
         ''' Tests that UnloadFromSelect handles parallel being set to off
         '''
-        expected_result = "unload ('SELECT count(t1.id) AS count_1 \nFROM t1') to 'cookies' credentials 'aws_access_key_id=cookies;aws_secret_access_key=cookies' delimiter ',' addquotes allowoverwrite parallel off"
-        insert = UnloadFromSelect(select([func.count(self.t1.c.id)]), 'cookies', 'cookies', 'cookies', 'off')
-        self.assertEqual(expected_result, str(insert))
+        expected_result = "UNLOAD ('SELECT count(t1.id) AS count_1 \nFROM t1') TO 's3://bucket/key'\n           " \
+                          "CREDENTIALS 'aws_access_key_id=cookies;aws_secret_access_key=cookies;token=cookies'\n" \
+                          "           DELIMITER ','\n           ADDQUOTES\n           ALLOWOVERWRITE\n          " \
+                          " PARALLEL OFF;"
+        insert = UnloadFromSelect(select([func.count(self.t1.c.id)]), 's3://bucket/key', 'cookies', 'cookies',
+                                  'cookies', {'parallel': 'OFF'})
+        self.assertEqual(expected_result, str(insert).strip())