From 752e77c74041614d652cf404bc6022ff74bc3c58 Mon Sep 17 00:00:00 2001
From: Saxon Fletcher <saxonafletcher@gmail.com>
Date: Mon, 23 Dec 2024 15:26:18 +1000
Subject: [PATCH] unify queryblock and sqlsnippet

---
 .../interfaces/Reports/GridResize.tsx         |  81 ++-
 .../interfaces/Reports/ParametersPopover.tsx  |   2 +-
 .../interfaces/Reports/QueryBlock.tsx         | 556 ++++++++++--------
 .../interfaces/Reports/QueryBlockWrapper.tsx  |  57 ++
 .../ui/AIAssistantPanel/Message.tsx           |   6 +-
 .../AIAssistantPanel/QueryBlockWithProps.tsx  | 130 ++++
 6 files changed, 553 insertions(+), 279 deletions(-)
 create mode 100644 apps/studio/components/interfaces/Reports/QueryBlockWrapper.tsx
 create mode 100644 apps/studio/components/ui/AIAssistantPanel/QueryBlockWithProps.tsx

diff --git a/apps/studio/components/interfaces/Reports/GridResize.tsx b/apps/studio/components/interfaces/Reports/GridResize.tsx
index 9b2142cc38a..19ed671210e 100644
--- a/apps/studio/components/interfaces/Reports/GridResize.tsx
+++ b/apps/studio/components/interfaces/Reports/GridResize.tsx
@@ -29,6 +29,7 @@ import BarChart from 'components/ui/Charts/BarChart'
 import dayjs from 'dayjs'
 import { ChartConfig } from '../SQLEditor/UtilityPanel/ChartConfig'
 import QueryBlock from 'components/interfaces/Reports/QueryBlock'
+import QueryBlockWrapper from './QueryBlockWrapper'
 
 const ReactGridLayout = WidthProvider(Responsive)
 
