insights/products/metrics/frontend/components/metricsSamplesLogic.tsx

import { MakeLogicType, actions, connect, kea, listeners, path, reducers, selectors } from 'kea'
import { loaders } from 'kea-loaders'

import { teamLogic } from 'scenes/teamLogic'

import { metricsSamplesCreate } from 'products/metrics/frontend/generated/api'
import type { _MetricEventSampleApi } from 'products/metrics/frontend/generated/api.schemas'
import { canViewMetrics } from 'products/metrics/frontend/metricsAccess'

import { formatSeriesName, seriesColor } from './metricsSeries'
import { type MetricsViewMode, type MetricsViewerSeries, metricsViewerLogic, resolveDate } from './metricsViewerLogic'

// The side panel next to the chart: per-series aggregates, or the raw emissions
// behind the chart with their trace linkage (the metric->trace pivot).
export type MetricsPanelTab = 'aggregates' | 'samples'

export const SAMPLES_LIMIT = 50

// One row of the Aggregates tab: a series with its headline numbers.
export interface MetricsAggregateRow {
    name: string
    color: string
    latest: number
    total: number
}

// A traced emission rendered as a clickable exemplar dot on the viewer chart —
// the metric->trace pivot straight from the time series.
export interface MetricsTraceExemplar {
    timestamp: string
    value: number
    traceId: string
    spanId: string
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface metricsSamplesLogicValues {
    dateFrom: string | null // metricsViewerLogic
    dateTo: string | null // metricsViewerLogic
    metricName: string // metricsViewerLogic
    queryResults: MetricsViewerSeries[] // metricsViewerLogic
    viewMode: MetricsViewMode // metricsViewerLogic
    currentTeamId: number | null // teamLogic
    activeTab: MetricsPanelTab
    aggregateRows: MetricsAggregateRow[]
    samples: _MetricEventSampleApi[]
    samplesLoading: boolean
    traceExemplars: MetricsTraceExemplar[]
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface metricsSamplesLogicActions {
    fetchQueryResultsSuccess: (
        queryResults: import('products/metrics/frontend/generated/api.schemas')._MetricSeriesApi[],
        payload?: any
    ) => {
        payload?: any
        queryResults: import('products/metrics/frontend/generated/api.schemas')._MetricSeriesApi[]
    } // metricsViewerLogic
    setDateFrom: (dateFrom: string | null) => {
        dateFrom: string | null
    } // metricsViewerLogic
    setDateTo: (dateTo: string | null) => {
        dateTo: string | null
    } // metricsViewerLogic
    setMetricName: (metricName: string) => {
        metricName: string
    } // metricsViewerLogic
    loadSamples: (_: any) => any
    loadSamplesFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    loadSamplesSuccess: (
        samples: _MetricEventSampleApi[],
        payload?: any
    ) => {
        samples: _MetricEventSampleApi[]
        payload?: any
    }
    setActiveTab: (activeTab: MetricsPanelTab) => {
        activeTab: MetricsPanelTab
    }
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface metricsSamplesLogicMeta {
    __keaTypeGenInternalSelectorTypes: {
        aggregateRows: (
            queryResults: import('products/metrics/frontend/generated/api.schemas')._MetricSeriesApi[],
            metricName: string
        ) => MetricsAggregateRow[]
        traceExemplars: (samples: _MetricEventSampleApi[]) => MetricsTraceExemplar[]
    }
}

export type metricsSamplesLogicType = MakeLogicType<
    metricsSamplesLogicValues,
    metricsSamplesLogicActions,
    Record<string, any>,
    metricsSamplesLogicMeta
>

export const metricsSamplesLogic = kea<metricsSamplesLogicType>([
    path(['products', 'metrics', 'frontend', 'components', 'metricsSamplesLogic']),
    connect(() => ({
        values: [
            teamLogic,
            ['currentTeamId'],
            metricsViewerLogic,
            ['metricName', 'dateFrom', 'dateTo', 'queryResults', 'viewMode'],
        ],
        actions: [metricsViewerLogic, ['setMetricName', 'setDateFrom', 'setDateTo', 'fetchQueryResultsSuccess']],
    })),
    actions({
        setActiveTab: (activeTab: MetricsPanelTab) => ({ activeTab }),
    }),
    reducers({
        activeTab: ['aggregates' as MetricsPanelTab, { setActiveTab: (_, { activeTab }) => activeTab }],
    }),
    loaders(({ values }) => ({
        samples: [
            [] as _MetricEventSampleApi[],
            {
                loadSamples: async (_, breakpoint) => {
                    if (!canViewMetrics()) {
                        return []
                    }
                    const metricName = values.metricName.trim()
                    const dateFrom = resolveDate(values.dateFrom)
                    if (!metricName || !dateFrom) {
                        return []
                    }
                    await breakpoint(300)
                    const dateTo = resolveDate(values.dateTo) ?? undefined
                    const response = await metricsSamplesCreate(String(values.currentTeamId), {
                        query: {
                            metricName,
                            dateFrom,
                            ...(dateTo ? { dateTo } : {}),
                            limit: SAMPLES_LIMIT,
                        },
                    })
                    breakpoint()
                    return response.results
                },
            },
        ],
    })),
    selectors({
        // Rows for the Aggregates tab, derived from the chart's series so both
        // views always describe the same query (colors match the chart legend).
        aggregateRows: [
            (s) => [s.queryResults, s.metricName],
            (queryResults: MetricsViewerSeries[], metricName: string): MetricsAggregateRow[] =>
                queryResults.map((series, index) => {
                    // Null points are gaps (unrepresentable buckets) — skip them
                    // rather than counting them as zero.
                    const values = series.points.map((point) => point.value).filter((value) => value !== null)
                    return {
                        name: formatSeriesName(series, metricName),
                        color: seriesColor(index),
                        latest: values.length ? values[values.length - 1] : 0,
                        total: values.reduce((sum, value) => sum + value, 0),
                    }
                }),
        ],
        // Exemplar dots for the chart: only emissions that carry trace context.
        traceExemplars: [
            (s) => [s.samples],
            (samples: _MetricEventSampleApi[]): MetricsTraceExemplar[] =>
                samples
                    .filter((sample) => sample.trace_id)
                    .map((sample) => ({
                        timestamp: sample.timestamp,
                        value: sample.value,
                        traceId: sample.trace_id,
                        spanId: sample.span_id,
                    })),
        ],
    }),
    listeners(({ actions, values }) => ({
        setActiveTab: ({ activeTab }) => {
            if (activeTab === 'samples') {
                actions.loadSamples({})
            }
        },
        // The chart's exemplar overlay renders from these samples, so every chart
        // redraw refreshes them; stat mode has no chart to dot, so skip it there.
        fetchQueryResultsSuccess: () => {
            if (values.viewMode === 'chart') {
                actions.loadSamples({})
            }
        },
        // The viewer's filters are the samples' filters: any change that redraws
        // the chart refreshes the visible samples too, but only when they're shown.
        setMetricName: () => {
            if (values.activeTab === 'samples') {
                actions.loadSamples({})
            }
        },
        setDateFrom: () => {
            if (values.activeTab === 'samples') {
                actions.loadSamples({})
            }
        },
        setDateTo: () => {
            if (values.activeTab === 'samples') {
                actions.loadSamples({})
            }
        },
    })),
])