TableIndexAdvisorContext.tsx158 lines · main
| 1 | // @ts-nocheck |
| 2 | import { useQueryClient } from '@tanstack/react-query' |
| 3 | import { createContext, PropsWithChildren, useCallback, useContext, useState } from 'react' |
| 4 | import { Sheet, SheetContent, SheetHeader, SheetTitle } from 'ui' |
| 5 | |
| 6 | import { useIndexAdvisorStatus } from '@/components/interfaces/QueryPerformance/hooks/useIsIndexAdvisorStatus' |
| 7 | import { QueryIndexes } from '@/components/interfaces/QueryPerformance/QueryIndexes' |
| 8 | import { databaseKeys } from '@/data/database/keys' |
| 9 | import { |
| 10 | cleanIndexColumnName, |
| 11 | IndexAdvisorSuggestion, |
| 12 | TableIndexAdvisorData, |
| 13 | useTableIndexAdvisorQuery, |
| 14 | } from '@/data/database/table-index-advisor-query' |
| 15 | import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject' |
| 16 | |
| 17 | interface TableIndexAdvisorContextValue { |
| 18 | isLoading: boolean |
| 19 | isAvailable: boolean |
| 20 | isEnabled: boolean |
| 21 | columnsWithSuggestions: string[] |
| 22 | suggestions: TableIndexAdvisorData['suggestions'] |
| 23 | openSheet: (columnName: string) => void |
| 24 | getSuggestionsForColumn: (columnName: string) => IndexAdvisorSuggestion[] |
| 25 | invalidate: () => Promise<void> |
| 26 | } |
| 27 | |
| 28 | const TableIndexAdvisorContext = createContext<TableIndexAdvisorContextValue>({ |
| 29 | isLoading: false, |
| 30 | isAvailable: false, |
| 31 | isEnabled: false, |
| 32 | columnsWithSuggestions: [], |
| 33 | suggestions: [], |
| 34 | openSheet: () => {}, |
| 35 | getSuggestionsForColumn: () => [], |
| 36 | invalidate: async () => {}, |
| 37 | }) |
| 38 | |
| 39 | interface TableIndexAdvisorProviderProps { |
| 40 | schema: string |
| 41 | table: string |
| 42 | } |
| 43 | |
| 44 | export function TableIndexAdvisorProvider({ |
| 45 | children, |
| 46 | schema, |
| 47 | table, |
| 48 | }: PropsWithChildren<TableIndexAdvisorProviderProps>) { |
| 49 | const { data: project } = useSelectedProjectQuery() |
| 50 | const { isIndexAdvisorAvailable, isIndexAdvisorEnabled } = useIndexAdvisorStatus() |
| 51 | const queryClient = useQueryClient() |
| 52 | const [isSheetOpen, setIsSheetOpen] = useState(false) |
| 53 | const [selectedColumn, setSelectedColumn] = useState<string | undefined>(undefined) |
| 54 | |
| 55 | const { data, isLoading } = useTableIndexAdvisorQuery( |
| 56 | { |
| 57 | projectRef: project?.ref, |
| 58 | connectionString: project?.connectionString, |
| 59 | schema, |
| 60 | table, |
| 61 | }, |
| 62 | { |
| 63 | enabled: isIndexAdvisorEnabled && !!schema && !!table, |
| 64 | } |
| 65 | ) |
| 66 | |
| 67 | const openSheet = useCallback((columnName: string) => { |
| 68 | setSelectedColumn(columnName) |
| 69 | setIsSheetOpen(true) |
| 70 | }, []) |
| 71 | |
| 72 | const closeSheet = useCallback(() => { |
| 73 | setIsSheetOpen(false) |
| 74 | setSelectedColumn(undefined) |
| 75 | }, []) |
| 76 | |
| 77 | const getSuggestionsForColumn = useCallback( |
| 78 | (columnName: string): IndexAdvisorSuggestion[] => { |
| 79 | if (!data?.suggestions) return [] |
| 80 | // Filter suggestions that include this column in their index statements |
| 81 | return data.suggestions.filter((suggestion) => |
| 82 | suggestion.index_statements.some((stmt) => { |
| 83 | const match = stmt.match(/USING\s+\w+\s*\(([^)]+)\)/i) |
| 84 | if (match) { |
| 85 | const columns = match[1].split(',').map((c) => cleanIndexColumnName(c)) |
| 86 | return columns.includes(columnName) |
| 87 | } |
| 88 | return false |
| 89 | }) |
| 90 | ) |
| 91 | }, |
| 92 | [data?.suggestions] |
| 93 | ) |
| 94 | |
| 95 | const invalidate = useCallback(async () => { |
| 96 | if (project?.ref && schema && table) { |
| 97 | await queryClient.invalidateQueries({ |
| 98 | queryKey: databaseKeys.tableIndexAdvisor(project.ref, schema, table), |
| 99 | }) |
| 100 | } |
| 101 | }, [queryClient, project?.ref, schema, table]) |
| 102 | |
| 103 | // Get the first suggestion for the selected column to pass to QueryIndexes |
| 104 | const selectedSuggestion = selectedColumn ? getSuggestionsForColumn(selectedColumn)[0] : null |
| 105 | |
| 106 | const prefetchedIndexAdvisorResult = selectedSuggestion |
| 107 | ? { |
| 108 | errors: [], |
| 109 | index_statements: selectedSuggestion.index_statements, |
| 110 | startup_cost_before: selectedSuggestion.startup_cost_before, |
| 111 | startup_cost_after: selectedSuggestion.startup_cost_after, |
| 112 | total_cost_before: selectedSuggestion.total_cost_before, |
| 113 | total_cost_after: selectedSuggestion.total_cost_after, |
| 114 | } |
| 115 | : null |
| 116 | |
| 117 | const value: TableIndexAdvisorContextValue = { |
| 118 | isLoading, |
| 119 | isAvailable: isIndexAdvisorAvailable, |
| 120 | isEnabled: isIndexAdvisorEnabled, |
| 121 | columnsWithSuggestions: data?.columnsWithSuggestions ?? [], |
| 122 | suggestions: data?.suggestions ?? [], |
| 123 | openSheet, |
| 124 | getSuggestionsForColumn, |
| 125 | invalidate, |
| 126 | } |
| 127 | |
| 128 | return ( |
| 129 | <TableIndexAdvisorContext.Provider value={value}> |
| 130 | {children} |
| 131 | <Sheet open={isSheetOpen} onOpenChange={(open) => !open && closeSheet()}> |
| 132 | <SheetContent className="flex flex-col gap-0 p-0 lg:w-[calc(100vw-802px)]! max-w-[700px]"> |
| 133 | <SheetHeader className="border-b px-5 py-3"> |
| 134 | <SheetTitle>Index Recommendation</SheetTitle> |
| 135 | </SheetHeader> |
| 136 | {selectedSuggestion && ( |
| 137 | <QueryIndexes |
| 138 | selectedRow={{ query: selectedSuggestion.query }} |
| 139 | columnName={selectedColumn} |
| 140 | suggestedSelectQuery={selectedSuggestion.query} |
| 141 | prefetchedIndexAdvisorResult={prefetchedIndexAdvisorResult} |
| 142 | onClose={closeSheet} |
| 143 | /> |
| 144 | )} |
| 145 | </SheetContent> |
| 146 | </Sheet> |
| 147 | </TableIndexAdvisorContext.Provider> |
| 148 | ) |
| 149 | } |
| 150 | |
| 151 | export function useTableIndexAdvisor() { |
| 152 | return useContext(TableIndexAdvisorContext) |
| 153 | } |
| 154 | |
| 155 | export function useColumnHasIndexSuggestion(columnName: string): boolean { |
| 156 | const { columnsWithSuggestions } = useTableIndexAdvisor() |
| 157 | return columnsWithSuggestions.includes(columnName) |
| 158 | } |