From 2034a1b929fbee2d0010531ff17c96a3cd03fd3c Mon Sep 17 00:00:00 2001
From: Joshen Lim <joshenlimek@gmail.com>
Date: Wed, 19 Aug 2026 16:20:03 +0800
Subject: [PATCH] Use DiffEditor for QueryCell for logs migration (#49238)

## Context

Previously we added the clickhouse logs migration banner for the Query
Cell in Notebooks
But rewriting was doing a direct swap of the content

Changes here opt to use the DiffEditor instead to maintain the same UX
for query editing that's not done by the user directly

<img width="972" height="423" alt="image"
src="https://github.com/user-attachments/assets/6863531a-3b53-4756-b134-12ae16191b80"
/>


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
  * Added a review workflow for legacy SQL rewrites.
  * View proposed rewrites in a full-editor comparison overlay.
* Accept rewrites to update and save the SQL, or discard them without
applying changes.
* **Bug Fixes**
* Prevented query execution, source changes, and visibility toggling
while a rewrite is under review.
* Prevented outdated rewrite proposals from overwriting newer SQL edits.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 .../Explorer/QueryEditor/QuerySourceMenu.tsx  |  3 +
 .../interfaces/Explorer/QueryEditor/index.tsx | 59 ++++++++++++++++---
 2 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/apps/studio/components/interfaces/Explorer/QueryEditor/QuerySourceMenu.tsx b/apps/studio/components/interfaces/Explorer/QueryEditor/QuerySourceMenu.tsx
index 450df7cec6f..61f6ffb8b6a 100644
--- a/apps/studio/components/interfaces/Explorer/QueryEditor/QuerySourceMenu.tsx
+++ b/apps/studio/components/interfaces/Explorer/QueryEditor/QuerySourceMenu.tsx
@@ -26,6 +26,7 @@ import {
 import { type RoleImpersonationController } from '@/state/role-impersonation-state'
 
 export type QuerySourceMenuProps = {
+  disabled?: boolean
   rowLimit?: number
   onRowLimitChange?: (val: number) => void
   roleImpersonationState?: RoleImpersonationController
@@ -43,6 +44,7 @@ export type QuerySourceMenuProps = {
  * has SQL to preserve or discard and a fresh draft does not.
  */
 export const QuerySourceMenu = ({
+  disabled = false,
   rowLimit = 100,
   onRowLimitChange,
   roleImpersonationState,
@@ -74,6 +76,7 @@ export const QuerySourceMenu = ({
       <DropdownMenu>
         <DropdownMenuTrigger asChild>
           <Button
+            disabled={disabled}
             variant="text"
             size="tiny"
             aria-label={`Query source: ${QUERY_SOURCE_LABELS[source._tag]}`}
diff --git a/apps/studio/components/interfaces/Explorer/QueryEditor/index.tsx b/apps/studio/components/interfaces/Explorer/QueryEditor/index.tsx
index 29b970fa242..150df0f12d3 100644
--- a/apps/studio/components/interfaces/Explorer/QueryEditor/index.tsx
+++ b/apps/studio/components/interfaces/Explorer/QueryEditor/index.tsx
@@ -2,7 +2,7 @@ import { acceptUntrustedSql, untrustedSql, type UntrustedSqlFragment } from '@su
 import { useFlag } from 'common'
 import { CodeSquare, Eye, EyeOff, Play } from 'lucide-react'
 import { forwardRef, useImperativeHandle, useState, type ReactNode } from 'react'
-import { cn } from 'ui'
+import { Button, cn } from 'ui'
 
 import { resolveLogTimeRange } from '../../QuerySources/LogTimeRange.utils'
 import {
@@ -25,6 +25,7 @@ import { QueryResultRenderer } from './QueryResultRenderer'
 import { QuerySourceMenu } from './QuerySourceMenu'
 import { LegacyLogsRewriteBanner } from '@/components/interfaces/Settings/Logs/LegacyLogsRewriteBanner'
 import { CodeEditor } from '@/components/ui/CodeEditor/CodeEditor'
+import { DiffEditor } from '@/components/ui/DiffEditor'
 import {
   type DatabaseSourceParameters,
   type LogsSourceParameters,
@@ -44,6 +45,7 @@ import {
 import { useReadReplicasQuery } from '@/data/read-replicas/replicas-query'
 import { useExecuteSqlMutation } from '@/data/sql/execute-sql-mutation'
 import { applyAutoLimit } from '@/data/sql/utils'
+import { type LegacyLogsRewriteProposal } from '@/hooks/analytics/useLegacyLogsRewrite'
 import { useLatest } from '@/hooks/misc/useLatest'
 import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
 import { wrapWithRoleImpersonation } from '@/lib/role-impersonation'
@@ -129,6 +131,7 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
   const databaseIdentifier = query._tag === 'database' ? query.database_identifier : undefined
 
   const [showQuery, setShowQuery] = useState(true)
+  const [rewriteProposal, setRewriteProposal] = useState<LegacyLogsRewriteProposal | null>(null)
 
   const { data: databases, isPending: isLoadingDatabases } = useReadReplicasQuery(
     { projectRef: project?.ref },
@@ -163,7 +166,7 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
    * Postgres SQL cannot reach the analytics wire or vice versa.
    */
   const handleRunQuery = async (rawSql: string = sql) => {
-    if (!project || isBusy || rawSql.trim().length === 0) return
+    if (!project || isBusy || rewriteProposal || rawSql.trim().length === 0) return
 
     onSqlCommit?.(rawSql)
 
@@ -208,6 +211,17 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
     }).catch(() => {})
   }
 
+  const acceptRewrite = () => {
+    if (!rewriteProposal) return
+    if (sql === rewriteProposal.original) {
+      onSqlChange(rewriteProposal.modified)
+      onSqlCommit?.(rewriteProposal.modified)
+    }
+    setRewriteProposal(null)
+  }
+
+  const discardRewrite = () => setRewriteProposal(null)
+
   useImperativeHandle(ref, () => ({ run: () => handleRunQuery() }))
 
   const Shell = variant === 'viewport' ? ExplorerQueryViewport : ExplorerQuery
@@ -223,8 +237,12 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
           {toolbarActions}
           {onSourceChange && (
             <QuerySourceMenu
+              disabled={rewriteProposal !== null}
               source={toQuerySourceBinding(query)}
-              onSourceChange={onSourceChange}
+              onSourceChange={(source) => {
+                setRewriteProposal(null)
+                onSourceChange(source)
+              }}
               rowLimit={rowLimit}
               onRowLimitChange={onRowLimitChange}
               roleImpersonationState={roleImpersonationState}
@@ -241,6 +259,7 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
           )}
           <ExplorerToolbarAction
             icon={showQuery ? <EyeOff /> : <Eye />}
+            disabled={rewriteProposal !== null}
             tooltip={showQuery ? 'Hide query' : 'Show query'}
             onClick={() => setShowQuery((value) => !value)}
           />
@@ -248,7 +267,9 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
             loading={isExecuting || isLoadingProject}
             icon={<Play />}
             tooltip="Run query"
-            disabled={isLoadingProject || isExecuting || sql.trim().length === 0}
+            disabled={
+              isLoadingProject || isExecuting || rewriteProposal !== null || sql.trim().length === 0
+            }
             onClick={() => handleRunQuery()}
           >
             Run
@@ -262,10 +283,8 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
             isLogsSource={query._tag === 'logs'}
             sql={sql}
             readSql={() => sqlRef.current}
-            onProposal={({ modified }) => {
-              onSqlChange(modified)
-              onSqlCommit?.(modified)
-            }}
+            onProposal={setRewriteProposal}
+            hidden={rewriteProposal !== null}
           />
           <ExplorerQueryEditor
             className={cn('relative', variant === 'viewport' ? 'h-[45%] min-h-48' : undefined)}
@@ -284,6 +303,30 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
                 editor.onDidBlurEditorWidget(() => onSqlCommitRef.current?.(sqlRef.current))
               }}
             />
+            {rewriteProposal && (
+              <div className="absolute inset-0 z-10 flex flex-col bg-studio">
+                <div className="flex items-center justify-between gap-2 border-b bg-surface-100 px-3 py-2">
+                  <span className="text-xs text-foreground-light">
+                    Review the ClickHouse SQL rewrite before accepting it
+                  </span>
+                  <div className="flex items-center gap-2">
+                    <Button variant="default" size="tiny" onClick={discardRewrite}>
+                      Discard
+                    </Button>
+                    <Button variant="primary" size="tiny" onClick={acceptRewrite}>
+                      Accept
+                    </Button>
+                  </div>
+                </div>
+                <div className="min-h-0 flex-1">
+                  <DiffEditor
+                    language="pgsql"
+                    original={rewriteProposal.original}
+                    modified={rewriteProposal.modified}
+                  />
+                </div>
+              </div>
+            )}
           </ExplorerQueryEditor>
         </>
       )}