Added UnloadFromSelect and associated tests. Also added build/ to .gitignore

Signed-off-by: Jason Myers <[email protected]>

Author: jasonamyers

.gitignore

@@ -3,3 +3,4 @@
 .DS_Store
 dist
 *.egg-info
+build

redshift_sqlalchemy/dialect.py

@@ -2,6 +2,9 @@
 from sqlalchemy.engine import reflection
 from sqlalchemy import util, exc
 from sqlalchemy.types import VARCHAR, NullType
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.expression import Executable, ClauseElement
+
 
 class RedshiftDialect(PGDialect_psycopg2):
     @reflection.cache
@@ -51,3 +54,20 @@ def _get_column_info(self, name, format_type, default,
         if isinstance(column_info['type'], VARCHAR) and column_info['type'].length is None:
             column_info['type'] = NullType()
         return column_info
+
+
+class UnloadFromSelect(Executable, ClauseElement):
+    def __init__(self, select, bucket, access_key, secret_key):
+        self.select = select
+        self.bucket = bucket
+        self.access_key = access_key
+        self.secret_key = secret_key
+
+@compiles(UnloadFromSelect)
+def visit_unload_from_select(element, compiler, **kw):
+    return "unload ('%(query)s') to '%(bucket)s' credentials 'aws_access_key_id=%(access_key)s;aws_secret_access_key=%(secret_key)s'" % {
+        'query': element.select,
+        'bucket': element.bucket,
+        'access_key': element.access_key,
+        'secret_key': element.secret_key,
+    }

tests/test_unload_from_select.py

@@ -0,0 +1,18 @@
+from unittest import TestCase
+from sqlalchemy import Table, Column, Integer, String, MetaData, ForeignKey, create_engine
+from sqlalchemy.sql import select, func
+from sqlalchemy.types import NullType, VARCHAR
+from redshift_sqlalchemy.dialect import RedshiftDialect, UnloadFromSelect, visit_unload_from_select
+
+
+class TestUnloadFromSelect(TestCase):
+    def setUp(self):
+        self.metadata = MetaData()
+        self.t1 = Table('t1', self.metadata, Column('id', Integer, primary_key=True), Column('name', String))
+        self.engine = create_engine('sqlite:///:memory:', echo=True)
+        self.metadata.create_all(self.engine)
+
+    def test_basic_unload_case(self):
+        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'"
+        insert = UnloadFromSelect(select([func.count(self.t1.c.id)]), 'cookies', 'cookies', 'cookies')
+        self.assertEqual(expected_result, str(insert))