Base duplicate index on column names (#4600)

* increase duplicate column names based on the original column name and not the number of duplicates detected

* add fetch columns test for base query runner

---------

Co-authored-by: Omer Lachish <[email protected]>
Co-authored-by: Guido Petri <[email protected]>

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 @@ def run_query(self, query, user):
 
     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]})

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