konnosDashboard
DisplayBlockRenderer.tsx265 lines · main
1import { acceptUntrustedSql, type UntrustedSqlFragment } from '@supabase/pg-meta'
2import { PermissionAction } from '@supabase/shared-types/out/constants'
3import { useQueryClient } from '@tanstack/react-query'
4import type { ToolUIPart } from 'ai'
5import { useParams } from 'common'
6import { useRouter } from 'next/router'
7import { useRef, useState, type DragEvent, type PropsWithChildren } from 'react'
8
9import { DEFAULT_CHART_CONFIG, QueryBlock } from '../QueryBlock/QueryBlock'
10import { identifyQueryType } from './AIAssistant.utils'
11import { ConfirmFooter } from './ConfirmFooter'
12import { ChartConfig } from '@/components/interfaces/SQLEditor/UtilityPanel/ChartConfig'
13import { entityTypeKeys } from '@/data/entity-types/keys'
14import { lintKeys } from '@/data/lint/keys'
15import { usePrimaryDatabase } from '@/data/read-replicas/replicas-query'
16import { useExecuteSqlMutation } from '@/data/sql/execute-sql-mutation'
17import { useChangedSync } from '@/hooks/misc/useChanged'
18import { useAsyncCheckPermissions } from '@/hooks/misc/useCheckPermissions'
19import { useProfile } from '@/lib/profile'
20import { useTrack } from '@/lib/telemetry/track'
21
22interface DisplayBlockRendererProps {
23 messageId: string
24 toolCallId: string
25 initialArgs: {
26 sql: UntrustedSqlFragment
27 label?: string
28 isWriteQuery?: boolean
29 view?: 'table' | 'chart'
30 xAxis?: string
31 yAxis?: string
32 }
33 initialResults?: unknown
34 /** Called when locally running SQL fails before or during client-side execution. */
35 onError?: (args: { messageId: string; errorText: string }) => void
36 /** Responds affirmatively to an AI SDK tool approval request; does not run SQL directly. */
37 onApprove?: () => void
38 /** Responds negatively to an AI SDK tool approval request; does not run SQL directly. */
39 onDeny?: () => void
40 /** AI SDK tool state used to show approval UI for pending tool calls. */
41 toolState?: ToolUIPart['state']
42 toolApprovalRespondedApproved?: boolean
43 isLastPart?: boolean
44 isLastMessage?: boolean
45 showConfirmFooter?: boolean
46 onChartConfigChange?: (chartConfig: ChartConfig) => void
47 /** Called when the user clicks the query block play button to run SQL locally. */
48 onQueryRun?: (queryType: 'select' | 'mutation') => void
49}
50
51export const DisplayBlockRenderer = ({
52 messageId,
53 toolCallId,
54 initialArgs,
55 initialResults,
56 onError,
57 onApprove,
58 onDeny,
59 toolState,
60 toolApprovalRespondedApproved,
61 isLastPart = false,
62 isLastMessage = false,
63 showConfirmFooter = true,
64 onChartConfigChange,
65 onQueryRun,
66}: PropsWithChildren<DisplayBlockRendererProps>) => {
67 const queryClient = useQueryClient()
68
69 const savedInitialArgs = useRef(initialArgs)
70 const savedInitialResults = useRef(initialResults)
71 const savedInitialConfig = useRef<ChartConfig>({
72 ...DEFAULT_CHART_CONFIG,
73 view: initialArgs.view === 'chart' ? 'chart' : 'table',
74 xKey: initialArgs.xAxis ?? '',
75 yKey: initialArgs.yAxis ?? '',
76 })
77
78 const router = useRouter()
79 const { ref } = useParams()
80 const { profile } = useProfile()
81
82 const track = useTrack()
83 const { can: canCreateSQLSnippet } = useAsyncCheckPermissions(
84 PermissionAction.CREATE,
85 'user_content',
86 {
87 resource: { type: 'sql', owner_id: profile?.id },
88 subject: { id: profile?.id },
89 }
90 )
91
92 const [chartConfig, setChartConfig] = useState<ChartConfig>(() => ({
93 ...DEFAULT_CHART_CONFIG,
94 view: initialArgs.view === 'chart' ? 'chart' : 'table',
95 xKey: initialArgs.xAxis ?? '',
96 yKey: initialArgs.yAxis ?? '',
97 }))
98
99 const [rows, setRows] = useState<any[] | undefined>(
100 Array.isArray(initialResults) ? initialResults : undefined
101 )
102 const isReportsPage = router.pathname.endsWith('/reports/[id]')
103 const isHomePage = router.pathname === '/project/[ref]'
104 const isDraggableToReports = canCreateSQLSnippet && (isReportsPage || isHomePage)
105 const label = initialArgs.label || 'SQL Results'
106 const [isWriteQuery, setIsWriteQuery] = useState<boolean>(initialArgs.isWriteQuery || false)
107 const sqlQuery = initialArgs.sql
108
109 const { database: primaryDatabase } = usePrimaryDatabase({ projectRef: ref })
110
111 const readOnlyConnectionString = primaryDatabase?.connection_string_read_only
112 const postgresConnectionString = primaryDatabase?.connectionString
113
114 const {
115 mutate: executeSql,
116 error: executeSqlError,
117 isPending: executeSqlLoading,
118 } = useExecuteSqlMutation({
119 onError: () => {
120 // Suppress toast because error message is displayed inline
121 },
122 })
123
124 const toolCallIdChanged = useChangedSync(toolCallId)
125 if (toolCallIdChanged) {
126 setChartConfig(savedInitialConfig.current)
127 onChartConfigChange?.(savedInitialConfig.current)
128 setIsWriteQuery(savedInitialArgs.current.isWriteQuery || false)
129 setRows(Array.isArray(savedInitialResults.current) ? savedInitialResults.current : undefined)
130 }
131
132 const initialResultsChanged = useChangedSync(initialResults)
133 if (initialResultsChanged) {
134 const normalized = Array.isArray(initialResults) ? initialResults : undefined
135 if (!normalized || normalized === rows) return
136 setRows(normalized)
137 }
138
139 const handleRunQuery = (queryType: 'select' | 'mutation') => {
140 if (!sqlQuery) return
141
142 onQueryRun?.(queryType)
143
144 track('assistant_suggestion_run_query_clicked', {
145 queryType,
146 ...(queryType === 'mutation'
147 ? { mutationType: identifyQueryType(sqlQuery) ?? 'unknown' }
148 : {}),
149 })
150 }
151
152 const runQuery = (queryType: 'select' | 'mutation') => {
153 if (!ref || !sqlQuery) return
154
155 const connectionString =
156 queryType === 'mutation'
157 ? postgresConnectionString
158 : (readOnlyConnectionString ?? postgresConnectionString)
159
160 if (!connectionString) {
161 const fallbackMessage = 'Unable to find a database connection to execute this query.'
162 onError?.({ messageId, errorText: fallbackMessage })
163 return
164 }
165
166 if (queryType === 'mutation') {
167 setIsWriteQuery(true)
168 }
169 executeSql(
170 { projectRef: ref, connectionString, sql: acceptUntrustedSql(sqlQuery) },
171 {
172 onSuccess: (data) => {
173 setRows(Array.isArray(data.result) ? data.result : undefined)
174 setIsWriteQuery(queryType === 'mutation' || initialArgs.isWriteQuery || false)
175 if (queryType === 'mutation') {
176 queryClient.invalidateQueries({ queryKey: lintKeys.lint(ref) })
177 queryClient.invalidateQueries({ queryKey: entityTypeKeys.list(ref) })
178 }
179 },
180 onError: (error) => {
181 const lowerMessage = error.message.toLowerCase()
182 const isReadOnlyError =
183 lowerMessage.includes('read-only transaction') ||
184 lowerMessage.includes('permission denied') ||
185 lowerMessage.includes('must be owner')
186
187 if (queryType === 'select' && isReadOnlyError) {
188 setIsWriteQuery(true)
189 }
190
191 onError?.({ messageId, errorText: error.message })
192 },
193 }
194 )
195 }
196
197 const handleExecute = (queryType: 'select' | 'mutation') => {
198 handleRunQuery(queryType)
199 runQuery(queryType)
200 }
201
202 const handleUpdateChartConfig = ({
203 chartConfig: updatedValues,
204 }: {
205 chartConfig: Partial<ChartConfig>
206 }) => {
207 setChartConfig((prev) => {
208 const next = { ...prev, ...updatedValues }
209 onChartConfigChange?.(next)
210 return next
211 })
212 }
213
214 const handleDragStart = (e: DragEvent<Element>) => {
215 e.dataTransfer.setData(
216 'application/json',
217 JSON.stringify({ label, sql: sqlQuery, config: chartConfig })
218 )
219 }
220
221 const isApprovalRequested = toolState === 'approval-requested'
222 const isApprovalResponded = toolState === 'approval-responded'
223 const isApprovalDenied = isApprovalResponded && toolApprovalRespondedApproved === false
224 const shouldShowConfirmFooter =
225 showConfirmFooter &&
226 (isApprovalRequested || (isApprovalResponded && !isApprovalDenied)) &&
227 isLastPart &&
228 isLastMessage &&
229 (isApprovalResponded || (!!onApprove && !!onDeny))
230 const isRunningApprovedTool = (isApprovalResponded && !isApprovalDenied) || executeSqlLoading
231
232 return (
233 <div className="display-block w-auto overflow-x-hidden">
234 <div className="relative z-10">
235 <QueryBlock
236 label={label}
237 isWriteQuery={isWriteQuery}
238 sql={sqlQuery}
239 results={rows}
240 errorText={executeSqlError?.message}
241 chartConfig={chartConfig}
242 onExecute={handleExecute}
243 onUpdateChartConfig={handleUpdateChartConfig}
244 draggable={isDraggableToReports}
245 onDragStart={handleDragStart}
246 disabled={shouldShowConfirmFooter}
247 isExecuting={isRunningApprovedTool}
248 />
249 </div>
250 {shouldShowConfirmFooter && (
251 <div className="mx-4">
252 <ConfirmFooter
253 message="Assistant wants to run this query"
254 cancelLabel="Skip"
255 confirmLabel="Run Query"
256 confirmLabelLoading="Running..."
257 isLoading={isApprovalResponded || executeSqlLoading}
258 onCancel={isApprovalRequested ? onDeny : undefined}
259 onConfirm={isApprovalRequested ? onApprove : undefined}
260 />
261 </div>
262 )}
263 </div>
264 )
265}