More options for UnloadCommand and more resilient test cases

Author: Robert Engel

redshift_sqlalchemy/__init__.py

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

redshift_sqlalchemy/dialect.py

@@ -195,6 +195,7 @@ def __init__(self, select, unload_location, access_key, secret_key, session_toke
                     '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
@@ -215,6 +216,7 @@ def visit_unload_from_select(element, compiler, **kw):
            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;
            """ % \
@@ -224,6 +226,7 @@ def visit_unload_from_select(element, compiler, **kw):
             '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')}
 
@@ -245,7 +248,7 @@ def __init__(self, schema_name, table_name, data_location, access_key, secret_ke
             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 '---'
+                null - Optional string value denoting what to interpret as a NULL value from the file
                 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
@@ -271,7 +274,7 @@ def visit_copy_command(element, compiler, **kw):
            TRUNCATECOLUMNS
            DELIMITER '%(delimiter)s'
            IGNOREHEADER %(ignore_header)s
-           NULL '%(null)s'
+           %(null)s
            %(empty_as_null)s
            %(blanks_as_null)s;
            """ % \
@@ -281,7 +284,7 @@ def visit_copy_command(element, compiler, **kw):
             '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', '---'),
+            '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),
             'empty_as_null': 'EMPTYASNULL' if bool(element.options.get('empty_as_null', True)) else '',

setup.py

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

tests/test_copy_command.py

@@ -1,18 +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 = "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 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,33 +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 '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',
+            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')
-        self.assertEqual(expected_result, str(insert).strip())
 
-    def test_parallel_off_unload_case(self):
-        ''' Tests that UnloadFromSelect handles parallel being set to off
+        unload_str = re.sub(r'\s+', ' ', str(unload)).strip()
+
+        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 '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())
+        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)