Merge pull request #236 from alison985/m10_rc

Athena array(), map() or row()

redash/query_runner/athena.py

@@ -122,7 +122,7 @@ def __init__(self, configuration):
     def get_schema(self, get_stats=False):
         schema = {}
         query = """
-        SELECT table_schema, table_name, column_name, data_type as column_type, extra_info
+        SELECT table_schema, table_name, column_name, data_type as column_type, comment as extra_info
         FROM information_schema.columns
         WHERE table_schema NOT IN ('information_schema')
         """
@@ -137,12 +137,12 @@ def get_schema(self, get_stats=False):
             if table_name not in schema:
                 schema[table_name] = {'name': table_name, 'columns': []}
 
-            if row['extra_info'] == 'partition key':
+            if row['extra_info'] == 'Partition Key':
                 schema[table_name]['columns'].append('[P] ' + row['column_name'] + ' (' + row['column_type'] + ')')
             elif row['column_type'] == 'integer' or row['column_type'] == 'varchar' or row['column_type'] == 'timestamp' or row['column_type'] == 'boolean' or row['column_type'] == 'bigint':
                 schema[table_name]['columns'].append(row['column_name'] + ' (' + row['column_type'] + ')')
-            elif row['column_type'][0:2] == 'row' and row['column_type'][0:2] == 'map':
-                schema[table_name]['columns'].append(row['column_name'] + ' (row or map)')
+            elif row['column_type'][0:2] == 'row' or row['column_type'][0:2] == 'map' or row['column_type'][0:2] == 'arr':
+                schema[table_name]['columns'].append(row['column_name'] + ' (row or map or array)')
             else:
                 schema[table_name]['columns'].append(row['column_name'])