Skip to content

Commit 17b9a0a

Browse files
committed
move format_schema function for athena/presto
1 parent b7a7b67 commit 17b9a0a

File tree

3 files changed

+21
-15
lines changed

3 files changed

+21
-15
lines changed

redash/query_runner/__init__.py

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -178,18 +178,3 @@ def import_query_runners(query_runner_imports):
178178
for runner_import in query_runner_imports:
179179
__import__(runner_import)
180180

181-
182-
# for athena and presto
183-
def format_schema(results):
184-
schema = {}
185-
for row in results['rows']:
186-
table_name = '{}.{}'.format(row['table_schema'], row['table_name'])
187-
if table_name not in schema:
188-
schema[table_name] = {'name': table_name, 'columns': []}
189-
190-
row_to_add = row['column_name'] + ' (' + row['column_type'] + ')'
191-
if row['extra_info'] == 'partition key':
192-
row_to_add = '[P] ' + row_to_add
193-
schema[table_name]['columns'].append(row_to_add)
194-
195-
return schema

redash/query_runner/athena.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from redash.query_runner import *
66
from redash.settings import parse_boolean
77
from redash.utils import JSONEncoder
8+
from .presto import format_schema
89

910
logger = logging.getLogger(__name__)
1011
ANNOTATE_QUERY = parse_boolean(os.environ.get('ATHENA_ANNOTATE_QUERY', 'true'))

redash/query_runner/presto.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,26 @@
3131
}
3232

3333

34+
# for athena and presto
35+
def format_schema(results):
36+
"""
37+
This function formats the schema, table, and columsn of Athena and Presto
38+
for display in the UI schema browser.
39+
"""
40+
schema = {}
41+
for row in results['rows']:
42+
table_name = '{}.{}'.format(row['table_schema'], row['table_name'])
43+
if table_name not in schema:
44+
schema[table_name] = {'name': table_name, 'columns': []}
45+
46+
row_to_add = row['column_name'] + ' (' + row['column_type'] + ')'
47+
if row['extra_info'] == 'partition key':
48+
row_to_add = '[P] ' + row_to_add
49+
schema[table_name]['columns'].append(row_to_add)
50+
51+
return schema
52+
53+
3454
class Presto(BaseQueryRunner):
3555
noop_query = 'SHOW TABLES'
3656

0 commit comments

Comments
 (0)