Merge pull request #15 from jasonamyers/feature/parallel

Adding option for parallel to unload command

Author: binarymatt

redshift_sqlalchemy/dialect.py

@@ -179,7 +179,7 @@ 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):
+    def __init__(self, select, bucket, access_key, secret_key, parallel='on'):
         ''' Initializes an UnloadFromSelect instance
 
         Args:
@@ -188,22 +188,27 @@ def __init__(self, select, bucket, access_key, secret_key):
             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
         '''
         self.select = select
         self.bucket = bucket
         self.access_key = access_key
         self.secret_key = secret_key
+        self.parallel = parallel
 
 
 @compiles(UnloadFromSelect)
 def visit_unload_from_select(element, compiler, **kw):
     ''' Returns the actual sql query for the UnloadFromSelect 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" % {
+    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,
     }
 
 @compiles(BindParameter)

tests/test_unload_from_select.py

@@ -14,6 +14,13 @@ 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"
+        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))
+
+    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))