Merge pull request #22 from robertengel/extended_unload_and_copy_commands

Extended unload and copy commands

Author: binarymatt

redshift_sqlalchemy/__init__.py

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

redshift_sqlalchemy/dialect.py

@@ -111,6 +111,7 @@ def _fetch_redshift_column_attributes(self, column):
             text += " SORTKEY"
         return text
 
+
 class RedshiftDialect(PGDialect_psycopg2):
     name = 'redshift'
     ddl_compiler = RedShiftDDLCompiler
@@ -182,37 +183,119 @@ 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
+                null_as: optional string that represents a null value in unload output
+                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
+           %(null_as)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 '',
+            'null_as': ("NULL '%s'" % element.options.get('null_as')) if element.options.get('null_as') 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 - or a manifest file if 'manifest' option is used
+            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 - Optional string value denoting what to interpret as a NULL value from the file
+                manifest - Boolean value denoting whether data_location is a manifest file; defaults to False
+                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)s
+           %(manifest)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': ("NULL '%s'" % element.options.get('null')) if element.options.get('null') else '',
+            'delimiter': element.options.get('delimiter', ','),
+            'ignore_header': element.options.get('ignore_header', 0),
+            'manifest': 'MANIFEST' if bool(element.options.get('manifest', False)) else '',
+            '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.1a',
     description='Amazon Redshift Dialect for sqlalchemy',
     long_description=open("README.rst").read(),
     author='Matt George',

tests/test_copy_command.py

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

tests/test_unload_from_select.py

@@ -2,25 +2,46 @@
 from sqlalchemy import Table, Column, Integer, String, MetaData
 from sqlalchemy.sql import select, func
 from redshift_sqlalchemy.dialect import UnloadFromSelect
+import re
 
 
 class TestUnloadFromSelect(TestCase):
+
     def setUp(self):
-        ''' Sets up a table and associate meta data for the test queries to build against
+        '''
+            Sets up a table and associate meta data for the test queries to build against
         '''
         self.metadata = MetaData()
         self.t1 = Table('t1', self.metadata, Column('id', Integer, primary_key=True), Column('name', String))
 
     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))
+            Tests that the simplest type of UnloadFromSelect works
+        '''
+        expected_result = re.sub(r'\s+', ' ',
+                                 "UNLOAD ('SELECT count(t1.id) AS count_1 \nFROM t1') TO 's3://bucket/key' "
+                                 "CREDENTIALS 'aws_access_key_id=cookies;aws_secret_access_key=cookies;token=cookies' "
+                                 "DELIMITER ',' ADDQUOTES ALLOWOVERWRITE PARALLEL ON;").strip()
+
+        unload = UnloadFromSelect(select([func.count(self.t1.c.id)]), 's3://bucket/key', 'cookies', 'cookies',
+                                  'cookies')
+
+        unload_str = re.sub(r'\s+', ' ', str(unload)).strip()
 
-    def test_parallel_off_unload_case(self):
-        ''' Tests that UnloadFromSelect handles parallel being set to off
+        self.assertEqual(expected_result, unload_str)
+
+    def test_unload_with_options(self):
+        '''
+            Tests that UnloadFromSelect handles options correctly
         '''
-        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 = re.sub(r'\s+', ' ',
+                                 "UNLOAD ('SELECT count(t1.id) AS count_1 \nFROM t1') TO 's3://bucket/key' "
+                                 "CREDENTIALS 'aws_access_key_id=cookies;aws_secret_access_key=cookies;token=cookies' "
+                                 "DELIMITER ',' ADDQUOTES NULL '---' ALLOWOVERWRITE PARALLEL OFF;").strip()
+
+        unload = UnloadFromSelect(select([func.count(self.t1.c.id)]), 's3://bucket/key', 'cookies', 'cookies',
+                                  'cookies', {'parallel': 'OFF', 'null_as': '---'})
+
+        unload_str = re.sub(r'\s+', ' ', str(unload)).strip()
+
+        self.assertEqual(expected_result, unload_str)