Merge pull request #62 from keboola/PS-2017-table-columns

PS-2017 fix header columns

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'))

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))