From a587ca25f7bd2c7fc0d637cbe2c2f57736c9552b Mon Sep 17 00:00:00 2001
From: Alex Fedotyev <alex.fedotyev@clickhouse.com>
Date: Tue, 14 Apr 2026 19:01:25 -0700
Subject: [PATCH] fix: restrict heatmap data source picker to trace sources
 only

Heatmap charts require trace data (duration expression) to work properly.
Non-trace sources would error out with no guidance. Instead of adding
complex error messaging, filter the source dropdown to only show trace
sources when heatmap display type is selected.
---
 .../ChartEditorControls.tsx                   |  5 +++++
 .../DBEditTimeChartForm/EditTimeChartForm.tsx | 19 +++++++------------
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/packages/app/src/components/DBEditTimeChartForm/ChartEditorControls.tsx b/packages/app/src/components/DBEditTimeChartForm/ChartEditorControls.tsx
index acb2a9b7..9cc73983 100644
--- a/packages/app/src/components/DBEditTimeChartForm/ChartEditorControls.tsx
+++ b/packages/app/src/components/DBEditTimeChartForm/ChartEditorControls.tsx
@@ -99,6 +99,11 @@ export function ChartEditorControls({
             control={control}
             name="source"
             data-testid="source-selector"
+            allowedSourceKinds={
+              displayType === DisplayType.Heatmap
+                ? [SourceKind.Trace]
+                : undefined
+            }
             sourceSchemaPreview={
               <SourceSchemaPreview source={tableSource} variant="text" />
             }
diff --git a/packages/app/src/components/DBEditTimeChartForm/EditTimeChartForm.tsx b/packages/app/src/components/DBEditTimeChartForm/EditTimeChartForm.tsx
index 1bed1fef..6b35a513 100644
--- a/packages/app/src/components/DBEditTimeChartForm/EditTimeChartForm.tsx
+++ b/packages/app/src/components/DBEditTimeChartForm/EditTimeChartForm.tsx
@@ -417,14 +417,11 @@ export default function EditTimeChartForm({
       }
 
       if (displayType === DisplayType.Heatmap && Array.isArray(select)) {
-        const traceSource =
+        const defaultValue =
           tableSource?.kind === SourceKind.Trace &&
           tableSource.durationExpression
-            ? tableSource
-            : undefined;
-        const defaultValue = traceSource
-          ? getDurationMsExpression(traceSource)
-          : (select[0]?.valueExpression ?? '');
+            ? getDurationMsExpression(tableSource)
+            : (select[0]?.valueExpression ?? '');
         const heatmapSeries: SavedChartConfigWithSelectArray['select'] = [
           {
             aggFn: 'count',
@@ -436,12 +433,10 @@ export default function EditTimeChartForm({
         setValue('select', heatmapSeries);
         setValue('series', heatmapSeries);
         setValue('series.0.countExpression' as any, 'count()');
-        if (traceSource) {
-          setValue('numberFormat', {
-            output: 'duration' as any,
-            factor: 0.001,
-          });
-        }
+        setValue('numberFormat', {
+          output: 'duration' as any,
+          factor: 0.001,
+        });
       }
 
       // Don't auto-submit when config type changes, to avoid clearing form state (like source)