From 23cfebbb9f407dde25d38b2c9941fc86fbefbaa8 Mon Sep 17 00:00:00 2001
From: Jonathan Summers-Muir <MildTomato@users.noreply.github.com>
Date: Fri, 11 Apr 2025 13:33:16 +0800
Subject: [PATCH] selected query UI

---
 .../QueryInsights/QueryInsights.tsx           |  24 +-
 .../QueryInsights/components/MetricPill.tsx   |  59 ++++
 .../components/MetricToggleButtons.tsx        | 106 +++++++
 .../QueryInsights/components/MetricsChart.tsx | 263 ++++++++----------
 .../components/SelectedQueryBadge.tsx         |  64 +++++
 5 files changed, 364 insertions(+), 152 deletions(-)
 create mode 100644 apps/studio/components/interfaces/QueryInsights/components/MetricPill.tsx
 create mode 100644 apps/studio/components/interfaces/QueryInsights/components/MetricToggleButtons.tsx
 create mode 100644 apps/studio/components/interfaces/QueryInsights/components/SelectedQueryBadge.tsx

diff --git a/apps/studio/components/interfaces/QueryInsights/QueryInsights.tsx b/apps/studio/components/interfaces/QueryInsights/QueryInsights.tsx
index c5aba4d36bc..cf1c2746ea9 100644
--- a/apps/studio/components/interfaces/QueryInsights/QueryInsights.tsx
+++ b/apps/studio/components/interfaces/QueryInsights/QueryInsights.tsx
@@ -16,6 +16,7 @@ import {
 } from 'data/query-insights/query-insights-query'
 import { useDatabaseSelectorStateSnapshot } from 'state/database-selector'
 import { useReadReplicasQuery } from 'data/read-replicas/replicas-query'
+import { SelectedQueryBadge } from './components/SelectedQueryBadge'
 
 export type MetricType = 'rows_read' | 'query_latency' | 'calls' | 'cache_hits'
 
