From 81348455528ca8a35f70c3424d56c6a4d458a6fc Mon Sep 17 00:00:00 2001
From: Allen Short <ashort@mozilla.com>
Date: Sat, 9 Dec 2017 05:48:56 +0000
Subject: [PATCH] add partition key marker to Athena and Presto columns (re
 #185)

---
 redash/query_runner/athena.py | 14 ++++++++++++--
 redash/query_runner/presto.py | 16 ++++++++++++++--
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/redash/query_runner/athena.py b/redash/query_runner/athena.py
index 6437af413d..2f3bcf2b45 100644
--- a/redash/query_runner/athena.py
+++ b/redash/query_runner/athena.py
@@ -156,9 +156,10 @@ def get_schema(self, get_stats=False):
 
         schema = {}
         query = """
-        SELECT table_schema, table_name, column_name, data_type as column_type
+        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')
+        ORDER BY 1, 5 DESC
         """
 
         results, error = self.run_query(query, None)
@@ -170,7 +171,16 @@ def get_schema(self, get_stats=False):
             table_name = '{0}.{1}'.format(row['table_schema'], row['table_name'])
             if table_name not in schema:
                 schema[table_name] = {'name': table_name, 'columns': []}
-            schema[table_name]['columns'].append(row['column_name'] + ' (' + row['column_type'] + ')')
+
+            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' 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'])
+
 
         return schema.values()
 
diff --git a/redash/query_runner/presto.py b/redash/query_runner/presto.py
index 509e5b9b56..ea94108394 100644
--- a/redash/query_runner/presto.py
+++ b/redash/query_runner/presto.py
@@ -1,4 +1,5 @@
 import json
+from markupsafe import Markup, escape
 
 from redash.utils import JSONEncoder
 from redash.query_runner import *
@@ -84,9 +85,10 @@ 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
+        SELECT table_schema, table_name, column_name, data_type as column_type, extra_info
         FROM information_schema.columns
         WHERE table_schema NOT IN ('pg_catalog', 'information_schema')
+        ORDER BY 1, 5 DESC
         """
 
         results, error = self.run_query(query, None)
@@ -102,7 +104,14 @@ def get_schema(self, get_stats=False):
             if table_name not in schema:
                 schema[table_name] = {'name': table_name, 'columns': []}
 
-            schema[table_name]['columns'].append(row['column_name'] + ' (' + row['column_type'] + ')')
+            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' 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'])
 
         return schema.values()
 
@@ -122,6 +131,9 @@ def run_query(self, query, user):
             column_tuples = [(i[0], PRESTO_TYPES_MAPPING.get(i[1], None)) for i in cursor.description]
             columns = self.fetch_columns(column_tuples)
             rows = [dict(zip(([c['name'] for c in columns]), r)) for i, r in enumerate(cursor.fetchall())]
+            for row in rows:
+                for field in row:
+                    field = escape(field)
             data = {'columns': columns, 'rows': rows}
             json_data = json.dumps(data, cls=JSONEncoder)
             error = None