TableIndexAdvisorContext.tsx158 lines · main
1// @ts-nocheck
2import { useQueryClient } from '@tanstack/react-query'
3import { createContext, PropsWithChildren, useCallback, useContext, useState } from 'react'
4import { Sheet, SheetContent, SheetHeader, SheetTitle } from 'ui'
5
6import { useIndexAdvisorStatus } from '@/components/interfaces/QueryPerformance/hooks/useIsIndexAdvisorStatus'
7import { QueryIndexes } from '@/components/interfaces/QueryPerformance/QueryIndexes'
8import { databaseKeys } from '@/data/database/keys'
9import {
10 cleanIndexColumnName,
11 IndexAdvisorSuggestion,
12 TableIndexAdvisorData,
13 useTableIndexAdvisorQuery,
14} from '@/data/database/table-index-advisor-query'
15import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
16
17interface 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
28const 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
39interface TableIndexAdvisorProviderProps {
40 schema: string
41 table: string
42}
43
44export 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
151export function useTableIndexAdvisor() {
152 return useContext(TableIndexAdvisorContext)
153}
154
155export function useColumnHasIndexSuggestion(columnName: string): boolean {
156 const { columnsWithSuggestions } = useTableIndexAdvisor()
157 return columnsWithSuggestions.includes(columnName)
158}