From 11a40a417ecb70d3fa8b121f4cfb257cefc248cc Mon Sep 17 00:00:00 2001
From: Mayur <mayurnewase111@gmail.com>
Date: Mon, 10 Oct 2022 09:38:33 +0530
Subject: [PATCH] fix: allow adhoc columns in non-aggregate query (#21729)

---
 superset/superset_typing.py | 2 ++
 superset/utils/core.py      | 9 +++++++++
 2 files changed, 11 insertions(+)

diff --git a/superset/superset_typing.py b/superset/superset_typing.py
index 00b76cd78caef..8eaea54176ed3 100644
--- a/superset/superset_typing.py
+++ b/superset/superset_typing.py
@@ -55,6 +55,8 @@ class AdhocColumn(TypedDict, total=False):
     hasCustomLabel: Optional[bool]
     label: str
     sqlExpression: str
+    columnType: Optional[Literal["BASE_AXIS", "SERIES"]]
+    timeGrain: Optional[str]
 
 
 class ResultSetColumnType(TypedDict):
diff --git a/superset/utils/core.py b/superset/utils/core.py
index 84ab5256b8d9b..76e566a895f86 100644
--- a/superset/utils/core.py
+++ b/superset/utils/core.py
@@ -1258,6 +1258,15 @@ def is_adhoc_column(column: Column) -> TypeGuard[AdhocColumn]:
     )
 
 
+def get_base_axis_labels(columns: Optional[List[Column]]) -> Tuple[str, ...]:
+    axis_cols = [
+        col
+        for col in columns or []
+        if is_adhoc_column(col) and col.get("columnType") == "BASE_AXIS"
+    ]
+    return tuple(get_column_name(col) for col in axis_cols)
+
+
 def get_column_name(
     column: Column, verbose_map: Optional[Dict[str, Any]] = None
 ) -> str: