From d7a089764f5f03a51bd2cfe0986b6869d2dedf90 Mon Sep 17 00:00:00 2001
From: pivnicek <marc@keboola.com>
Date: Mon, 12 Apr 2021 12:04:44 +0200
Subject: [PATCH 1/2] test for column header

---
 tests/functional/test_tables.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/tests/functional/test_tables.py b/tests/functional/test_tables.py
index cd0ffc5..cefa1b3 100644
--- a/tests/functional/test_tables.py
+++ b/tests/functional/test_tables.py
@@ -307,3 +307,23 @@ def test_table_export_sliced(self):
         self.assertEqual(['"col1","col2"\n', '"foo","bar"\n',
                           '"ping","pong"\n'],
                          sorted(lines))
+
+    def test_table_columns(self):
+        file, path = tempfile.mkstemp(prefix='sapi-test')
+        with open(path, 'w') as csv_file:
+            writer = csv.DictWriter(csv_file, fieldnames=['col1', 'col2', 'col3', 'col4'],
+                                    lineterminator='\n', delimiter=',',
+                                    quotechar='"')
+            writer.writeheader()
+            writer.writerow({'col1': 'ping', 'col2': 'pong', 'col3': 'king', 'col4': 'kong'})
+        os.close(file)
+        table_id = self.tables.create(name='some-table', file_path=path, bucket_id='in.c-py-test-tables')
+        temp_path = tempfile.TemporaryDirectory()
+        local_path = self.tables.export_to_file(table_id=table_id,
+                                                path_name=temp_path.name,
+                                                is_gzip=False,
+                                                columns=['col3', 'col2'])
+
+        with open(local_path, mode='rt') as file:
+            lines = file.readlines()
+        self.assertEqual(['"col3","col2"\n', '"king","pong"\n'], sorted(lines))

From fc95ff8914001d29a0df4e7466f34fffd080a0bc Mon Sep 17 00:00:00 2001
From: pivnicek <marc@keboola.com>
Date: Mon, 12 Apr 2021 12:06:45 +0200
Subject: [PATCH 2/2] use columns if given

---
 kbcstorage/tables.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/kbcstorage/tables.py b/kbcstorage/tables.py
index 50124ce..324c459 100644
--- a/kbcstorage/tables.py
+++ b/kbcstorage/tables.py
@@ -419,7 +419,8 @@ def export_to_file(self, table_id, path_name, limit=None,
 
         with open(local_file, mode='rb') as in_file, \
                 open(destination_file, mode='wb') as out_file:
-            columns = table_detail['columns']
+            if columns is None:
+                columns = table_detail['columns']
             columns = ['"{}"'.format(col) for col in columns]
             header = ",".join(columns) + '\n'
             out_file.write(header.encode('utf-8'))