diff --git a/superset/sql_lab.py b/superset/sql_lab.py index addbc7268f..6dcda0bcac 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -11,7 +11,7 @@ from sqlalchemy.pool import NullPool from sqlalchemy.orm import sessionmaker from superset import ( - app, db, models, utils, dataframe, results_backend, sql_parse, sm) + app, db, models, utils, dataframe, results_backend, sql_parse) from superset.db_engine_specs import LimitMethod from superset.jinja_context import get_template_processor QueryStatus = models.QueryStatus @@ -19,6 +19,27 @@ QueryStatus = models.QueryStatus celery_app = celery.Celery(config_source=app.config.get('CELERY_CONFIG')) +def dedup(l, suffix='__'): + """De-duplicates a list of string by suffixing a counter + + Always returns the same number of entries as provided, and always returns + unique values. + + >>> dedup(['foo', 'bar', 'bar', 'bar']) + ['foo', 'bar', 'bar__1', 'bar__2'] + """ + new_l = [] + seen = {} + for s in l: + if s in seen: + seen[s] += 1 + s += suffix + str(seen[s]) + else: + seen[s] = 0 + new_l.append(s) + return new_l + + def create_table_as(sql, table_name, schema=None, override=False): """Reformats the query into the create table as query. @@ -117,6 +138,7 @@ def get_sql_results(self, query_id, return_results=True, store_results=False): cdf = None if result_proxy.cursor: column_names = [col[0] for col in result_proxy.cursor.description] + column_names = dedup(column_names) if db_engine_spec.limit_method == LimitMethod.FETCH_MANY: data = result_proxy.fetchmany(query.limit) else: diff --git a/tests/base_tests.py b/tests/base_tests.py index 36f3d8f3a9..10e2e25cac 100644 --- a/tests/base_tests.py +++ b/tests/base_tests.py @@ -197,7 +197,7 @@ class SupersetTestCase(unittest.TestCase): perm.view_menu and table.perm in perm.view_menu.name): appbuilder.sm.del_permission_role(public_role, perm) - def run_sql(self, sql, client_id, user_name=None): + def run_sql(self, sql, client_id, user_name=None, raise_on_error=False): if user_name: self.logout() self.login(username=(user_name if user_name else 'admin')) @@ -208,6 +208,8 @@ class SupersetTestCase(unittest.TestCase): data=dict(database_id=dbid, sql=sql, select_as_create_as=False, client_id=client_id), ) + if raise_on_error and 'error' in resp: + raise Exception("run_sql failed") return resp def test_gamma_permissions(self): diff --git a/tests/core_tests.py b/tests/core_tests.py index 2c2267a3bd..3f73a975df 100644 --- a/tests/core_tests.py +++ b/tests/core_tests.py @@ -13,7 +13,7 @@ import unittest from flask import escape -from superset import db, models, utils, appbuilder, sm, jinja_context +from superset import db, models, utils, appbuilder, sm, jinja_context, sql_lab from superset.views import DatabaseView from .base_tests import SupersetTestCase @@ -165,7 +165,7 @@ class CoreTests(SupersetTestCase): assert escape(title) in self.client.get(url).data.decode('utf-8') def test_doctests(self): - modules = [utils, models] + modules = [utils, models, sql_lab] for mod in modules: failed, tests = doctest.testmod(mod) if failed: diff --git a/tests/sqllab_tests.py b/tests/sqllab_tests.py index 00c574dc74..2597017cd5 100644 --- a/tests/sqllab_tests.py +++ b/tests/sqllab_tests.py @@ -190,6 +190,13 @@ class SqlLabTests(SupersetTestCase): self.assertLess(int(first_query_time), v['startDttm']) self.assertLess(v['startDttm'], int(second_query_time)) + def test_alias_duplicate(self): + self.run_sql( + "SELECT username as col, id as col, username FROM ab_user", + client_id='2e2df3', + user_name='admin', + raise_on_error=True) + if __name__ == '__main__': unittest.main()