diff --git a/redash/query_runner/__init__.py b/redash/query_runner/__init__.py index 361b7863d..9e7539042 100644 --- a/redash/query_runner/__init__.py +++ b/redash/query_runner/__init__.py @@ -1,4 +1,5 @@ import logging +from collections import defaultdict from contextlib import ExitStack from functools import wraps @@ -213,14 +214,14 @@ class BaseQueryRunner(object): def fetch_columns(self, columns): column_names = set() - duplicates_counter = 1 + duplicates_counters = defaultdict(int) 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 + while column_name in column_names: + duplicates_counters[col[0]] += 1 + column_name = "{}{}".format(col[0], duplicates_counters[col[0]]) column_names.add(column_name) new_columns.append({"name": column_name, "friendly_name": column_name, "type": col[1]}) diff --git a/tests/query_runner/test_basequeryrunner.py b/tests/query_runner/test_basequeryrunner.py new file mode 100644 index 000000000..09b741933 --- /dev/null +++ b/tests/query_runner/test_basequeryrunner.py @@ -0,0 +1,34 @@ +import unittest + +from redash.query_runner import BaseQueryRunner + + +class TestBaseQueryRunner(unittest.TestCase): + def setUp(self): + self.query_runner = BaseQueryRunner({}) + + def test_duplicate_column_names_assigned_correctly(self): + original_column_names = [ + ("name", bool), + ("created_at", bool), + ("updated_at", bool), + ("name", bool), + ("created_at", bool), + ("updated_at", bool), + ] + expected = [ + {"name": "name", "friendly_name": "name", "type": bool}, + {"name": "created_at", "friendly_name": "created_at", "type": bool}, + {"name": "updated_at", "friendly_name": "updated_at", "type": bool}, + {"name": "name1", "friendly_name": "name1", "type": bool}, + {"name": "created_at1", "friendly_name": "created_at1", "type": bool}, + {"name": "updated_at1", "friendly_name": "updated_at1", "type": bool}, + ] + + new_columns = self.query_runner.fetch_columns(original_column_names) + + self.assertEqual(new_columns, expected) + + +if __name__ == "__main__": + unittest.main()