add partition key marker to partition columns

getredash#185 added to the front of the column name so they are in a consistent ordinal position for each table tests pass
washort · Aug 14, 2017 · f4d5830 · f4d5830
1 parent 84bbae9
commit f4d5830
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/redash/query_runner/athena.py b/redash/query_runner/athena.py
@@ -129,7 +129,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
+        SELECT table_schema, table_name, column_name, data_type as column_type, extra_info
         FROM information_schema.columns
         WHERE table_schema NOT IN ('information_schema')
         """
@@ -143,7 +143,12 @@ 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'] + ')')
+            else:
+                schema[table_name]['columns'].append(row['column_name'] + ' (' + row['column_type'] + ')')
+
 
         return schema.values()
 

diff --git a/redash/query_runner/presto.py b/redash/query_runner/presto.py
@@ -86,7 +86,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
+        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')
         """
@@ -104,7 +104,10 @@ 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'] + ')')
+            else:
+                schema[table_name]['columns'].append(row['column_name'] + ' (' + row['column_type'] + ')')
 
         return schema.values()