feat: implement connection edit (#6)

components/ConnectionSidebar.tsx

@@ -7,13 +7,11 @@ import Icon from "./Icon";
 import EngineIcon from "./EngineIcon";
 import CreateConnectionModal from "./CreateConnectionModal";
 import SettingModal from "./SettingModal";
-import ActionConfirmModal, { ActionConfirmModalProps } from "./ActionConfirmModal";
 import EditChatTitleModal from "./EditChatTitleModal";
 
 interface State {
   showCreateConnectionModal: boolean;
   showSettingModal: boolean;
-  showDeleteConnectionModal: boolean;
   showEditChatTitleModal: boolean;
 }
 
@@ -24,10 +22,9 @@ const ConnectionSidebar = () => {
   const [state, setState] = useState<State>({
     showCreateConnectionModal: false,
     showSettingModal: false,
-    showDeleteConnectionModal: false,
     showEditChatTitleModal: false,
   });
-  const [deleteConnectionModalContext, setDeleteConnectionModalContext] = useState<ActionConfirmModalProps>();
+  const [editConnectionModalContext, setEditConnectionModalContext] = useState<Connection>();
   const [editChatTitleModalContext, setEditChatTitleModalContext] = useState<Chat>();
   const connectionList = connectionStore.connectionList;
   const currentConnectionCtx = connectionStore.currentConnectionCtx;
@@ -41,6 +38,7 @@ const ConnectionSidebar = () => {
       ...state,
       showCreateConnectionModal: show,
     });
+    setEditConnectionModalContext(undefined);
   };
 
   const toggleSettingModal = (show = true) => {
@@ -65,28 +63,12 @@ const ConnectionSidebar = () => {
     });
   };
 