@@ -60,12 +61,11 @@ const DEFAULT_CHART_CONFIG: ChartConfig = {
 interface LayoutItem {
   id: string
   label: string
-  sql?: string
-  isSnippet?: boolean
-  results?: any
-  isChart?: boolean
+  sql: string
+  isChart: boolean
   chartConfig?: ChartConfig
-  // ... other existing fields ...
+  results?: any
+  // ... any other existing fields ...
 }
 
 const GridResize = ({
@@ -81,6 +81,44 @@ const GridResize = ({
 }: GridResizeProps) => {
   const [collectedParams, setCollectedParams] = useState<Parameter[]>([])
 
+  // Merge parameters from all blocks
+  const handleBlockParameters = (params: Parameter[]) => {
+    setCollectedParams((prev) => {
+      const merged = [...prev]
+      params.forEach((param) => {
+        const existingIndex = merged.findIndex((p) => p.name === param.name)
+        if (existingIndex === -1) {
+          merged.push(param)
+        }
+      })
+      return merged
+    })
+  }
+
+  const handleToggleChart = (itemId: string) => {
+    const updatedLayout = editableReport.layout.map((x: LayoutItem) => {
+      if (x.id === itemId) {
+        return {
+          ...x,
+          isChart: !x.isChart,
+          chartConfig: x.chartConfig || DEFAULT_CHART_CONFIG,
+        }
+      }
+      return x
+    })
+    setEditableReport({ ...editableReport, layout: updatedLayout })
+  }
+
+  const handleUpdateChartConfig = (itemId: string, config: ChartConfig) => {
+    const updatedLayout = editableReport.layout.map((x: LayoutItem) => {
+      if (x.id === itemId) {
+        return { ...x, chartConfig: config }
+      }
+      return x
+    })
+    setEditableReport({ ...editableReport, layout: updatedLayout })
+  }
+
   function onLayoutChange(layout: any) {
     let updatedLayout = [...editableReport.layout]
     layout.map((item: any) => {
@@ -96,21 +134,6 @@ const GridResize = ({
     })
   }
 
-  // Merge parameters from all blocks and notify parent once all are collected
-  const handleBlockParameters = (params: Parameter[]) => {
-    console.log('collected:', params)
-    setCollectedParams((prev) => {
-      const merged = [...prev]
-      params.forEach((param) => {
-        const existingIndex = merged.findIndex((p) => p.name === param.name)
-        if (existingIndex === -1) {
-          merged.push(param)
-        }
-      })
-      return merged
-    })
-  }
-
   // Notify parent when all parameters are collected
   useEffect(() => {
     onSetParameter(collectedParams)
@@ -132,32 +155,32 @@ const GridResize = ({
       isResizable={true}
       draggableHandle=".grid-item-drag-handle"
     >
-      {editableReport.layout.map((x: LayoutItem) => (
+      {editableReport.layout.map((item: LayoutItem) => (
         <div
-          key={x.id}
+          key={item.id}
           data-grid={{
-            ...x,
+            ...item,
             minH: 3,
             minW: 1,
             maxH: 6,
             maxW: LAYOUT_COLUMN_COUNT,
           }}
-          className="react-grid-layout__report-item bg-surface-100 border-overlay group relative rounded border shadow-sm hover:border-green-900"
+          className="react-grid-layout__report-item hover:border-green-900 group"
         >
-          <div className="grid-item-drag-handle absolute left-1 top-1 z-20 cursor-move opacity-0 transition-opacity group-hover:opacity-100">
+          <div className="grid-item-drag-handle absolute left-1 top-1 z-40 cursor-move opacity-0 transition-opacity group-hover:opacity-100">
             <GripHorizontal size={14} strokeWidth={1.5} />
           </div>
-          <QueryBlock
-            item={x}
+          <QueryBlockWrapper
+            item={item}
             disableUpdate={disableUpdate}
             onRemoveChart={onRemoveChart}
             startDate={startDate}
             endDate={endDate}
             interval={interval}
-            editableReport={editableReport}
-            setEditableReport={setEditableReport}
             parameterValues={parameterValues}
             onSetParameter={handleBlockParameters}
+            onToggleChart={() => handleToggleChart(item.id)}
+            onUpdateChartConfig={(config) => handleUpdateChartConfig(item.id, config)}
           />
         </div>
       ))}
diff --git a/apps/studio/components/interfaces/Reports/ParametersPopover.tsx b/apps/studio/components/interfaces/Reports/ParametersPopover.tsx
index 6ee157395e1..25aa3def334 100644
--- a/apps/studio/components/interfaces/Reports/ParametersPopover.tsx
+++ b/apps/studio/components/interfaces/Reports/ParametersPopover.tsx
@@ -132,7 +132,7 @@ const ParametersPopover = ({
   return (
     <Popover_Shadcn_ modal={false}>
       <PopoverTrigger_Shadcn_ asChild>
-        <Button icon={<ArrowUpDown size={14} />} type="outline" size="tiny" className="w-7 h-7" />
+        <Button icon={<ArrowUpDown size={14} />} type="text" size="tiny" className="w-7 h-7" />
       </PopoverTrigger_Shadcn_>
       <PopoverContent_Shadcn_ side="bottom" align="end" className="w-[300px] p-4">
         <div className="space-y-4">
diff --git a/apps/studio/components/interfaces/Reports/QueryBlock.tsx b/apps/studio/components/interfaces/Reports/QueryBlock.tsx
index c3b187e2bf2..6cc4665203e 100644
--- a/apps/studio/components/interfaces/Reports/QueryBlock.tsx
+++ b/apps/studio/components/interfaces/Reports/QueryBlock.tsx
@@ -1,15 +1,10 @@
-import { useContentIdQuery } from 'data/content/content-id-query'
 import { useParams } from 'common'
 import { useEffect, useState, useCallback, useMemo } from 'react'
 import 'react-grid-layout/css/styles.css'
 import 'react-resizable/css/styles.css'
-import { X, Settings2, ArrowUpDown, Play } from 'lucide-react'
-import { Responsive, WidthProvider } from 'react-grid-layout'
+import { X, Settings2, ArrowUpDown, Play, Code } from 'lucide-react'
 
-import ChartHandler from 'components/to-be-cleaned/Charts/ChartHandler'
 import { ButtonTooltip } from 'components/ui/ButtonTooltip'
-import { useDatabaseSelectorStateSnapshot } from 'state/database-selector'
-import { LAYOUT_COLUMN_COUNT } from './Reports.constants'
 import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
 import Results from '../SQLEditor/UtilityPanel/Results'
 import {
@@ -27,6 +22,9 @@ import {
   ToggleGroup,
   ToggleGroupItem,
   Input_Shadcn_,
+  cn,
+  CodeBlock,
+  SQL_ICON,
 } from 'ui'
 import { useProjectContext } from 'components/layouts/ProjectLayout/ProjectContext'
 import {
@@ -43,6 +41,11 @@ import { BarChart2, Table } from 'lucide-react'
 import { toast } from 'sonner'
 import { parseParameters, processParameterizedSql, Parameter } from 'lib/sql-parameters'
 import ParametersPopover from './ParametersPopover'
+import {
+  isReadOnlySelect,
+  containsUnknownFunction,
+} from 'components/ui/AIAssistantPanel/AIAssistant.utils'
+import { Admonition } from 'ui-patterns'
 
 // Add helper function for cumulative results
 const getCumulativeResults = (results: { rows: any[] }, config: ChartConfig) => {
@@ -71,66 +74,59 @@ const DEFAULT_CHART_CONFIG: ChartConfig = {
   showGrid: false,
 }
 
-// Update layout item type
-interface LayoutItem {
-  id: string
-  label: string
-  sql?: string
-  isSnippet?: boolean
-  results?: any
-  isChart?: boolean
-  chartConfig?: ChartConfig
-  // ... other existing fields ...
-}
-
 interface QueryBlockProps {
-  item: LayoutItem
+  sql: string
+  isChart: boolean
+  chartConfig?: ChartConfig
+  results?: any
+  label: string
+  id: string
   disableUpdate: boolean
-  onRemoveChart: ({ metric }: { metric: { key: string } }) => void
+  onRemoveChart?: ({ metric }: { metric: { key: string } }) => void
   startDate?: string
   endDate?: string
   interval?: string
-  editableReport: any
-  setEditableReport: (payload: any) => void
+  onToggleChart: () => void
+  onUpdateChartConfig: (config: ChartConfig) => void
   parameterValues?: Record<string, string>
   onSetParameter?: (params: Parameter[]) => void
+  isLoading?: boolean
+  maxHeight?: number
+  actions?: React.ReactNode
+  runQuery?: boolean
 }
 
 const QueryBlock = ({
-  item,
+  sql,
+  isChart = false,
+  chartConfig = DEFAULT_CHART_CONFIG,
+  results,
+  label,
+  id,
   disableUpdate,
   onRemoveChart,
   startDate,
   endDate,
   interval,
-  editableReport,
-  setEditableReport,
+  onToggleChart,
+  onUpdateChartConfig,
   parameterValues: externalParameterValues,
   onSetParameter,
+  isLoading = false,
+  maxHeight,
+  actions,
+  runQuery = false,
 }: QueryBlockProps) => {
   const { ref } = useParams()
   const { project } = useProjectContext()
-  const state = useDatabaseSelectorStateSnapshot()
 
-  const [sql, setSql] = useState<string>()
   const [queryResult, setQueryResult] = useState<any>()
   const [parameterValues, setParameterValues] = useState<Record<string, string>>({})
+  const [showSql, setShowSql] = useState(false)
+  const [showWarning, setShowWarning] = useState<'hasWriteOperation' | 'hasUnknownFunctions'>()
 
   const combinedParameterValues = { ...externalParameterValues, ...parameterValues }
 
-  const { data: snippetData } = useContentIdQuery(
-    { projectRef: ref, id: item.id },
-    { enabled: item.isSnippet }
-  )
-
-  console.log('item', item, snippetData)
-
-  useEffect(() => {
-    if (snippetData?.content?.sql) {
-      setSql(snippetData.content.sql)
-    }
-  }, [snippetData])
-
   const handleSuccess = useCallback((data: any) => {
     setQueryResult(data.result)
   }, [])
@@ -139,28 +135,12 @@ const QueryBlock = ({
     onSuccess: handleSuccess,
   })
 
-  const onToggleChart = (item: LayoutItem) => {
-    const updatedLayout = editableReport.layout.map((x: LayoutItem) => {
-      if (x.id === item.id) {
-        return {
-          ...x,
-          isChart: !x.isChart,
-          chartConfig: x.chartConfig || DEFAULT_CHART_CONFIG,
-        }
-      }
-      return x
-    })
-    setEditableReport({ ...editableReport, layout: updatedLayout })
+  const handleToggleChart = () => {
+    onToggleChart()
   }
 
-  const onUpdateChartConfig = (item: LayoutItem, config: ChartConfig) => {
-    const updatedLayout = editableReport.layout.map((x: LayoutItem) => {
-      if (x.id === item.id) {
-        return { ...x, chartConfig: config }
-      }
-      return x
-    })
-    setEditableReport({ ...editableReport, layout: updatedLayout })
+  const handleUpdateChartConfig = (config: ChartConfig) => {
+    onUpdateChartConfig(config)
   }
 
   // Run once on mount to parse parameters and notify parent
@@ -184,12 +164,16 @@ const QueryBlock = ({
 
   // Update execute call to include parameters
   const handleExecute = () => {
-    if (!sql) return
+    if (!sql || isLoading) return
+
+    if (!isReadOnlySelect(sql)) {
+      const hasUnknownFunctions = containsUnknownFunction(sql)
+      setShowWarning(hasUnknownFunctions ? 'hasUnknownFunctions' : 'hasWriteOperation')
+      return
+    }
 
     try {
       const processedSql = processParameterizedSql(sql, combinedParameterValues)
-      console.log('processedSql', processedSql)
-
       execute({
         projectRef: ref,
         connectionString: project?.connectionString,
@@ -201,196 +185,274 @@ const QueryBlock = ({
   }
 
   useEffect(() => {
-    handleExecute()
-  }, [sql, parameterValues, externalParameterValues])
+    if (!sql || isLoading) return
 
-  console.log('localParameters', item.label, parameterValues)
+    if (runQuery && isReadOnlySelect(sql)) {
+      handleExecute()
+    }
+  }, [sql, isLoading, runQuery])
 
-  if (item.isSnippet) {
-    return (
-      <div className="h-full flex flex-col">
-        <div className="flex justify-between px-5 py-2 items-center z-10 border-b">
-          <h3 className="text-sm text-foreground-light">{item.label}</h3>
-          <div className="flex gap-2">
-            {queryResult && (
-              <>
-                {parameters.length > 0 && (
-                  <ParametersPopover
-                    parameters={parameters}
-                    parameterValues={combinedParameterValues}
-                    onSubmit={handleParametersSubmit}
-                  />
-                )}
-                <Popover_Shadcn_ modal={false}>
-                  <PopoverTrigger_Shadcn_ asChild>
-                    <Button
-                      icon={<Settings2 size={14} />}
-                      type="outline"
-                      size="tiny"
-                      className="w-7 h-7"
-                    />
-                  </PopoverTrigger_Shadcn_>
-                  <PopoverContent_Shadcn_ side="bottom" align="end" className="w-[300px] p-3">
-                    <form className="grid gap-2">
-                      <ToggleGroup
-                        type="single"
-                        value={item.isChart ? 'chart' : 'table'}
-                        className="w-full"
-                        onValueChange={(value) => {
-                          if (value) onToggleChart(item)
-                        }}
-                      >
-                        <ToggleGroupItem
-                          className="w-full"
-                          value="table"
-                          aria-label="Show as table"
-                        >
-                          <Table className="h-4 w-4" />
-                        </ToggleGroupItem>
-                        <ToggleGroupItem
-                          className="w-full"
-                          value="chart"
-                          aria-label="Show as chart"
-                        >
-                          <BarChart2 className="h-4 w-4" />
-                        </ToggleGroupItem>
-                      </ToggleGroup>
-
-                      {item.isChart && (
-                        <>
-                          <Select_Shadcn_
-                            value={item.chartConfig?.xKey}
-                            onValueChange={(value) =>
-                              onUpdateChartConfig(item, { ...item.chartConfig, xKey: value })
-                            }
-                          >
-                            <SelectTrigger_Shadcn_>
-                              X Axis {item.chartConfig?.xKey && `- ${item.chartConfig.xKey}`}
-                            </SelectTrigger_Shadcn_>
-                            <SelectContent_Shadcn_>
-                              <SelectGroup_Shadcn_>
-                                {Object.keys(queryResult[0] || {}).map((key) => (
-                                  <SelectItem_Shadcn_ value={key} key={key}>
-                                    {key}
-                                  </SelectItem_Shadcn_>
-                                ))}
-                              </SelectGroup_Shadcn_>
-                            </SelectContent_Shadcn_>
-                          </Select_Shadcn_>
-
-                          <Select_Shadcn_
-                            value={item.chartConfig?.yKey}
-                            onValueChange={(value) =>
-                              onUpdateChartConfig(item, { ...item.chartConfig, yKey: value })
-                            }
-                          >
-                            <SelectTrigger_Shadcn_>
-                              Y Axis {item.chartConfig?.yKey && `- ${item.chartConfig.yKey}`}
-                            </SelectTrigger_Shadcn_>
-                            <SelectContent_Shadcn_>
-                              <SelectGroup_Shadcn_>
-                                {Object.keys(queryResult[0] || {}).map((key) => (
-                                  <SelectItem_Shadcn_ value={key} key={key}>
-                                    {key}
-                                  </SelectItem_Shadcn_>
-                                ))}
-                              </SelectGroup_Shadcn_>
-                            </SelectContent_Shadcn_>
-                          </Select_Shadcn_>
-
-                          <div className="*:flex *:gap-2 *:items-center grid gap-2 *:text-foreground-light *:p-1.5 *:pl-0">
-                            <Label_Shadcn_ htmlFor="cumulative">
-                              <Checkbox_Shadcn_
-                                id="cumulative"
-                                checked={item.chartConfig?.cumulative}
-                                onClick={() =>
-                                  onUpdateChartConfig(item, {
-                                    ...item.chartConfig,
-                                    cumulative: !item.chartConfig?.cumulative,
-                                  })
-                                }
-                              />
-                              Cumulative
-                            </Label_Shadcn_>
-                          </div>
-                        </>
-                      )}
-                    </form>
-                  </PopoverContent_Shadcn_>
-                </Popover_Shadcn_>
-              </>
-            )}
-            <Button
-              onClick={handleExecute}
-              loading={isExecuting}
-              icon={<Play size={14} />}
-              type="outline"
-              size="tiny"
-              className="w-7 h-7"
-            />
-
-            {!disableUpdate && (
-              <ButtonTooltip
-                type="outline"
-                size="tiny"
-                icon={<X />}
-                className="w-7 h-7"
-                onClick={() => onRemoveChart({ metric: { key: item.attribute, id: item.id } })}
-                tooltip={{ content: { side: 'bottom', text: 'Remove chart' } }}
-              />
-            )}
-          </div>
-        </div>
-        {item.isChart && queryResult && item.chartConfig?.xKey && item.chartConfig?.yKey ? (
-          <div className="flex-1">
-            <ChartContainer config={{}} className="aspect-auto h-full">
-              <BarChart
-                accessibilityLayer
-                data={
-                  item.chartConfig?.cumulative
-                    ? getCumulativeResults({ rows: queryResult }, item.chartConfig)
-                    : queryResult
-                }
-              >
-                <CartesianGrid vertical={false} />
-                <XAxis
-                  dataKey={item.chartConfig?.xKey}
-                  tickLine={false}
-                  tickMargin={10}
-                  axisLine={false}
-                />
-                <ChartTooltip content={<ChartTooltipContent />} />
-                <Bar dataKey={item.chartConfig?.yKey} fill="var(--chart-1)" radius={4} />
-              </BarChart>
-            </ChartContainer>
-          </div>
-        ) : (
-          <>{queryResult && <Results rows={queryResult} />}</>
-        )}
-      </div>
-    )
-  }
+  console.log('localParameters', label, parameterValues)
 
   return (
-    <ChartHandler
-      startDate={startDate}
-      endDate={endDate}
-      interval={interval}
-      attribute={item.attribute}
-      provider={item.provider}
-      label={`${item.label}${ref !== state.selectedDatabaseId ? (item.provider === 'infra-monitoring' ? ' of replica' : ' on project') : ''}`}
-      customDateFormat={'MMM D, YYYY'}
-    >
-      {!disableUpdate && (
-        <ButtonTooltip
-          type="outline"
-          icon={<X />}
-          className="w-7 h-7"
-          onClick={() => onRemoveChart({ metric: { key: item.attribute } })}
-          tooltip={{ content: { side: 'bottom', text: 'Remove chart' } }}
+    <div className="h-full flex flex-col overflow-hidden bg-surface-100 border-overlay rounded border shadow-sm">
+      <div className="flex pl-4 py-2 pr-2 items-center gap-2 z-10 shrink-0">
+        <SQL_ICON
+          className={cn(
+            'transition-colors',
+            'fill-foreground-muted',
+            'group-aria-selected:fill-foreground',
+            'w-5 h-5 shrink-0 grow-0',
+            '-ml-0.5'
+          )}
+          size={16}
+          strokeWidth={1.5}
         />
+        <h3 className="text-sm text-foreground-light flex-1">{label}</h3>
+        <div className="flex">
+          <Button
+            icon={<Code size={14} />}
+            type="text"
+            size="tiny"
+            className="w-7 h-7"
+            onClick={() => setShowSql(!showSql)}
+          />
+          {queryResult && (
+            <>
+              {parameters.length > 0 && (
+                <ParametersPopover
+                  parameters={parameters}
+                  parameterValues={combinedParameterValues}
+                  onSubmit={handleParametersSubmit}
+                />
+              )}
+
+              <Popover_Shadcn_ modal={false}>
+                <PopoverTrigger_Shadcn_ asChild>
+                  <Button
+                    icon={<Settings2 size={14} />}
+                    type="text"
+                    size="tiny"
+                    className="w-7 h-7"
+                  />
+                </PopoverTrigger_Shadcn_>
+                <PopoverContent_Shadcn_ side="bottom" align="end" className="w-[300px] p-3">
+                  <form className="grid gap-2">
+                    <ToggleGroup
+                      type="single"
+                      value={isChart ? 'chart' : 'table'}
+                      className="w-full"
+                      onValueChange={(value) => {
+                        if (value) handleToggleChart()
+                      }}
+                    >
+                      <ToggleGroupItem className="w-full" value="table" aria-label="Show as table">
+                        <Table className="h-4 w-4" />
+                      </ToggleGroupItem>
+                      <ToggleGroupItem className="w-full" value="chart" aria-label="Show as chart">
+                        <BarChart2 className="h-4 w-4" />
+                      </ToggleGroupItem>
+                    </ToggleGroup>
+
+                    {isChart && (
+                      <>
+                        <Select_Shadcn_
+                          value={chartConfig?.xKey}
+                          onValueChange={(value) =>
+                            handleUpdateChartConfig({ ...chartConfig, xKey: value })
+                          }
+                        >
+                          <SelectTrigger_Shadcn_>
+                            X Axis {chartConfig?.xKey && `- ${chartConfig.xKey}`}
+                          </SelectTrigger_Shadcn_>
+                          <SelectContent_Shadcn_>
+                            <SelectGroup_Shadcn_>
+                              {Object.keys(queryResult[0] || {}).map((key) => (
+                                <SelectItem_Shadcn_ value={key} key={key}>
+                                  {key}
+                                </SelectItem_Shadcn_>
+                              ))}
+                            </SelectGroup_Shadcn_>
+                          </SelectContent_Shadcn_>
+                        </Select_Shadcn_>
+
+                        <Select_Shadcn_
+                          value={chartConfig?.yKey}
+                          onValueChange={(value) =>
+                            handleUpdateChartConfig({ ...chartConfig, yKey: value })
+                          }
+                        >
+                          <SelectTrigger_Shadcn_>
+                            Y Axis {chartConfig?.yKey && `- ${chartConfig.yKey}`}
+                          </SelectTrigger_Shadcn_>
+                          <SelectContent_Shadcn_>
+                            <SelectGroup_Shadcn_>
+                              {Object.keys(queryResult[0] || {}).map((key) => (
+                                <SelectItem_Shadcn_ value={key} key={key}>
+                                  {key}
+                                </SelectItem_Shadcn_>
+                              ))}
+                            </SelectGroup_Shadcn_>
+                          </SelectContent_Shadcn_>
+                        </Select_Shadcn_>
+
+                        <div className="*:flex *:gap-2 *:items-center grid gap-2 *:text-foreground-light *:p-1.5 *:pl-0">
+                          <Label_Shadcn_ htmlFor="cumulative">
+                            <Checkbox_Shadcn_
+                              id="cumulative"
+                              checked={chartConfig?.cumulative}
+                              onClick={() =>
+                                handleUpdateChartConfig({
+                                  ...chartConfig,
+                                  cumulative: !chartConfig?.cumulative,
+                                })
+                              }
+                            />
+                            Cumulative
+                          </Label_Shadcn_>
+                        </div>
+                      </>
+                    )}
+                  </form>
+                </PopoverContent_Shadcn_>
+              </Popover_Shadcn_>
+            </>
+          )}
+
+          {actions}
+
+          <ButtonTooltip
+            type="text"
+            size="tiny"
+            className="w-7 h-7"
+            icon={<Play size={14} />}
+            loading={isExecuting || isLoading}
+            disabled={isLoading}
+            onClick={handleExecute}
+            tooltip={{
+              content: {
+                side: 'bottom',
+                className: 'max-w-56 text-center',
+                text: isExecuting ? (
+                  <p>{`Query is running. You may cancel ongoing queries via the [SQL Editor](/project/${ref}/sql?viewOngoingQueries=true).`}</p>
+                ) : (
+                  'Run query'
+                ),
+              },
+            }}
+          />
+
+          {!disableUpdate && onRemoveChart && (
+            <ButtonTooltip
+              type="text"
+              size="tiny"
+              icon={<X />}
+              className="w-7 h-7"
+              onClick={() => onRemoveChart({ metric: { key: id, id: id } })}
+              tooltip={{ content: { side: 'bottom', text: 'Remove chart' } }}
+            />
+          )}
+        </div>
+      </div>
+      {showWarning && (
+        <Admonition
+          type="warning"
+          className="mb-0 rounded-none border-0 shrink-0 bg-background-100 border-t"
+        >
+          <p>
+            {showWarning === 'hasWriteOperation'
+              ? 'This query contains write operations.'
+              : 'This query involves running a function.'}{' '}
+            Are you sure you want to execute it?
+          </p>
+          <p className="text-foreground-light">
+            Make sure you are not accidentally removing something important.
+          </p>
+          <div className="flex justify-stretch mt-2 gap-2">
+            <Button
+              type="outline"
+              size="tiny"
+              className="w-full flex-1"
+              onClick={() => setShowWarning(undefined)}
+            >
+              Cancel
+            </Button>
+            <Button
+              type="danger"
+              size="tiny"
+              className="w-full flex-1"
+              onClick={() => {
+                setShowWarning(undefined)
+                const processedSql = processParameterizedSql(sql!, combinedParameterValues)
+                execute({
+                  projectRef: ref,
+                  connectionString: project?.connectionString,
+                  sql: processedSql,
+                })
+              }}
+            >
+              Run
+            </Button>
+          </div>
+        </Admonition>
       )}
-    </ChartHandler>
+      {showSql && (
+        <div className="shrink-0 max-h-96 overflow-y-auto border-t">
+          <CodeBlock
+            hideLineNumbers
+            wrapLines={false}
+            value={sql}
+            language="sql"
+            className={cn(
+              'max-w-none block !bg-transparent !py-3 !px-3.5 prose dark:prose-dark border-0 border-t text-foreground !rounded-none w-full',
+              '[&>code]:m-0 [&>code>span]:text-foreground'
+            )}
+          />
+        </div>
+      )}
+      {isChart && queryResult && chartConfig?.xKey && chartConfig?.yKey ? (
+        <div className={cn('flex-1 shrink-0 border-t p-4')}>
+          <ChartContainer
+            config={{}}
+            className="aspect-auto h-full"
+            style={{
+              height: maxHeight ? `${maxHeight}px` : undefined,
+              minHeight: maxHeight ? `${maxHeight}px` : undefined,
+            }}
+          >
+            <BarChart
+              accessibilityLayer
+              data={
+                chartConfig?.cumulative
+                  ? getCumulativeResults({ rows: queryResult }, chartConfig)
+                  : queryResult
+              }
+            >
+              <CartesianGrid vertical={false} />
+              <XAxis
+                dataKey={chartConfig?.xKey}
+                tickLine={false}
+                tickMargin={10}
+                axisLine={false}
+              />
+              <ChartTooltip content={<ChartTooltipContent />} />
+              <Bar dataKey={chartConfig?.yKey} fill="var(--chart-1)" radius={4} />
+            </BarChart>
+          </ChartContainer>
+        </div>
+      ) : (
+        <>
+          {queryResult && (
+            <div
+              className={cn('flex-1 overflow-auto relative border-t')}
+              style={{ maxHeight: maxHeight ? `${maxHeight}px` : undefined }}
+            >
+              <Results rows={queryResult} />
+            </div>
+          )}
+        </>
+      )}
+    </div>
   )
 }
 
diff --git a/apps/studio/components/interfaces/Reports/QueryBlockWrapper.tsx b/apps/studio/components/interfaces/Reports/QueryBlockWrapper.tsx
new file mode 100644
index 00000000000..cd59ec1dda4
--- /dev/null
+++ b/apps/studio/components/interfaces/Reports/QueryBlockWrapper.tsx
@@ -0,0 +1,57 @@
+import { useParams } from 'common'
+import { useContentIdQuery } from 'data/content/content-id-query'
+import { useState, useEffect } from 'react'
+import QueryBlock from './QueryBlock'
+import { ChartConfig } from '../SQLEditor/UtilityPanel/ChartConfig'
+import { Parameter } from 'lib/sql-parameters'
+
+interface QueryBlockWrapperProps {
+  item: {
+    id: string
+    label: string
+    isSnippet?: boolean
+    sql?: string
+    isChart: boolean
+    chartConfig?: ChartConfig
+    results?: any
+  }
+  disableUpdate: boolean
+  onRemoveChart: ({ metric }: { metric: { key: string } }) => void
+  startDate?: string
+  endDate?: string
+  interval?: string
+  onToggleChart: () => void
+  onUpdateChartConfig: (config: ChartConfig) => void
+  parameterValues?: Record<string, string>
+  onSetParameter?: (params: Parameter[]) => void
+}
+
+const QueryBlockWrapper = ({ item, ...restProps }: QueryBlockWrapperProps) => {
+  const { ref } = useParams()
+  const [sql, setSql] = useState(item.sql || '')
+
+  const { data: snippetData } = useContentIdQuery(
+    { projectRef: ref, id: item.id },
+    { enabled: item.isSnippet }
+  )
+
+  useEffect(() => {
+    if (snippetData?.content?.sql) {
+      setSql(snippetData.content.sql)
+    }
+  }, [snippetData])
+
+  return (
+    <QueryBlock
+      sql={sql}
+      isChart={item.isChart}
+      chartConfig={item.chartConfig}
+      results={item.results}
+      label={item.label}
+      id={item.id}
+      {...restProps}
+    />
+  )
+}
+
+export default QueryBlockWrapper
diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.tsx
index 9cada1b99f5..10f058a8795 100644
--- a/apps/studio/components/ui/AIAssistantPanel/Message.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/Message.tsx
@@ -7,6 +7,8 @@ import remarkGfm from 'remark-gfm'
 import { AiIconAnimation, cn, CodeBlock, markdownComponents, WarningIcon } from 'ui'
 import CollapsibleCodeBlock from './CollapsibleCodeBlock'
 import { SqlSnippet } from './SqlSnippet'
+import QueryBlock from 'components/interfaces/Reports/QueryBlock'
+import QueryBlockWithProps from './QueryBlockWithProps'
 
 interface MessageProps {
   id: string
@@ -75,10 +77,10 @@ export const Message = function Message({
                         hideLineNumbers
                       />
                     ) : (
-                      <SqlSnippet
+                      <QueryBlockWithProps
+                        sql={props.children[0].props.children[0]}
                         readOnly={readOnly}
                         isLoading={isLoading}
-                        sql={props.children[0].props.children}
                       />
                     )
                   ) : (
diff --git a/apps/studio/components/ui/AIAssistantPanel/QueryBlockWithProps.tsx b/apps/studio/components/ui/AIAssistantPanel/QueryBlockWithProps.tsx
new file mode 100644
index 00000000000..a657d326cd2
--- /dev/null
+++ b/apps/studio/components/ui/AIAssistantPanel/QueryBlockWithProps.tsx
@@ -0,0 +1,130 @@
+import { useState } from 'react'
+import QueryBlock from 'components/interfaces/Reports/QueryBlock'
+import { ChartConfig } from 'components/interfaces/SQLEditor/UtilityPanel/ChartConfig'
+import { useRouter } from 'next/router'
+import { Edit } from 'lucide-react'
+import { useSqlEditorV2StateSnapshot } from 'state/sql-editor-v2'
+import useNewQuery from 'components/interfaces/SQLEditor/hooks'
+import { DiffType } from 'components/interfaces/SQLEditor/SQLEditor.types'
+import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
+import { TelemetryActions } from 'lib/constants/telemetry'
+import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger } from 'ui'
+
+import { ButtonTooltip } from 'components/ui/ButtonTooltip'
+
+interface QueryBlockWithPropsProps {
+  sql: string
+  isLoading?: boolean
+  readOnly?: boolean
+}
+
+const DEFAULT_CHART_CONFIG: ChartConfig = {
+  type: 'bar',
+  cumulative: false,
+  xKey: '',
+  yKey: '',
+  showLabels: false,
+  showGrid: false,
+}
+
+const QueryBlockWithProps = ({
+  sql,
+  isLoading = false,
+  readOnly = false,
+}: QueryBlockWithPropsProps) => {
+  const router = useRouter()
+  const snapV2 = useSqlEditorV2StateSnapshot()
+  const { newQuery } = useNewQuery()
+  const { mutate: sendEvent } = useSendEventMutation()
+
+  const isInSQLEditor = router.pathname.includes('/sql')
+  const isInNewSnippet = router.pathname.endsWith('/sql')
+
+  const formatted = sql || ''
+  const propsMatch = formatted.match(/--\s*props:\s*(\{[^}]+\})/)
+  const props = propsMatch ? JSON.parse(propsMatch[1]) : {}
+  const title = props.title || 'SQL Query'
+  const updatedSql = formatted?.replace(/--\s*props:\s*\{[^}]+\}/, '').trim()
+
+  const [isChart, setIsChart] = useState(props.isChart === 'true')
+  const [chartConfig, setChartConfig] = useState<ChartConfig>({
+    ...DEFAULT_CHART_CONFIG,
+    xKey: props.xAxis || '',
+    yKey: props.yAxis || '',
+  })
+
+  const handleToggleChart = () => {
+    setIsChart(!isChart)
+  }
+
+  const handleUpdateChartConfig = (config: Partial<ChartConfig>) => {
+    setChartConfig((prev) => ({ ...prev, ...config }))
+  }
+
+  const handleEditInSQLEditor = () => {
+    if (isInSQLEditor) {
+      snapV2.setDiffContent(updatedSql, DiffType.Addition)
+    } else {
+      newQuery(updatedSql, title)
+    }
+  }
+
+  const editButton =
+    !isInSQLEditor || isInNewSnippet ? (
+      <ButtonTooltip
+        type="text"
+        size="tiny"
+        className="w-7 h-7"
+        icon={<Edit size={14} />}
+        onClick={() => {
+          handleEditInSQLEditor()
+          sendEvent({ action: TelemetryActions.ASSISTANT_EDIT_SQL_CLICKED })
+        }}
+        tooltip={{ content: { side: 'bottom', text: 'Edit in SQL Editor' } }}
+      />
+    ) : (
+      <DropdownMenu>
+        <DropdownMenuTrigger asChild>
+          <ButtonTooltip
+            type="text"
+            size="tiny"
+            className="w-7 h-7"
+            icon={<Edit size={14} />}
+            tooltip={{ content: { side: 'bottom', text: 'Edit in SQL Editor' } }}
+          />
+        </DropdownMenuTrigger>
+        <DropdownMenuContent className="w-36">
+          <DropdownMenuItem onClick={() => snapV2.setDiffContent(updatedSql, DiffType.Addition)}>
+            Insert code
+          </DropdownMenuItem>
+          <DropdownMenuItem
+            onClick={() => snapV2.setDiffContent(updatedSql, DiffType.Modification)}
+          >
+            Replace code
+          </DropdownMenuItem>
+          <DropdownMenuItem onClick={() => newQuery(updatedSql, title)}>
+            Create new snippet
+          </DropdownMenuItem>
+        </DropdownMenuContent>
+      </DropdownMenu>
+    )
+
+  return (
+    <QueryBlock
+      sql={updatedSql}
+      isChart={isChart}
+      chartConfig={chartConfig}
+      label={title}
+      id={`sql-${Math.random().toString(36).substr(2, 9)}`}
+      disableUpdate={readOnly}
+      onToggleChart={handleToggleChart}
+      onUpdateChartConfig={handleUpdateChartConfig}
+      isLoading={isLoading}
+      maxHeight={300}
+      actions={editButton}
+      runQuery={props.runQuery === 'true'}
+    />
+  )
+}
+
+export default QueryBlockWithProps