From 879b1c225584d30f1f4f03e0ae66ce5bb30c75d9 Mon Sep 17 00:00:00 2001
From: alexandrusoare <alexandrusoare98@gmail.com>
Date: Wed, 6 May 2026 11:28:03 +0300
Subject: [PATCH] taking care of PR comments

---
 superset/mcp_service/chart/tool/get_chart_data.py            | 2 +-
 .../unit_tests/mcp_service/chart/tool/test_get_chart_data.py | 5 ++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/superset/mcp_service/chart/tool/get_chart_data.py b/superset/mcp_service/chart/tool/get_chart_data.py
index 1978a0a61d5..d5f19c3dbd0 100644
--- a/superset/mcp_service/chart/tool/get_chart_data.py
+++ b/superset/mcp_service/chart/tool/get_chart_data.py
@@ -800,7 +800,7 @@ async def get_chart_data(  # noqa: C901
                 # Use SQL-derived GenericDataType when available,
                 # fall back to Python isinstance heuristic
                 data_type = "string"
-                if idx < len(coltypes):
+                if coltypes:
                     data_type = _GENERIC_TYPE_MAP.get(coltypes[idx], "string")
                 elif sample_values:
                     if all(isinstance(v, bool) for v in sample_values):
diff --git a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
index c289b127d8c..c54f42817a1 100644
--- a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
+++ b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
@@ -30,12 +30,14 @@ from superset.mcp_service.chart.schemas import (
     PerformanceMetadata,
 )
 from superset.mcp_service.chart.tool.get_chart_data import (
+    _GENERIC_TYPE_MAP,
     _MAX_RECOMMENDATIONS,
     _recommend_visualizations,
     _sanitize_chart_data_for_llm_context,
 )
 from superset.mcp_service.utils import sanitize_for_llm_context
 from superset.mcp_service.utils.sanitization import LLM_CONTEXT_ESCAPED_CLOSE_DELIMITER
+from superset.utils.core import GenericDataType
 
 
 def _collect_groupby_extras(
@@ -1099,9 +1101,6 @@ def test_recommend_single_numeric_high_cardinality_suggests_histogram():
 
 def test_coltypes_populates_data_type():
     """Verify that GenericDataType values from coltypes are mapped correctly."""
-    from superset.mcp_service.chart.tool.get_chart_data import _GENERIC_TYPE_MAP
-    from superset.utils.core import GenericDataType
-
     assert _GENERIC_TYPE_MAP[GenericDataType.NUMERIC] == "numeric"
     assert _GENERIC_TYPE_MAP[GenericDataType.STRING] == "string"
     assert _GENERIC_TYPE_MAP[GenericDataType.TEMPORAL] == "temporal"