supabase/apps/studio/data/database/table-index-advisor-query.ts

import { getTableIndexAdvisorSql, type SafeSqlFragment } from '@supabase/pg-meta'
import { useQuery } from '@tanstack/react-query'

import { databaseKeys } from './keys'
import { filterProtectedSchemaIndexStatements } from '@/components/interfaces/QueryPerformance/IndexAdvisor/index-advisor.utils'
import { executeSql } from '@/data/sql/execute-sql-query'
import type { ResponseError, UseCustomQueryOptions } from '@/types'

export type TableIndexAdvisorVariables = {
  projectRef?: string
  connectionString?: string | null
  schema: string
  table: string
}

export type IndexAdvisorSuggestion = {
  query: SafeSqlFragment
  calls: number
  total_time: number
  mean_time: number
  index_statements: string[]
  startup_cost_before: number
  startup_cost_after: number
  total_cost_before: number
  total_cost_after: number
  improvement_percentage: number
}

export type TableIndexAdvisorResponse = {
  suggestions: IndexAdvisorSuggestion[]
  columnsWithSuggestions: string[]
}

// Strips ordering modifiers and outer quotes from a raw column token from an index statement.
// e.g. '"created_at" DESC NULLS LAST' -> 'created_at'
export function cleanIndexColumnName(raw: string): string {
  return raw
    .trim()
    .replace(/\s+(asc|desc)(\s+nulls\s+(first|last))?$/i, '')
    .replace(/\s+nulls\s+(first|last)$/i, '')
    .trim()
    .replace(/^"(.+)"$/, '$1')
}

//Extracts column names from index statements
//e.g. "CREATE INDEX ON public.users USING btree (email)" -> "email"
function extractColumnsFromIndexStatements(indexStatements: string[]): string[] {
  const columns = new Set<string>()

  for (const statement of indexStatements) {
    // Match patterns like "USING btree (column_name)" or "USING btree (column1, column2)"
    const match = statement.match(/USING\s+\w+\s*\(([^)]+)\)/i)
    if (match) {
      const columnPart = match[1]
      // Split by comma and clean up each column name
      columnPart.split(',').forEach((col) => {
        const cleanedCol = cleanIndexColumnName(col)
        // Only add simple identifiers — skip expressions like lower(col)
        if (cleanedCol && /^[a-z_][a-z0-9_$]*$/i.test(cleanedCol)) {
          columns.add(cleanedCol)
        }
      })
    }
  }

  return Array.from(columns)
}

export async function getTableIndexAdvisorSuggestions({
  projectRef,
  connectionString,
  schema,
  table,
}: TableIndexAdvisorVariables): Promise<TableIndexAdvisorResponse> {
  if (!projectRef) throw new Error('Project ref is required')
  if (!schema) throw new Error('Schema is required')
  if (!table) throw new Error('Table is required')

  const sql = getTableIndexAdvisorSql(schema, table)

  const { result } = await executeSql<
    Array<{
      query: string
      calls: number
      total_time: number
      mean_time: number
      index_statements: string[]
      startup_cost_before: number
      startup_cost_after: number
      total_cost_before: number
      total_cost_after: number
    }>
  >({
    projectRef,
    connectionString,
    sql,
  })

  const suggestions: IndexAdvisorSuggestion[] = (result || [])
    .filter((row) => row.index_statements && row.index_statements.length > 0)
    .map((row) => {
      // Filter out protected schema index statements
      const filteredStatements = filterProtectedSchemaIndexStatements(row.index_statements)

      // Skip this suggestion if all statements were filtered out
      if (filteredStatements.length === 0) {
        return null
      }

      const improvement =
        row.total_cost_before > 0
          ? ((row.total_cost_before - row.total_cost_after) / row.total_cost_before) * 100
          : 0

      return {
        query: row.query as SafeSqlFragment,
        calls: row.calls,
        total_time: row.total_time,
        mean_time: row.mean_time,
        index_statements: filteredStatements,
        startup_cost_before: row.startup_cost_before,
        startup_cost_after: row.startup_cost_after,
        total_cost_before: row.total_cost_before,
        total_cost_after: row.total_cost_after,
        improvement_percentage: Math.round(improvement * 100) / 100,
      }
    })
    .filter((suggestion): suggestion is IndexAdvisorSuggestion => suggestion !== null)

  // Extract all unique columns from suggestions
  const allIndexStatements = suggestions.flatMap((s) => s.index_statements)
  const columnsWithSuggestions = extractColumnsFromIndexStatements(allIndexStatements)

  return {
    suggestions,
    columnsWithSuggestions,
  }
}

export type TableIndexAdvisorData = Awaited<ReturnType<typeof getTableIndexAdvisorSuggestions>>
export type TableIndexAdvisorError = ResponseError

export function useTableIndexAdvisorQuery<TData = TableIndexAdvisorData>(
  { projectRef, connectionString, schema, table }: TableIndexAdvisorVariables,
  {
    enabled = true,
    ...options
  }: UseCustomQueryOptions<TableIndexAdvisorData, TableIndexAdvisorError, TData> = {}
) {
  return useQuery<TableIndexAdvisorData, TableIndexAdvisorError, TData>({
    queryKey: databaseKeys.tableIndexAdvisor(projectRef, schema, table),
    queryFn: () =>
      getTableIndexAdvisorSuggestions({
        projectRef,
        connectionString,
        schema,
        table,
      }),
    enabled: enabled && typeof projectRef !== 'undefined' && !!schema && !!table,
    retry: false,
    staleTime: 5 * 60 * 1000,
    ...options,
  })
}