Merge branch 'main' into payment

src/components/CodeBlock.tsx

@@ -6,14 +6,16 @@ import { oneDark } from "react-syntax-highlighter/dist/cjs/styles/prism";
 import { useConnectionStore, useQueryStore } from "@/store";
 import Icon from "./Icon";
 import Tooltip from "./kit/Tooltip";
+import { Id } from "@/types";
 
 interface Props {
   language: string;
   value: string;
+  messageId: Id;
 }
 
 export const CodeBlock = (props: Props) => {
-  const { language, value } = props;
+  const { language, value, messageId } = props;
   const { t } = useTranslation();
   const connectionStore = useConnectionStore();
   const queryStore = useQueryStore();
@@ -37,6 +39,7 @@ export const CodeBlock = (props: Props) => {
     queryStore.setContext({
       connection: currentConnectionCtx.connection,
       database: currentConnectionCtx.database,
+      messageId: messageId,
       statement: value,
     });
     queryStore.toggleDrawer(true);

src/components/ConversationView/MessageView.tsx

@@ -125,6 +125,7 @@ const MessageView = (props: Props) => {
                         <pre className={`${className || ""} w-full p-0 my-1`} {...props}>
                           <CodeBlock
                             key={Math.random()}
+                            messageId={message.id}
                             language={language || "SQL"}
                             value={String(child.props.children).replace(/\n$/, "")}
                             {...props}

src/components/QueryDrawer.tsx

@@ -3,7 +3,7 @@ import { useEffect, useState } from "react";
 import { toast } from "react-hot-toast";
 import { useTranslation } from "react-i18next";
 import TextareaAutosize from "react-textarea-autosize";
-import { useQueryStore } from "@/store";
+import { useQueryStore, useMessageStore } from "@/store";
 import { ExecutionResult, ResponseObject } from "@/types";
 import { checkStatementIsSelect, getMessageFromExecutionResult } from "@/utils";
 import Tooltip from "./kit/Tooltip";
@@ -16,7 +16,9 @@ import ExecutionWarningBanner from "./ExecutionView/ExecutionWarningBanner";
 const QueryDrawer = () => {
   const { t } = useTranslation();
   const queryStore = useQueryStore();
+  const messageStore = useMessageStore();
   const [executionResult, setExecutionResult] = useState<ExecutionResult | undefined>(undefined);
+  const [originalStatement, setOriginalStatement] = useState<string>("");
   const [statement, setStatement] = useState<string>("");
   const [isLoading, setIsLoading] = useState(false);
   const context = queryStore.context;
@@ -30,12 +32,24 @@ const QueryDrawer = () => {
 
     const statement = context?.statement || "";
     setStatement(statement);
+    // Save original statement when open QueryDrawer.
+    if (!originalStatement) {
+      setOriginalStatement(statement);
+    }
+
     if (statement !== "" && checkStatementIsSelect(statement)) {
       executeStatement(statement);
     }
     setExecutionResult(undefined);
   }, [context, queryStore.showDrawer]);
 
+  // Reset original statement to "" when close QueryDrawer.
+  useEffect(() => {
+    if (!queryStore.showDrawer) {
+      setOriginalStatement("");
+    }
+  }, [queryStore.showDrawer]);
+
   const executeStatement = async (statement: string) => {
     if (!statement) {
       toast.error("Please enter a statement.");
@@ -81,7 +95,12 @@ const QueryDrawer = () => {
     }
   };
 
-  const close = () => queryStore.toggleDrawer(false);
+  const close = () => {
+    if (originalStatement !== statement && context?.messageId) {
+      messageStore.updateStatement(context?.messageId, originalStatement, statement);
+    }
+    queryStore.toggleDrawer(false);
+  };
 
   return (
     <Drawer open={queryStore.showDrawer} anchor="right" className="w-full" onClose={close}>

src/store/message.ts

@@ -7,6 +7,7 @@ interface MessageState {
   getState: () => MessageState;
   addMessage: (message: Message) => void;
   updateMessage: (messageId: Id, message: Partial<Message>) => void;
+  updateStatement: (messageId: Id, originalStatement: string, replacementStatement: string) => void;
   clearMessage: (filter: (message: Message) => boolean) => void;
 }
 
@@ -22,6 +23,16 @@ export const useMessageStore = create<MessageState>()(
           messageList: state.messageList.map((item) => (item.id === messageId ? { ...item, ...message } : item)),
         }));
       },
+      updateStatement: (messageId: Id, originalStatement: string, replacementStatement: string) => {
+        if (!originalStatement) return;
+        const newMessage = get().messageList.find((message) => message.id == messageId);
+        if (!newMessage) return;
+        newMessage.content = newMessage.content.replace(originalStatement, replacementStatement);
+        set((state) => ({
+          ...state,
+          messageList: state.messageList.map((item) => (item.id === messageId ? newMessage : item)),
+        }));
+      },
       clearMessage: (filter: (message: Message) => boolean) => set((state) => ({ messageList: state.messageList.filter(filter) })),
     }),
     {

src/store/query.ts

@@ -1,11 +1,12 @@
 import { merge } from "lodash-es";
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
-import { Connection, Database, Timestamp } from "@/types";
+import { Connection, Database, Id, Timestamp } from "@/types";
 
 interface ExecuteQueryContext {
   connection: Connection;
   database?: Database;
+  messageId?: Id;
   statement: string;
 }