[sqllab] bugfix where a query has the same alias twice as output (#1734)

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:

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):

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:

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()