Prevent SQL snippet from running until message is finished

apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx

@@ -47,18 +47,20 @@ import AIOnboarding from './AIOnboarding'
 import CollapsibleCodeBlock from './CollapsibleCodeBlock'
 import { Message } from './Message'
 
-const MemoizedMessage = memo(({ message }: { message: MessageType }) => {
-  return (
-    <Message
-      key={message.id}
-      name={message.name}
-      role={message.role}
-      content={message.content}
-      createdAt={new Date(message.createdAt || new Date()).getTime()}
-      readOnly={message.role === 'user'}
-    />
-  )
-})
+const MemoizedMessage = memo(
+  ({ message, isLoading }: { message: MessageType; isLoading: boolean }) => {
+    return (
+      <Message
+        key={message.id}
+        id={message.id}
+        role={message.role}
+        content={message.content}
+        readOnly={message.role === 'user'}
+        isLoading={isLoading}
+      />
+    )
+  }
+)
 
 MemoizedMessage.displayName = 'MemoizedMessage'
 
@@ -133,8 +135,6 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP
   const canUpdateOrganization = useCheckPermissions(PermissionAction.UPDATE, 'organizations')
   const { mutate: updateOrganization, isLoading: isUpdating } = useOrganizationUpdateMutation()
 
-  const isLoading = isChatLoading
-
   const messages = useMemo(() => {
     const merged = [
       ...chatMessages,
@@ -151,9 +151,9 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP
   const renderedMessages = useMemo(
     () =>
       messages.map((message) => {
-        return <MemoizedMessage key={message.id} message={message} />
+        return <MemoizedMessage key={message.id} message={message} isLoading={isChatLoading} />
       }),
-    [messages]
+    [messages, isChatLoading]
   )
 
   const hasMessages = messages.length > 0
@@ -237,7 +237,7 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP
   }, [initialInput])
 
   useEffect(() => {
-    if (!isLoading) {
+    if (!isChatLoading) {
       if (inputRef.current) inputRef.current.focus()
     }
 
@@ -245,9 +245,9 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP
       () => {
         if (bottomRef.current) bottomRef.current.scrollIntoView({ behavior: 'smooth' })
       },
-      isLoading ? 100 : 500
+      isChatLoading ? 100 : 500
     )
-  }, [isLoading])
+  }, [isChatLoading])
 
   useEffect(() => {
     if (bottomRef.current) bottomRef.current.scrollIntoView({ behavior: 'smooth' })
@@ -278,7 +278,7 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP
               <div className="text-sm flex-1">{hasMessages ? 'New chat' : 'Assistant'}</div>
               <div className="flex gap-2">
                 {(hasMessages || suggestions || sqlSnippets) && (
-                  <Button type="default" disabled={isLoading} onClick={onResetConversation}>
+                  <Button type="default" disabled={isChatLoading} onClick={onResetConversation}>
                     Reset
                   </Button>
                 )}
@@ -516,8 +516,8 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP
             className={cn(
               'z-20 [&>textarea]:border-1 [&>textarea]:rounded-md [&>textarea]:!outline-none [&>textarea]:!ring-offset-0 [&>textarea]:!ring-0'
             )}
-            loading={isLoading}
-            disabled={!isApiKeySet || disablePrompts || isLoading}
+            loading={isChatLoading}
+            disabled={!isApiKeySet || disablePrompts || isChatLoading}
             placeholder={
               hasMessages
                 ? 'Reply to the assistant...'

apps/studio/components/ui/AIAssistantPanel/Message.tsx

@@ -8,25 +8,21 @@ import CollapsibleCodeBlock from './CollapsibleCodeBlock'
 import { SqlSnippet } from './SqlSnippet'
 
 interface MessageProps {
-  name?: string
+  id: string
   role: 'function' | 'system' | 'user' | 'assistant' | 'data' | 'tool'
   content?: string
-  createdAt?: number
-  isSelected?: boolean
-  isLoading?: boolean
+  isLoading: boolean
   readOnly?: boolean
   action?: React.ReactNode
   variant?: 'default' | 'warning'
 }
 
 export const Message = function Message({
-  name,
+  id,
   role,
   content,
-  createdAt,
   isLoading,
   readOnly,
-  isSelected = false,
   children,
   action = null,
   variant = 'default',
@@ -70,6 +66,7 @@ export const Message = function Message({
                 </div>
               ) : (
                 <SqlSnippet
+                  id={id}
                   readOnly={readOnly}
                   isLoading={isLoading}
                   sql={props.children[0].props.children}

apps/studio/components/ui/AIAssistantPanel/SqlSnippet.tsx

@@ -27,20 +27,13 @@ import { ButtonTooltip } from '../ButtonTooltip'
 import { isReadOnlySelect } from './AIAssistant.utils'
 
 interface SqlSnippetWrapperProps {
+  id: string
   sql: string
+  isLoading: boolean
   readOnly?: boolean
 }
 
-interface ParsedSqlProps {
-  sql: string
-  title: string
-  readOnly?: boolean
-  isChart: boolean
-  xAxis: string
-  yAxis: string
-}
-
-const SqlSnippetWrapper = ({ sql, readOnly = false }: SqlSnippetWrapperProps) => {
+const SqlSnippetWrapper = ({ id, sql, isLoading, readOnly = false }: SqlSnippetWrapperProps) => {
   const formatted = (sql || [''])[0]
   const propsMatch = formatted.match(/--\s*props:\s*(\{[^}]+\})/)
   const props = propsMatch ? JSON.parse(propsMatch[1]) : {}
@@ -50,24 +43,39 @@ const SqlSnippetWrapper = ({ sql, readOnly = false }: SqlSnippetWrapperProps) =>
   return (
     <div className="-mx-8 my-3 mt-2 border-b overflow-hidden">
       <SqlCard
+        id={id}
         sql={updatedFormatted}
         isChart={props.isChart}
         xAxis={props.xAxis}
         yAxis={props.yAxis}
         title={title}
         readOnly={readOnly}
+        isLoading={isLoading}
       />
     </div>
   )
 }
 
+interface ParsedSqlProps {
+  id: string
+  sql: string
+  title: string
+  isLoading: boolean
+  readOnly?: boolean
+  isChart: boolean
+  xAxis: string
+  yAxis: string
+}
+
 export const SqlCard = ({
+  id,
   sql,
   isChart,
   xAxis,
   yAxis,
   title,
   readOnly = false,
+  isLoading,
 }: ParsedSqlProps) => {
   const router = useRouter()
   const project = useSelectedProject()
@@ -76,6 +84,7 @@ export const SqlCard = ({
 
   const isInSQLEditor = router.pathname.includes('/sql')
   const isInNewSnippet = router.pathname.endsWith('/sql')
+
   const [showCode, setShowCode] = useState(readOnly || !isReadOnlySelect(sql))
   const [showResults, setShowResults] = useState(false)
   const [results, setResults] = useState<any[]>()
@@ -84,7 +93,6 @@ export const SqlCard = ({
 
   const { mutate: executeSql, isLoading: isExecuting } = useExecuteSqlMutation({
     onSuccess: (res) => {
-      console.log('SQL executed successfully:', res)
       setShowResults(true)
       setResults(res.result)
       setShowWarning(false)
@@ -118,13 +126,6 @@ export const SqlCard = ({
     })
   }, [project?.ref, project?.connectionString, sql, executeSql, readOnly])
 
-  useEffect(() => {
-    if (isReadOnlySelect(sql) && !results && !readOnly) {
-      handleExecute()
-    }
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [sql, handleExecute, readOnly])
-
   const handleEditInSQLEditor = () => {
     if (isInSQLEditor) {
       snapV2.setDiffContent(sql, DiffType.Addition)
@@ -136,6 +137,13 @@ export const SqlCard = ({
   const [errorHeader, ...errorContent] =
     (error?.formattedError?.split('\n') ?? [])?.filter((x: string) => x.length > 0) ?? []
 
+  useEffect(() => {
+    if (isReadOnlySelect(sql) && !results && !readOnly && !isLoading) {
+      handleExecute()
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [sql, readOnly, isLoading])
+
   return (
     <div className="overflow-hidden">
       <div className="flex items-center px-5 py-2 gap-2">