Repository: sqlalchemy-redshift/sqlalchemy-redshift

Issue:
UNLOAD now supports Parquet
As of December 2019 Redshift [now supports](https://aws.amazon.com/about-aws/whats-new/2019/12/announcing-amazon-redshift-data-lake-export/) the Parquet format in the `UNLOAD` command. Currently the implementation of `UNLOAD` throws an exception if it receives anything but CSV for the export format. We need to allow Parquet now as well.

Discussion:

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 0c78ae9..b39c53d 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -1,7 +1,8 @@
 0.7.7 (unreleased)
 ------------------
 
-- Nothing changed yet.
+- Add support for Parquet format in ``UNLOAD`` command
+  (`Issue #187 <https://github.com/sqlalchemy-redshift/sqlalchemy-redshift/issues/187>`_)
 
 
 0.7.6 (2020-01-17)
diff --git a/sqlalchemy_redshift/commands.py b/sqlalchemy_redshift/commands.py
index 583fe85..5fd65e4 100644
--- a/sqlalchemy_redshift/commands.py
+++ b/sqlalchemy_redshift/commands.py
@@ -1,8 +1,11 @@
-import collections
 import enum
 import numbers
 import re
 import warnings
+try:
+    from collections.abc import Iterable
+except ImportError:
+    from collections import Iterable
 
 import sqlalchemy as sa
 from sqlalchemy import exc as sa_exc
@@ -307,8 +310,20 @@ def visit_unload_from_select(element, compiler, **kw):
         if el.delimiter is not None or el.fixed_width is not None:
             raise ValueError(
                 'CSV format cannot be used with delimiter or fixed_width')
+    elif el.format == Format.parquet:
+        format_ = 'FORMAT AS {}'.format(el.format.value)
+        if any((
+            el.delimiter, el.fixed_width, el.add_quotes, el.escape, el.null,
+            el.header, el.gzip
+        )):
+            raise ValueError(
+                "Parquet format can't be used with `delimiter`, `fixed_width`,"
+                ' `add_quotes`, `escape`, `null`, `header`, or `gzip`.'
+            )
     else:
-        raise ValueError('Only CSV format is currently supported')
+        raise ValueError(
+            'Only CSV and Parquet formats are currently supported.'
+        )
 
     qs = template.format(
         manifest='MANIFEST' if el.manifest else '',
@@ -581,7 +596,7 @@ class CopyCommand(_ExecutableClause):
 
         table = None
         columns = []
-        if isinstance(to, collections.Iterable):
+        if isinstance(to, Iterable):
             for column in to:
                 if table is not None and table != column.table:
                     raise ValueError(


Test patch:
diff --git a/tests/test_unload_from_select.py b/tests/test_unload_from_select.py
index a8d4ddd..96d0724 100644
--- a/tests/test_unload_from_select.py
+++ b/tests/test_unload_from_select.py
@@ -1,3 +1,4 @@
+import pytest
 import sqlalchemy as sa
 
 from sqlalchemy_redshift import commands
@@ -157,7 +158,7 @@ def test_all_redshift_options_with_header():
     assert clean(compile_query(unload)) == clean(expected_result)
 
 
-def test_csv_format():
+def test_csv_format__basic():
     """Tests that UnloadFromSelect uses the format option correctly."""
 
     unload = dialect.UnloadFromSelect(
@@ -176,3 +177,69 @@ def test_csv_format():
         """.format(creds=creds)
 
     assert clean(compile_query(unload)) == clean(expected_result)
+
+
+@pytest.mark.parametrize('delimiter,fixed_width', (
+    ('\t', None),
+    (None, 'id:8,name:32'),
+    (';', 'id:8,name:32'),
+))
+def test_csv_format__bad_options_crash(delimiter, fixed_width):
+    """Test that UnloadFromSelect crashes if you try to use DELIMITER and/or
+    FIXEDWIDTH with the CSV format.
+    """
+    unload = dialect.UnloadFromSelect(
+        select=sa.select([sa.func.count(table.c.id)]),
+        unload_location='s3://bucket/key',
+        access_key_id=access_key_id,
+        secret_access_key=secret_access_key,
+        format=commands.Format.csv,
+        delimiter=delimiter,
+        fixed_width=fixed_width
+    )
+
+    with pytest.raises(ValueError):
+        compile_query(unload)
+
+
+def test_parquet_format__basic():
+    """Basic successful test of unloading with the Parquet format."""
+    unload = dialect.UnloadFromSelect(
+        select=sa.select([sa.func.count(table.c.id)]),
+        unload_location='s3://bucket/key',
+        access_key_id=access_key_id,
+        secret_access_key=secret_access_key,
+        format=commands.Format.parquet,
+    )
+
+    expected_result = """
+        UNLOAD ('SELECT count(t1.id) AS count_1 FROM t1')
+        TO 's3://bucket/key'
+        CREDENTIALS '{creds}'
+        FORMAT AS PARQUET
+    """.format(creds=creds)
+
+    assert clean(compile_query(unload)) == clean(expected_result)
+
+
+@pytest.mark.parametrize('kwargs', (
+    {'delimiter': '\t'},
+    {'fixed_width': 'id:8,name:32'},
+    {'gzip': True},
+    {'add_quotes': True, 'escape': True},
+    {'null': '\\N'},
+    {'header': True},
+))
+def test_parquet_format__bad_options_crash(kwargs):
+    """Verify we crash if we use the Parquet format with a bad option."""
+    unload = dialect.UnloadFromSelect(
+        select=sa.select([sa.func.count(table.c.id)]),
+        unload_location='s3://bucket/key',
+        access_key_id=access_key_id,
+        secret_access_key=secret_access_key,
+        format=commands.Format.parquet,
+        **kwargs
+    )
+
+    with pytest.raises(ValueError):
+        compile_query(unload)