feat(mgmt-api): update reports API to use multiple attributes (#40473)

* feat(mgmt-api): update reports API

Related to https://github.com/supabase/infrastructure/pull/27519 and https://linear.app/supabase/issue/API-551/extend-infra-monitoring-with-multi-attribute-response

* tests: infra-monitoring-queries

---------

Co-authored-by: Ali Waseem <waseema393@gmail.com>

apps/studio/data/analytics/infra-monitoring-queries.test.ts

@@ -0,0 +1,72 @@
+import { describe, expect, it } from 'vitest'
+
+import type { InfraMonitoringMultiData } from './infra-monitoring-query'
+import { mapMultiResponseToAnalyticsData } from './infra-monitoring-queries'
+
+const mockResponse: InfraMonitoringMultiData = {
+  data: [
+    {
+      period_start: '2024-01-02 03:04:00',
+      values: {
+        max_cpu_usage: '1.5',
+        ram_usage: undefined,
+      },
+    },
+    {
+      period_start: '2024-01-02 04:04:00',
+      values: {
+        max_cpu_usage: undefined,
+        ram_usage: '3',
+      },
+    },
+  ],
+  series: {
+    max_cpu_usage: { format: 'percent', total: 2.5, totalAverage: 1.25, yAxisLimit: 100 },
+    ram_usage: { format: 'bytes', total: 3, totalAverage: 1.5, yAxisLimit: 200 },
+  },
+}
+
+describe('mapMultiResponseToAnalyticsData', () => {
+  it('maps attribute series and coerces values to numbers', () => {
+    const result = mapMultiResponseToAnalyticsData(mockResponse, ['max_cpu_usage', 'ram_usage'])
+
+    expect(result.max_cpu_usage?.data).toStrictEqual([
+      {
+        period_start: '2024-01-02 03:04:00',
+        periodStartFormatted: '03:04 02 Jan',
+        max_cpu_usage: 1.5,
+      },
+      {
+        period_start: '2024-01-02 04:04:00',
+        periodStartFormatted: '04:04 02 Jan',
+        max_cpu_usage: 0,
+      },
+    ])
+
+    expect(result.ram_usage?.data[0].ram_usage).toBe(0)
+    expect(result.ram_usage?.data[1].ram_usage).toBe(3)
+    expect(result.max_cpu_usage).toMatchObject({
+      format: 'percent',
+      total: 2.5,
+      yAxisLimit: 100,
+    })
+  })
+
+  it('omits attributes that are missing series metadata', () => {
+    const result = mapMultiResponseToAnalyticsData(mockResponse, [
+      'max_cpu_usage',
+      'ram_usage',
+      'pg_stat_database_num_backends',
+    ])
+
+    expect(result.max_cpu_usage).toBeDefined()
+    expect(result.ram_usage).toBeDefined()
+    expect(result.pg_stat_database_num_backends).toBeUndefined()
+  })
+
+  it('uses a custom date format when provided', () => {
+    const result = mapMultiResponseToAnalyticsData(mockResponse, ['max_cpu_usage'], 'YYYY/MM/DD')
+
+    expect(result.max_cpu_usage?.data[0].periodStartFormatted).toBe('2024/01/02')
+  })
+})

apps/studio/data/analytics/infra-monitoring-queries.ts

@@ -1,6 +1,24 @@
-import { AnalyticsInterval } from './constants'
+import { useMemo } from 'react'
+import type { UseQueryResult } from '@tanstack/react-query'
+
+import dayjs from 'dayjs'
+
+import type { AnalyticsData, AnalyticsInterval } from './constants'
 import type { InfraMonitoringAttribute } from './infra-monitoring-query'
-import { useInfraMonitoringQuery } from './infra-monitoring-query'
+import {
+  InfraMonitoringMultiData,
+  InfraMonitoringError,
+  useInfraMonitoringAttributesQuery,
+  useInfraMonitoringQuery,
+} from './infra-monitoring-query'
+
+const DEFAULT_DATE_FORMAT = 'HH:mm DD MMM'
+const DEFAULT_ATTRIBUTE: InfraMonitoringAttribute = 'max_cpu_usage'
+
+type InfraQueryResult = Pick<
+  UseQueryResult<AnalyticsData, InfraMonitoringError>,
+  'data' | 'error' | 'isError' | 'isFetching' | 'isLoading' | 'status'
+>
 
 export function useInfraMonitoringQueries(
   attributes: InfraMonitoringAttribute[],
@@ -12,18 +30,83 @@ export function useInfraMonitoringQueries(
   data: any,
   isVisible: boolean
 ) {
-  return attributes.map((attribute) =>
-    // eslint-disable-next-line react-hooks/rules-of-hooks
-    useInfraMonitoringQuery(
-      {
-        projectRef: ref as string,
-        attribute,
-        startDate,
-        endDate,
-        interval,
-        databaseIdentifier,
-      },
-      { enabled: data === undefined && isVisible }
-    )
+  const shouldFetch = data === undefined && isVisible
+  const hasAttributes = attributes.length > 0
+  const isSingleAttribute = attributes.length === 1
+  const primaryAttribute = attributes[0] ?? DEFAULT_ATTRIBUTE
+
+  const singleAttributeQuery = useInfraMonitoringQuery(
+    {
+      projectRef: ref as string,
+      attribute: primaryAttribute,
+      startDate,
+      endDate,
+      interval,
+      databaseIdentifier,
+    },
+    { enabled: shouldFetch && isSingleAttribute && hasAttributes }
   )
+
+  const multiQuery = useInfraMonitoringAttributesQuery(
+    {
+      projectRef: ref as string,
+      attributes,
+      startDate,
+      endDate,
+      interval,
+      databaseIdentifier,
+    },
+    { enabled: shouldFetch && attributes.length > 1 }
+  )
+
+  const seriesByAttribute = useMemo(() => {
+    if (!multiQuery.data) return undefined
+    return mapMultiResponseToAnalyticsData(multiQuery.data, attributes)
+  }, [multiQuery.data, attributes])
+
+  if (!hasAttributes) {
+    return []
+  }
+
+  if (isSingleAttribute) {
+    return [singleAttributeQuery]
+  }
+
+  return attributes.map<InfraQueryResult>((attribute) => ({
+    data: seriesByAttribute?.[attribute],
+    error: multiQuery.error,
+    isError: multiQuery.isError,
+    isFetching: multiQuery.isFetching,
+    isLoading: multiQuery.isLoading,
+    status: multiQuery.status,
+  }))
+}
+
+export function mapMultiResponseToAnalyticsData(
+  response: InfraMonitoringMultiData,
+  attributes: InfraMonitoringAttribute[],
+  dateFormat: string = DEFAULT_DATE_FORMAT
+) {
+  return attributes.reduce<Record<string, AnalyticsData>>((acc, attribute) => {
+    const metadata = response.series?.[attribute]
+    if (!metadata) return acc
+
+    const dataPoints = response.data.map((point) => {
+      const value = point.values?.[attribute]
+      return {
+        period_start: point.period_start,
+        periodStartFormatted: dayjs(point.period_start).format(dateFormat),
+        [attribute]: value === undefined ? 0 : Number(value),
+      }
+    })
+
+    acc[attribute] = {
+      data: dataPoints,
+      format: metadata.format,
+      total: metadata.total,
+      yAxisLimit: metadata.yAxisLimit,
+    }
+
+    return acc
+  }, {})
 }

apps/studio/data/analytics/infra-monitoring-query.ts

@@ -35,6 +35,28 @@ export type InfraMonitoringVariables = {
   modifier?: (x: number) => number
 }
 
+export type InfraMonitoringSeriesMetadata = {
+  yAxisLimit: number
+  format: string
+  total: number
+  totalAverage: number
+}
+
+export type InfraMonitoringMultiResponse = {
+  data: {
+    period_start: string
+    values: Record<string, string | undefined>
+  }[]
+  series: Record<string, InfraMonitoringSeriesMetadata>
+}
+
+export type InfraMonitoringMultiVariables = Omit<
+  InfraMonitoringVariables,
+  'attribute' | 'modifier'
+> & {
+  attributes: InfraMonitoringAttribute[]
+}
+
 export async function getInfraMonitoring(
   {
     projectRef,
@@ -69,8 +91,44 @@ export async function getInfraMonitoring(
   return data as unknown as AnalyticsData
 }
 
+export async function getInfraMonitoringAttributes(
+  {
+    projectRef,
+    attributes,
+    startDate,
+    endDate,
+    interval = '1h',
+    databaseIdentifier,
+  }: InfraMonitoringMultiVariables,
+  signal?: AbortSignal
+) {
+  if (!projectRef) throw new Error('Project ref is required')
+  if (!attributes?.length) throw new Error('At least one attribute is required')
+  if (!startDate) throw new Error('Start date is required')
+  if (!endDate) throw new Error('End date is required')
+
+  const { data, error } = await get('/platform/projects/{ref}/infra-monitoring', {
+    params: {
+      path: { ref: projectRef },
+      // Attributes support is not yet reflected in the generated client types.
+      query: {
+        attributes,
+        startDate,
+        endDate,
+        interval,
+        databaseIdentifier,
+      } as any,
+    },
+    signal,
+  })
+
+  if (error) handleError(error)
+  return data as unknown as InfraMonitoringMultiResponse
+}
+
 export type InfraMonitoringData = Awaited<ReturnType<typeof getInfraMonitoring>>
 export type InfraMonitoringError = unknown
+export type InfraMonitoringMultiData = Awaited<ReturnType<typeof getInfraMonitoringAttributes>>
 
 export const useInfraMonitoringQuery = <TData = InfraMonitoringData>(
   {
@@ -123,3 +181,40 @@ export const useInfraMonitoringQuery = <TData = InfraMonitoringData>(
     staleTime: 1000 * 60,
     ...options,
   })
+
+export const useInfraMonitoringAttributesQuery = <TData = InfraMonitoringMultiData>(
+  {
+    projectRef,
+    attributes,
+    startDate,
+    endDate,
+    interval = '1h',
+    databaseIdentifier,
+  }: InfraMonitoringMultiVariables,
+  {
+    enabled = true,
+    ...options
+  }: UseCustomQueryOptions<InfraMonitoringMultiData, InfraMonitoringError, TData> = {}
+) =>
+  useQuery<InfraMonitoringMultiData, InfraMonitoringError, TData>({
+    queryKey: analyticsKeys.infraMonitoringGroup(projectRef, {
+      attributes,
+      startDate,
+      endDate,
+      interval,
+      databaseIdentifier,
+    }),
+    queryFn: ({ signal }) =>
+      getInfraMonitoringAttributes(
+        { projectRef, attributes, startDate, endDate, interval, databaseIdentifier },
+        signal
+      ),
+    enabled:
+      enabled &&
+      typeof projectRef !== 'undefined' &&
+      !!attributes?.length &&
+      typeof startDate !== 'undefined' &&
+      typeof endDate !== 'undefined',
+    staleTime: 1000 * 60,
+    ...options,
+  })

apps/studio/data/analytics/keys.ts

@@ -121,6 +121,35 @@ export const analyticsKeys = {
       'infra-monitoring',
       { attribute, startDate, endDate, interval, databaseIdentifier },
     ] as const,
+  infraMonitoringGroup: (
+    projectRef: string | undefined,
+    {
+      attributes,
+      startDate,
+      endDate,
+      interval,
+      databaseIdentifier,
+    }: {
+      attributes?: string[]
+      startDate?: string
+      endDate?: string
+      interval?: string
+      databaseIdentifier?: string
+    }
+  ) =>
+    [
+      'projects',
+      projectRef,
+      'infra-monitoring',
+      'group',
+      {
+        attributes: attributes ? [...attributes].sort() : undefined,
+        startDate,
+        endDate,
+        interval,
+        databaseIdentifier,
+      },
+    ] as const,
   projectMetrics: (projectRef: string | undefined, { interval }: { interval?: string }) =>
     ['projects', projectRef, 'project.metrics', { interval }] as const,
   usageApiCounts: (projectRef: string | undefined, interval: string | undefined) =>