From 2d6f5b091cd4072f433f02ea2cd8f2b691017f30 Mon Sep 17 00:00:00 2001
From: Omer Lachish <289488+rauchy@users.noreply.github.com>
Date: Thu, 24 Aug 2023 07:18:05 +0300
Subject: [PATCH] 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 <omer@rauchy.net>
Co-authored-by: Guido Petri <18634426+guidopetri@users.noreply.github.com>
---
 redash/query_runner/__init__.py            |  9 +++---
 tests/query_runner/test_basequeryrunner.py | 34 ++++++++++++++++++++++
 2 files changed, 39 insertions(+), 4 deletions(-)
 create mode 100644 tests/query_runner/test_basequeryrunner.py

diff --git a/redash/query_runner/__init__.py b/redash/query_runner/__init__.py
index 361b7863df..9e75390422 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 @@ 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]})
diff --git a/tests/query_runner/test_basequeryrunner.py b/tests/query_runner/test_basequeryrunner.py
new file mode 100644
index 0000000000..09b7419333
--- /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()