@@ -78,6 +79,27 @@ export const QueryInsights = () => {
 
   const selectedDatabase = databases?.find((db) => db.identifier === state.selectedDatabaseId)
 
+  // Event listener for clearing the selected query
+  useEffect(() => {
+    const handleClearSelectedQuery = (event: CustomEvent) => {
+      if (event.detail?.clearQuery) {
+        setSelectedQuery(null)
+      }
+    }
+
+    window.addEventListener(
+      'clearSelectedQueryInsightsQuery',
+      handleClearSelectedQuery as EventListener
+    )
+
+    return () => {
+      window.removeEventListener(
+        'clearSelectedQueryInsightsQuery',
+        handleClearSelectedQuery as EventListener
+      )
+    }
+  }, [])
+
   useEffect(() => {
     // Update the time range when selectedTimeRange changes
     switch (selectedTimeRange) {
@@ -209,7 +231,7 @@ export const QueryInsights = () => {
         onValueChange={(value) => setSelectedMetric(value as MetricType)}
         className="space-y-4"
       >
-        <TabsList_Shadcn_ className="px-5 flex gap-5">
+        <TabsList_Shadcn_ className="w-full flex justify-start gap-5 px-5">
           {METRICS.map((metric) => (
             <TabsTrigger_Shadcn_ key={metric.id} value={metric.id}>
               {metric.label}
diff --git a/apps/studio/components/interfaces/QueryInsights/components/MetricPill.tsx b/apps/studio/components/interfaces/QueryInsights/components/MetricPill.tsx
new file mode 100644
index 00000000000..e96b380fafb
--- /dev/null
+++ b/apps/studio/components/interfaces/QueryInsights/components/MetricPill.tsx
@@ -0,0 +1,59 @@
+import { cn } from 'ui'
+import { MetricType } from '../QueryInsights'
+
+interface MetricPillProps {
+  label: string
+  value: string | number
+  color?: string
+  metricType?: MetricType
+  isActive?: boolean
+  onClick?: () => void
+}
+
+export const MetricPill = ({
+  label,
+  value,
+  color,
+  metricType,
+  isActive = true,
+  onClick,
+}: MetricPillProps) => {
+  // Get the dot color based on metric type if no color is provided
+  const getDotColor = () => {
+    if (color) return color
+
+    switch (metricType) {
+      case 'query_latency':
+        return 'hsl(var(--chart-1))'
+      case 'rows_read':
+        return 'hsl(var(--chart-2))'
+      case 'calls':
+        return 'hsl(var(--chart-3))'
+      case 'cache_hits':
+        return 'hsl(var(--chart-4))'
+      default:
+        return 'currentColor'
+    }
+  }
+
+  return (
+    <button
+      type="button"
+      className={cn(
+        'text-xs px-2 py-0.5 rounded-full transition-colors inline-flex items-center gap-1.5 border',
+        isActive ? 'bg-surface-300' : 'hover:bg-surface-100 border-transparent'
+      )}
+      onClick={onClick}
+    >
+      <div
+        className={cn('w-2 h-2 rounded-full', !isActive && 'opacity-50')}
+        style={{
+          backgroundColor: getDotColor(),
+        }}
+      ></div>
+      <span className="text-foreground">
+        {label}: {value}
+      </span>
+    </button>
+  )
+}
diff --git a/apps/studio/components/interfaces/QueryInsights/components/MetricToggleButtons.tsx b/apps/studio/components/interfaces/QueryInsights/components/MetricToggleButtons.tsx
new file mode 100644
index 00000000000..c518b55bc79
--- /dev/null
+++ b/apps/studio/components/interfaces/QueryInsights/components/MetricToggleButtons.tsx
@@ -0,0 +1,106 @@
+import { motion, AnimatePresence } from 'framer-motion'
+import { cn } from 'ui'
+import { ChartConfigType, ChartConfigResult } from '../chartConfigs'
+import { QueryInsightsQuery } from 'data/query-insights/query-insights-query'
+import { MetricType } from '../QueryInsights'
+import { MetricPill } from './MetricPill'
+
+// Define the exact shape of the visibleMetrics state
+export interface VisibleMetricsState {
+  p50: boolean
+  p95: boolean
+  p99: boolean
+  p99_9: boolean
+  query_latency: boolean
+  rows: boolean
+  query_rows: boolean
+  shared_blks_hit: boolean
+  shared_blks_read: boolean
+  shared_blks_dirtied: boolean
+  shared_blks_written: boolean
+  cache_hit_ratio: boolean
+  cache_miss_ratio: boolean
+  calls: boolean
+  query_calls: boolean
+  [key: string]: boolean // Allow string indexing
+}
+
+interface MetricToggleButtonsProps {
+  chartConfig: ChartConfigResult<any>
+  visibleMetrics: VisibleMetricsState
+  setVisibleMetrics: (value: React.SetStateAction<VisibleMetricsState>) => void
+  metric: MetricType
+  selectedQuery: QueryInsightsQuery | null
+}
+
+export const MetricToggleButtons = ({
+  chartConfig,
+  visibleMetrics,
+  setVisibleMetrics,
+  metric,
+  selectedQuery,
+}: MetricToggleButtonsProps) => {
+  if (!chartConfig) return null
+
+  const { config, chartData } = chartConfig
+
+  return (
+    <AnimatePresence>
+      {Object.entries(config).map(([key, value]) => {
+        // If we have a selected query, hide the query-specific metrics since they'll
+        // be shown inside the Selected Query pill
+        if (
+          (key === 'query_rows' || key === 'query_latency' || key === 'query_calls') &&
+          selectedQuery
+        ) {
+          return null
+        }
+
+        // Only show query-specific metrics when a query is selected
+        if (
+          (key === 'query_rows' || key === 'query_latency' || key === 'query_calls') &&
+          !selectedQuery
+        ) {
+          return null
+        }
+
+        // Filter out the blocks metrics for cache hits chart
+        if (
+          metric === 'cache_hits' &&
+          (key === 'shared_blks_dirtied' || key === 'shared_blks_written')
+        ) {
+          return null
+        }
+
+        const metricValue = Number(
+          chartData[chartData.length - 1]?.[key as keyof (typeof chartData)[0]]
+        )
+        const formattedValue = !isNaN(metricValue) ? (value as any).formatter(metricValue) : '0'
+
+        return (
+          <motion.div
+            key={key}
+            initial={{ opacity: 0, scale: 0.95, x: -8 }}
+            animate={{ opacity: 1, scale: 1, x: 0 }}
+            exit={{ opacity: 0, scale: 0.95, x: -8 }}
+            transition={{ duration: 0.15 }}
+          >
+            <MetricPill
+              label={(value as any).label}
+              value={formattedValue}
+              color={(value as any).color}
+              metricType={metric}
+              isActive={visibleMetrics[key]}
+              onClick={() =>
+                setVisibleMetrics((prev) => ({
+                  ...prev,
+                  [key]: !prev[key],
+                }))
+              }
+            />
+          </motion.div>
+        )
+      })}
+    </AnimatePresence>
+  )
+}
diff --git a/apps/studio/components/interfaces/QueryInsights/components/MetricsChart.tsx b/apps/studio/components/interfaces/QueryInsights/components/MetricsChart.tsx
index f6cc5925d54..59966312d28 100644
--- a/apps/studio/components/interfaces/QueryInsights/components/MetricsChart.tsx
+++ b/apps/studio/components/interfaces/QueryInsights/components/MetricsChart.tsx
@@ -34,6 +34,8 @@ import {
   ChartConfigResult,
   getCacheHitsConfig,
 } from '../chartConfigs'
+import { MetricToggleButtons, VisibleMetricsState } from './MetricToggleButtons'
+import { SelectedQueryBadge } from './SelectedQueryBadge'
 
 // Chart opacity constants
 const CHART_OPACITY = {
@@ -65,23 +67,7 @@ export function MetricsChart({
   selectedQuery,
 }: MetricsChartProps) {
   const { ref } = useParams()
-  const [visibleMetrics, setVisibleMetrics] = useState<{
-    p50: boolean
-    p95: boolean
-    p99: boolean
-    p99_9: boolean
-    query_latency: boolean
-    rows: boolean
-    query_rows: boolean
-    shared_blks_hit: boolean
-    shared_blks_read: boolean
-    shared_blks_dirtied: boolean
-    shared_blks_written: boolean
-    cache_hit_ratio: boolean
-    cache_miss_ratio: boolean
-    calls: boolean
-    query_calls: boolean
-  }>({
+  const [visibleMetrics, setVisibleMetrics] = useState<VisibleMetricsState>({
     p50: true,
     p95: true,
     p99: true,
@@ -146,6 +132,16 @@ export function MetricsChart({
     }
   }, [data, metric, queryLatencyData, queryRowsData, queryCallsData])
 
+  // Handler for clearing the selected query
+  const handleClearSelectedQuery = () => {
+    // Find the appropriate onQuerySelect handler (it's actually in the parent component)
+    // This will be called via the parent's prop
+    const paramsToSend = { clearQuery: true }
+    window.dispatchEvent(
+      new CustomEvent('clearSelectedQueryInsightsQuery', { detail: paramsToSend })
+    )
+  }
+
   if (isLoading) {
     return (
       <div className="flex-1 flex items-center justify-center">
@@ -183,67 +179,62 @@ export function MetricsChart({
     const config = (chartConfig as ChartConfigResult<any>).config
     const chartData = (chartConfig as ChartConfigResult<any>).chartData
 
+    // Get metrics for the selected query badge
+    const getQueryMetrics = () => {
+      if (!selectedQuery) return []
+
+      const latestData = chartData[chartData.length - 1]
+      if (!latestData) return []
+
+      const metrics = []
+
+      if (metric === 'query_latency') {
+        // Add just the single query latency value
+        if (latestData.query_latency !== undefined) {
+          metrics.push({
+            label: 'Query latency',
+            value: formatLatency(latestData.query_latency),
+            color: config.query_latency?.color || 'hsl(var(--chart-1))',
+          })
+        }
+      } else if (metric === 'rows_read') {
+        // Add rows read for the query
+        if (latestData.query_rows !== undefined) {
+          metrics.push({
+            label: 'Rows',
+            value: formatMetricValue(latestData.query_rows),
+            color: config.query_rows?.color || 'hsl(var(--chart-2))',
+          })
+        }
+      } else if (metric === 'cache_hits') {
+        // Cache hits don't have query-specific metrics
+      }
+
+      return metrics
+    }
+
     return (
       <div className="h-[320px] flex flex-col">
         <div className="flex flex-col flex-1 min-h-0">
-          <div className="flex gap-2 px-5 mb-4">
-            <AnimatePresence>
-              {Object.entries(config).map(([key, value]) => {
-                // Only show specific buttons based on conditions
-                if ((key === 'query_rows' || key === 'query_latency') && !selectedQuery) {
-                  return null
-                }
+          <div className="flex flex-wrap items-center gap-2 px-5 mb-4">
+            <MetricToggleButtons
+              chartConfig={chartConfig as ChartConfigResult<any>}
+              visibleMetrics={visibleMetrics}
+              setVisibleMetrics={setVisibleMetrics}
+              metric={metric}
+              selectedQuery={selectedQuery}
+            />
 
-                // Filter out the blocks metrics for cache hits chart
-                if (
-                  metric === 'cache_hits' &&
-                  (key === 'shared_blks_dirtied' || key === 'shared_blks_written')
-                ) {
-                  return null
-                }
-
-                const metricValue = Number(
-                  chartData[chartData.length - 1]?.[key as keyof (typeof chartData)[0]]
-                )
-                const formattedValue = !isNaN(metricValue)
-                  ? (value as any).formatter(metricValue)
-                  : '0'
-
-                return (
-                  <motion.button
-                    key={key}
-                    initial={{ opacity: 0, scale: 0.95, x: -8 }}
-                    animate={{ opacity: 1, scale: 1, x: 0 }}
-                    exit={{ opacity: 0, scale: 0.95, x: -8 }}
-                    transition={{ duration: 0.15 }}
-                    className={cn(
-                      'text-xs px-2 py-1 rounded-full transition-colors inline-flex items-center gap-1.5 border',
-                      visibleMetrics[key as keyof typeof visibleMetrics]
-                        ? 'bg-surface-300'
-                        : 'hover:bg-surface-100 border-transparent'
-                    )}
-                    onClick={() =>
-                      setVisibleMetrics((prev) => ({
-                        ...prev,
-                        [key]: !prev[key as keyof typeof visibleMetrics],
-                      }))
-                    }
-                  >
-                    <div
-                      className={cn(
-                        'w-2 h-2 rounded-full',
-                        !visibleMetrics[key as keyof typeof visibleMetrics] && 'opacity-50'
-                      )}
-                      style={{ backgroundColor: (value as any).color }}
-                    />
-                    <span>
-                      {(value as any).label}: {formattedValue}
-                    </span>
-                  </motion.button>
-                )
-              })}
-            </AnimatePresence>
+            {selectedQuery && (
+              <SelectedQueryBadge
+                selectedQuery={selectedQuery}
+                onClear={handleClearSelectedQuery}
+                metric={metric}
+                metrics={getQueryMetrics()}
+              />
+            )}
           </div>
+
           <div className="flex-1 min-h-0">
             <ChartContainer className="h-full w-full" config={config}>
               <AreaChart
@@ -477,69 +468,41 @@ export function MetricsChart({
       ? chartData.reduce((sum, point) => sum + (point.query_calls ?? 0), 0)
       : 0
 
+    // Get metrics for the selected query badge
+    const getQueryMetrics = () => {
+      if (!selectedQuery) return []
+
+      return [
+        {
+          label: 'Calls',
+          value: formatMetricValue(totalQueryCalls),
+          color: config.query_calls?.color || 'hsl(var(--chart-3))',
+        },
+      ]
+    }
+
     return (
       <div className="h-[320px] flex flex-col">
         <div className="flex flex-col flex-1 min-h-0">
-          <div className="flex gap-2 px-5 mb-4">
-            <AnimatePresence>
-              <motion.button
-                key="calls"
-                initial={{ opacity: 0, scale: 0.95, x: -8 }}
-                animate={{ opacity: 1, scale: 1, x: 0 }}
-                exit={{ opacity: 0, scale: 0.95, x: -8 }}
-                transition={{ duration: 0.15 }}
-                className={cn(
-                  'text-xs px-2 py-1 rounded-full transition-colors inline-flex items-center gap-1.5 border',
-                  visibleMetrics.calls
-                    ? 'bg-surface-300'
-                    : 'hover:bg-surface-100 border-transparent'
-                )}
-                onClick={() =>
-                  setVisibleMetrics((prev) => ({
-                    ...prev,
-                    calls: !prev.calls,
-                  }))
-                }
-              >
-                <div
-                  className={cn('w-2 h-2 rounded-full', !visibleMetrics.calls && 'opacity-50')}
-                  style={{ backgroundColor: config.calls.color }}
-                />
-                <span>Total Calls: {formatMetricValue(totalCalls)}</span>
-              </motion.button>
+          <div className="flex flex-wrap items-center gap-2 px-5 mb-4">
+            <MetricToggleButtons
+              chartConfig={chartConfig as ChartConfigResult<any>}
+              visibleMetrics={visibleMetrics}
+              setVisibleMetrics={setVisibleMetrics}
+              metric={metric}
+              selectedQuery={selectedQuery}
+            />
 
-              {selectedQuery && (
-                <motion.button
-                  key="query_calls"
-                  initial={{ opacity: 0, scale: 0.95, x: -8 }}
-                  animate={{ opacity: 1, scale: 1, x: 0 }}
-                  exit={{ opacity: 0, scale: 0.95, x: -8 }}
-                  transition={{ duration: 0.15 }}
-                  className={cn(
-                    'text-xs px-2 py-1 rounded-full transition-colors inline-flex items-center gap-1.5 border',
-                    visibleMetrics.query_calls
-                      ? 'bg-surface-300'
-                      : 'hover:bg-surface-100 border-transparent'
-                  )}
-                  onClick={() =>
-                    setVisibleMetrics((prev) => ({
-                      ...prev,
-                      query_calls: !prev.query_calls,
-                    }))
-                  }
-                >
-                  <div
-                    className={cn(
-                      'w-2 h-2 rounded-full',
-                      !visibleMetrics.query_calls && 'opacity-50'
-                    )}
-                    style={{ backgroundColor: config.query_calls.color }}
-                  />
-                  <span>Query Calls: {formatMetricValue(totalQueryCalls)}</span>
-                </motion.button>
-              )}
-            </AnimatePresence>
+            {selectedQuery && (
+              <SelectedQueryBadge
+                selectedQuery={selectedQuery}
+                onClear={handleClearSelectedQuery}
+                metric={metric}
+                metrics={getQueryMetrics()}
+              />
+            )}
           </div>
+
           <div className="flex-1 min-h-0">
             <ChartContainer className="h-full w-full" config={config}>
               <AreaChart
@@ -641,27 +604,25 @@ export function MetricsChart({
   return (
     <div className="h-[320px] flex flex-col">
       <div className="flex flex-col flex-1 min-h-0">
-        <div className="flex gap-2 px-5 mb-4">
-          <AnimatePresence>
-            <motion.button
-              key="generic-metric"
-              initial={{ opacity: 0, scale: 0.95, x: -8 }}
-              animate={{ opacity: 1, scale: 1, x: 0 }}
-              exit={{ opacity: 0, scale: 0.95, x: -8 }}
-              transition={{ duration: 0.15 }}
-              className={cn(
-                'text-xs px-2 py-1 rounded-full transition-colors inline-flex items-center gap-1.5 border',
-                'bg-surface-300'
-              )}
-            >
-              <div
-                className="w-2 h-2 rounded-full"
-                style={{ backgroundColor: 'hsl(var(--chart-1))' }}
-              />
-              <span>{metric}</span>
-            </motion.button>
-          </AnimatePresence>
+        <div className="flex flex-wrap items-center gap-2 px-5 mb-4">
+          <MetricToggleButtons
+            chartConfig={chartConfig as ChartConfigResult<any>}
+            visibleMetrics={visibleMetrics}
+            setVisibleMetrics={setVisibleMetrics}
+            metric={metric}
+            selectedQuery={selectedQuery}
+          />
+
+          {selectedQuery && (
+            <SelectedQueryBadge
+              selectedQuery={selectedQuery}
+              onClear={handleClearSelectedQuery}
+              metric={metric}
+              metrics={[]}
+            />
+          )}
         </div>
+
         <div className="flex-1 min-h-0">
           <ChartContainer
             className="h-full w-full"
diff --git a/apps/studio/components/interfaces/QueryInsights/components/SelectedQueryBadge.tsx b/apps/studio/components/interfaces/QueryInsights/components/SelectedQueryBadge.tsx
new file mode 100644
index 00000000000..552a32f1449
--- /dev/null
+++ b/apps/studio/components/interfaces/QueryInsights/components/SelectedQueryBadge.tsx
@@ -0,0 +1,64 @@
+import { X } from 'lucide-react'
+import { QueryInsightsQuery } from 'data/query-insights/query-insights-query'
+import { Button, cn } from 'ui'
+import { MetricType } from '../QueryInsights'
+import { MetricPill } from './MetricPill'
+
+interface MetricInfo {
+  label: string
+  value: string | number
+  color?: string
+}
+
+interface SelectedQueryBadgeProps {
+  selectedQuery: QueryInsightsQuery | null
+  onClear: () => void
+  className?: string
+  metric?: MetricType
+  metrics?: MetricInfo[]
+}
+
+export const SelectedQueryBadge = ({
+  selectedQuery,
+  onClear,
+  className,
+  metric,
+  metrics = [],
+}: SelectedQueryBadgeProps) => {
+  if (!selectedQuery) return null
+
+  return (
+    <div
+      className={cn(
+        'inline-flex items-center px-3 pr-2 py-0.5 rounded-full text-xs border',
+        'border-surface-300 dark:border-surface-400',
+        'shadow-sm',
+        className
+      )}
+    >
+      <div className="flex items-center gap-2">
+        <span className="text-foreground-light">Selected Query</span>
+        <div className="flex items-center gap-2">
+          {metrics.map((metricInfo, index) => (
+            <MetricPill
+              key={index}
+              label={metricInfo.label}
+              value={metricInfo.value}
+              color={metricInfo.color}
+              metricType={metric}
+              isActive={true}
+            />
+          ))}
+        </div>
+      </div>
+
+      <button
+        onClick={onClear}
+        className="text-foreground-lighter hover:text-foreground-light ml-1"
+        title="Clear selected query filter"
+      >
+        <X size={12} strokeWidth={2.5} />
+      </button>
+    </div>
+  )
+}