Merge pull request #576 from joaofraga/fix/deduplicate-column-names

Fix: support for duplicate columns for MySQL query runner
getredash · Oct 1, 2015 · 5d1c75d · 5d1c75d
2 parents 1052528 + 65cc67d
commit 5d1c75d
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 31 deletions.
diff --git a/redash/query_runner/__init__.py b/redash/query_runner/__init__.py
@@ -67,6 +67,24 @@ def configuration_schema(cls):
     def run_query(self, query):
         raise NotImplementedError()
 
+    def fetch_columns(self, columns):
+        column_names = []
+        duplicates_counter = 1
+        new_columns = []
+
+        for col in columns:
+            column_name = col[0]
+            if column_name in column_names:
+                column_name = "{}{}".format(column_name, duplicates_counter)
+                duplicates_counter += 1
+
+            column_names.append(column_name)
+            new_columns.append({'name': column_name,
+                                'friendly_name': column_name,
+                                'type': col[1]})
+
+        return new_columns
+
     def get_schema(self):
         return []
 

diff --git a/redash/query_runner/mysql.py b/redash/query_runner/mysql.py
@@ -119,13 +119,8 @@ def run_query(self, query):
 
             # TODO - very similar to pg.py
             if cursor.description is not None:
-                columns_data = [(i[0], i[1]) for i in cursor.description]
-
-                rows = [dict(zip((c[0] for c in columns_data), row)) for row in data]
-
-                columns = [{'name': col[0],
-                            'friendly_name': col[0],
-                            'type': types_map.get(col[1], None)} for col in columns_data]
+                columns = self.fetch_columns([(i[0], types_map.get(i[1], None)) for i in cursor.description])
+                rows = [dict(zip((c['name'] for c in columns), row)) for row in data]
 
                 data = {'columns': columns, 'rows': rows}
                 json_data = json.dumps(data, cls=JSONEncoder)
@@ -149,4 +144,4 @@ def run_query(self, query):
 
         return json_data, error
 
-register(Mysql)
+register(Mysql)
diff --git a/redash/query_runner/pg.py b/redash/query_runner/pg.py
@@ -121,29 +121,9 @@ def run_query(self, query):
             cursor.execute(query)
             _wait(connection)
 
-            # While set would be more efficient here, it sorts the data which is not what we want, but due to the small
-            # size of the data we can assume it's ok.
-            column_names = []
-            columns = []
-            duplicates_counter = 1
-
             if cursor.description is not None:
-                for column in cursor.description:
-                    # TODO: this deduplication needs to be generalized and reused in all query runners.
-                    column_name = column.name
-                    if column_name in column_names:
-                        column_name += str(duplicates_counter)
-                        duplicates_counter += 1
-
-                    column_names.append(column_name)
-
-                    columns.append({
-                        'name': column_name,
-                        'friendly_name': column_name,
-                        'type': types_map.get(column.type_code, None)
-                    })
-
-                rows = [dict(zip(column_names, row)) for row in cursor]
+                columns = self.fetch_columns([(i[0], types_map.get(i[1], None)) for i in cursor.description])
+                rows = [dict(zip((c['name'] for c in columns), row)) for row in cursor]
 
                 data = {'columns': columns, 'rows': rows}
                 error = None
@@ -170,4 +150,4 @@ def run_query(self, query):
 
         return json_data, error
 
-register(PostgreSQL)
+register(PostgreSQL)