-  const handleDeleteConnection = (connection: Connection) => {
+  const handleEditConnection = (connection: Connection) => {
     setState({
       ...state,
-      showDeleteConnectionModal: true,
+      showCreateConnectionModal: true,
-    });
-    setDeleteConnectionModalContext({
-      title: "Delete Connection",
-      content: "Are you sure to delete this connection?",
-      confirmButtonStyle: "btn-error",
-      close: () => {
-        setState({
-          ...state,
-          showDeleteConnectionModal: false,
-        });
-      },
-      confirm: () => {
-        connectionStore.clearConnection((item) => item.id !== connection.id);
-        if (currentConnectionCtx?.connection.id === connection.id) {
-          connectionStore.setCurrentConnectionCtx(undefined);
-        }
-      },
     });
+    setEditConnectionModalContext(connection);
   };
 
   const handleDatabaseNameSelect = async (databaseName: string) => {
@@ -145,19 +127,19 @@ const ConnectionSidebar = () => {
               {connectionList.map((connection) => (
                 <button
                   key={connection.id}
-                  className={`w-full h-14 rounded-l-lg p-2 mt-2 group ${
+                  className={`relative w-full h-14 rounded-l-lg p-2 mt-2 group ${
                     currentConnectionCtx?.connection.id === connection.id && "bg-gray-100 shadow"
                   }`}
                   onClick={() => handleConnectionSelect(connection)}
                 >
                   <span
-                    className="absolute -ml-1.5 -mt-1.5 hidden opacity-60 group-hover:block hover:opacity-80"
+                    className="absolute right-0.5 -mt-1.5 opacity-60 hidden group-hover:block hover:opacity-80"
                     onClick={(e) => {
                       e.stopPropagation();
-                      handleDeleteConnection(connection);
+                      handleEditConnection(connection);
                     }}
                   >
-                    <Icon.IoClose className="w-4 h-auto" />
+                    <Icon.FiEdit3 className="w-3.5 h-auto" />
                   </span>
                   <EngineIcon engine={connection.engineType} className="w-auto h-full mx-auto" />
                 </button>
@@ -239,7 +221,7 @@ const ConnectionSidebar = () => {
                 New Chat
               </button>
             </div>
-            <div className="sticky bottom-0 w-full flex justify-center bg-gray-100 backdrop-blur bg-opacity-60 pb-4 py-2">
+            <div className="sticky bottom-0 w-full flex justify-center bg-gray-100 backdrop-blur bg-opacity-60 pb-6 py-2">
               <a
                 href="https://discord.com/invite/huyw7gRsyA"
                 className="text-indigo-600 text-sm font-medium flex flex-row justify-center items-center hover:underline"
@@ -254,24 +236,16 @@ const ConnectionSidebar = () => {
       </aside>
 
       {createPortal(
-        <CreateConnectionModal show={state.showCreateConnectionModal} close={() => toggleCreateConnectionModal(false)} />,
+        <CreateConnectionModal
+          show={state.showCreateConnectionModal}
+          connection={editConnectionModalContext}
+          close={() => toggleCreateConnectionModal(false)}
+        />,
         document.body
       )}
 
       {createPortal(<SettingModal show={state.showSettingModal} close={() => toggleSettingModal(false)} />, document.body)}
 
-      {state.showDeleteConnectionModal &&
-        createPortal(
-          <ActionConfirmModal
-            title={deleteConnectionModalContext?.title ?? ""}
-            content={deleteConnectionModalContext?.content ?? ""}
-            confirmButtonStyle={deleteConnectionModalContext?.confirmButtonStyle ?? ""}
-            close={deleteConnectionModalContext?.close ?? (() => {})}
-            confirm={deleteConnectionModalContext?.confirm ?? (() => {})}
-          />,
-          document.body
-        )}
-
       {state.showEditChatTitleModal &&
         editChatTitleModalContext &&
         createPortal(<EditChatTitleModal close={() => toggleEditChatTitleModal(false)} chat={editChatTitleModalContext} />, document.body)}

components/CreateConnectionModal.tsx

@@ -1,13 +1,16 @@
 import { cloneDeep, head } from "lodash-es";
 import { useEffect, useState } from "react";
+import { createPortal } from "react-dom";
 import { toast } from "react-hot-toast";
 import { testConnection, useConnectionStore } from "@/store";
 import { Connection, Engine } from "@/types";
 import Icon from "./Icon";
 import DataStorageBanner from "./DataStorageBanner";
+import ActionConfirmModal from "./ActionConfirmModal";
 
 interface Props {
   show: boolean;
+  connection?: Connection;
   close: () => void;
 }
 
@@ -22,15 +25,17 @@ const defaultConnection: Connection = {
 };
 
 const CreateConnectionModal = (props: Props) => {
-  const { show, close } = props;
+  const { show, connection: editConnection, close } = props;
   const connectionStore = useConnectionStore();
   const [connection, setConnection] = useState<Connection>(defaultConnection);
+  const [showDeleteConnectionModal, setShowDeleteConnectionModal] = useState(false);
   const [isRequesting, setIsRequesting] = useState(false);
   const showDatabaseField = connection.engineType === Engine.PostgreSQL;
+  const isEditing = editConnection !== undefined;
 
   useEffect(() => {
     if (show) {
-      setConnection(defaultConnection);
+      setConnection(isEditing ? editConnection : defaultConnection);
     }
   }, [show]);
 
@@ -47,22 +52,32 @@ const CreateConnectionModal = (props: Props) => {
     }
 
     setIsRequesting(true);
-    const connectionCreate = cloneDeep(connection);
+    const tempConnection = cloneDeep(connection);
     if (!showDatabaseField) {
-      connectionCreate.database = undefined;
+      tempConnection.database = undefined;
     }
+
     try {
-      const result = await testConnection(connectionCreate);
+      await testConnection(tempConnection);
-      if (!result) {
+    } catch (error) {
       setIsRequesting(false);
-        toast.error("Failed to test connection");
+      toast.error("Failed to test connection, please check your connection settings");
       return;
     }
-      const createdConnection = connectionStore.createConnection(connectionCreate);
+
+    try {
+      let connection: Connection;
+      if (isEditing) {
+        connectionStore.updateConnection(tempConnection.id, tempConnection);
+        connection = tempConnection;
+      } else {
+        connection = connectionStore.createConnection(tempConnection);
+      }
+
       // Set the created connection as the current connection.
-      const databaseList = await connectionStore.getOrFetchDatabaseList(createdConnection);
+      const databaseList = await connectionStore.getOrFetchDatabaseList(connection, true);
       connectionStore.setCurrentConnectionCtx({
-        connection: createdConnection,
+        connection: connection,
         database: head(databaseList),
       });
     } catch (error) {
@@ -76,10 +91,19 @@ const CreateConnectionModal = (props: Props) => {
     close();
   };
 
+  const handleDeleteConnection = () => {
+    connectionStore.clearConnection((item) => item.id !== connection.id);
+    if (connectionStore.currentConnectionCtx?.connection.id === connection.id) {
+      connectionStore.setCurrentConnectionCtx(undefined);
+    }
+    close();
+  };
+
   return (
+    <>
       <div className={`modal modal-middle ${show && "modal-open"}`}>
         <div className="modal-box relative">
-        <h3 className="font-bold text-lg">Create Connection</h3>
+          <h3 className="font-bold text-lg">{isEditing ? "Edit Connection" : "Create Connection"}</h3>
           <button className="btn btn-sm btn-circle absolute right-4 top-4" onClick={close}>
             <Icon.IoMdClose className="w-5 h-auto" />
           </button>
@@ -116,6 +140,18 @@ const CreateConnectionModal = (props: Props) => {
                 onChange={(e) => setPartialConnection({ port: e.target.value })}
               />
             </div>
+            {showDatabaseField && (
+              <div className="w-full flex flex-col">
+                <label className="block text-sm font-medium text-gray-700 mb-1">Database Name</label>
+                <input
+                  type="text"
+                  placeholder="Connect database"
+                  className="input input-bordered w-full"
+                  value={connection.database}
+                  onChange={(e) => setPartialConnection({ database: e.target.value })}
+                />
+              </div>
+            )}
             <div className="w-full flex flex-col">
               <label className="block text-sm font-medium text-gray-700 mb-1">Username</label>
               <input
@@ -136,20 +172,16 @@ const CreateConnectionModal = (props: Props) => {
                 onChange={(e) => setPartialConnection({ password: e.target.value })}
               />
             </div>
-          {showDatabaseField && (
-            <div className="w-full flex flex-col">
-              <label className="block text-sm font-medium text-gray-700 mb-1">Database Name</label>
-              <input
-                type="text"
-                placeholder="Connect database"
-                className="input input-bordered w-full"
-                value={connection.database}
-                onChange={(e) => setPartialConnection({ database: e.target.value })}
-              />
           </div>
+          <div className="modal-action w-full flex flex-row justify-between items-center space-x-2">
+            <div>
+              {isEditing && (
+                <button className="btn btn-ghost" onClick={() => setShowDeleteConnectionModal(true)}>
+                  Delete
+                </button>
               )}
             </div>
-        <div className="modal-action">
+            <div className="space-x-2 flex flex-row justify-center">
               <button className="btn btn-outline" onClick={close}>
                 Close
               </button>
@@ -160,6 +192,20 @@ const CreateConnectionModal = (props: Props) => {
             </div>
           </div>
         </div>
+      </div>
+
+      {showDeleteConnectionModal &&
+        createPortal(
+          <ActionConfirmModal
+            title="Delete Connection"
+            content="Are you sure you want to delete this connection?"
+            confirmButtonStyle="btn-error"
+            close={() => setShowDeleteConnectionModal(false)}
+            confirm={() => handleDeleteConnection()}
+          />,
+          document.body
+        )}
+    </>
   );
 };
 

lib/connectors/mysql/index.ts

@@ -11,13 +11,9 @@ const convertToConnectionUrl = (connection: Connection): string => {
 
 const testConnection = async (connection: Connection): Promise<boolean> => {
   const connectionUrl = convertToConnectionUrl(connection);
-  try {
   const conn = await mysql.createConnection(connectionUrl);
   conn.destroy();
   return true;
-  } catch (error) {
-    return false;
-  }
 };
 
 const execute = async (connection: Connection, databaseName: string, statement: string): Promise<any> => {

lib/connectors/postgres/index.ts

@@ -13,18 +13,10 @@ const newPostgresClient = (connection: Connection) => {
 };
 
 const testConnection = async (connection: Connection): Promise<boolean> => {
-  if (!connection.database) {
-    return false;
-  }
-
-  try {
   const client = newPostgresClient(connection);
   await client.connect();
   await client.end();
   return true;
-  } catch (error) {
-    return false;
-  }
 };
 
 const execute = async (connection: Connection, _: string, statement: string): Promise<any> => {

pages/api/connection/test.ts

@@ -13,10 +13,10 @@ const handler = async (req: NextApiRequest, res: NextApiResponse) => {
   const connection = req.body.connection as Connection;
   try {
     const connector = newConnector(connection);
-    const result = await connector.testConnection();
+    await connector.testConnection();
-    res.status(200).json(result);
+    res.status(200).json({});
   } catch (error) {
-    res.status(400).json(false);
+    res.status(400).json({});
   }
 };
 

store/chat.ts

@@ -1,5 +1,4 @@
 import dayjs from "dayjs";
-import { uniqBy } from "lodash-es";
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
 import { Chat, Id } from "@/types";
@@ -43,13 +42,9 @@ export const useChatStore = create<ChatState>()(
       },
       setCurrentChat: (chat: Chat | undefined) => set(() => ({ currentChat: chat })),
       updateChat: (chatId: Id, chat: Partial<Chat>) => {
-        const rawChat = get().chatList.find((chat) => chat.id === chatId);
-        if (!rawChat) {
-          return;
-        }
-        Object.assign(rawChat, chat);
         set((state) => ({
-          chatList: uniqBy([...state.chatList], (chat) => chat.id),
+          ...state,
+          chatList: state.chatList.map((item) => (item.id === chatId ? { ...item, ...chat } : item)),
         }));
       },
       clearChat: (filter: (chat: Chat) => boolean) => {

store/connection.ts

@@ -2,7 +2,7 @@ import axios from "axios";
 import { uniqBy } from "lodash-es";
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
-import { Connection, Database, Table } from "@/types";
+import { Connection, Database, Engine, Table } from "@/types";
 import { generateUUID } from "@/utils";
 
 interface ConnectionContext {
@@ -10,21 +10,33 @@ interface ConnectionContext {
   database?: Database;
 }
 
+const samplePGConnection: Connection = {
+  id: "sample-pg",
+  title: "Sample PostgreSQL",
+  engineType: Engine.PostgreSQL,
+  host: "db.aqbxmomjsyqbacfsujwd.supabase.co",
+  port: "",
+  username: "readonly_user",
+  password: "bytebase-sqlchat",
+  database: "employee",
+};
+
 interface ConnectionState {
   connectionList: Connection[];
   databaseList: Database[];
   currentConnectionCtx?: ConnectionContext;
   createConnection: (connection: Connection) => Connection;
   setCurrentConnectionCtx: (connectionCtx: ConnectionContext | undefined) => void;
-  getOrFetchDatabaseList: (connection: Connection) => Promise<Database[]>;
+  getOrFetchDatabaseList: (connection: Connection, skipCache?: boolean) => Promise<Database[]>;
   getOrFetchDatabaseSchema: (database: Database) => Promise<Table[]>;
+  updateConnection: (connectionId: string, connection: Partial<Connection>) => void;
   clearConnection: (filter: (connection: Connection) => boolean) => void;
 }
 
 export const useConnectionStore = create<ConnectionState>()(
   persist(
     (set, get) => ({
-      connectionList: [],
+      connectionList: [samplePGConnection],
       databaseList: [],
       createConnection: (connection: Connection) => {
         const createdConnection = {
@@ -42,11 +54,14 @@ export const useConnectionStore = create<ConnectionState>()(
           ...state,
           currentConnectionCtx: connectionCtx,
         })),
-      getOrFetchDatabaseList: async (connection: Connection) => {
+      getOrFetchDatabaseList: async (connection: Connection, skipCache = false) => {
         const state = get();
+
+        if (!skipCache) {
           if (state.databaseList.some((database) => database.connectionId === connection.id)) {
             return state.databaseList.filter((database) => database.connectionId === connection.id);
           }
+        }
 
         const { data } = await axios.post<string[]>("/api/connection/db", {
           connection,
@@ -82,6 +97,12 @@ export const useConnectionStore = create<ConnectionState>()(
         });
         return data;
       },
+      updateConnection: (connectionId: string, connection: Partial<Connection>) => {
+        set((state) => ({
+          ...state,
+          connectionList: state.connectionList.map((item) => (item.id === connectionId ? { ...item, ...connection } : item)),
+        }));
+      },
       clearConnection: (filter: (connection: Connection) => boolean) => {
         set((state) => ({
           ...state,

store/message.ts

@@ -1,4 +1,3 @@
-import { uniqBy } from "lodash-es";
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
 import { Id, Message } from "@/types";
@@ -18,21 +17,9 @@ export const useMessageStore = create<MessageState>()(
       getState: () => get(),
       addMessage: (message: Message) => set((state) => ({ messageList: [...state.messageList, message] })),
       updateMessage: (messageId: Id, message: Partial<Message>) => {
-        const rawMessage = get().messageList.find((message) => message.id === messageId);
-        if (!rawMessage) {
-          return;
-        }
         set((state) => ({
-          messageList: uniqBy(
+          ...state,
-            [
+          messageList: state.messageList.map((item) => (item.id === messageId ? { ...item, ...message } : item)),
-              ...state.messageList,
-              {
-                ...rawMessage,
-                ...message,
-              },
-            ],
-            (message) => message.id
-          ),
         }));
       },
       clearMessage: (filter: (message: Message) => boolean) => set((state) => ({ messageList: state.messageList.filter(filter) })),