feat: support run non-select SQL (#35)

src/components/ExecutionView/DataTableView.tsx

@@ -0,0 +1,39 @@
+import { head } from "lodash-es";
+import DataTable from "react-data-table-component";
+import { useTranslation } from "react-i18next";
+import { RawResult } from "@/types";
+import Icon from "../Icon";
+
+interface Props {
+  rawResults: RawResult[];
+}
+
+const DataTableView = (props: Props) => {
+  const { rawResults } = props;
+  const { t } = useTranslation();
+  const columns = Object.keys(head(rawResults) || {}).map((key) => {
+    return {
+      name: key,
+      sortable: true,
+      selector: (row: any) => row[key],
+    };
+  });
+
+  return rawResults.length === 0 ? (
+    <div className="w-full flex flex-col justify-center items-center py-6 pt-10">
+      <Icon.BsBox2 className="w-7 h-auto opacity-70" />
+      <span className="text-sm font-mono text-gray-500 mt-2">{t("execution.message.no-data")}</span>
+    </div>
+  ) : (
+    <DataTable
+      className="w-full border !rounded-lg dark:border-zinc-700"
+      columns={columns}
+      data={rawResults}
+      fixedHeader
+      pagination
+      responsive
+    />
+  );
+};
+
+export default DataTableView;

src/components/ExecutionView/ExecutionWarningBanner.tsx

@@ -0,0 +1,22 @@
+import { useTranslation } from "react-i18next";
+import Icon from "../Icon";
+
+interface Props {
+  className?: string;
+}
+
+const ExecutionWarningBanner = (props: Props) => {
+  const { className } = props;
+  const { t } = useTranslation();
+
+  return (
+    <div className={`${className || ""} relative w-full flex flex-row justify-start items-center px-4 py-2 bg-yellow-100 dark:bg-zinc-700`}>
+      <span className="text-sm leading-6 pr-4">
+        <Icon.IoInformationCircleOutline className="inline-block h-5 w-auto -mt-0.5 mr-0.5 opacity-80" />
+        {t("banner.non-select-sql-warning")}
+      </span>
+    </div>
+  );
+};
+
+export default ExecutionWarningBanner;

src/components/ExecutionView/NotificationView.tsx

@@ -0,0 +1,12 @@
+interface Props {
+  message: string;
+  style: "info" | "error";
+}
+
+const NotificationView = (props: Props) => {
+  const { message, style } = props;
+  const additionalStyle = style === "error" ? "text-red-500" : "text-gray-500";
+  return <p className={`${additionalStyle} w-full pl-4 mt-4 font-mono text-sm whitespace-pre-wrap`}>{message}</p>;
+};
+
+export default NotificationView;

src/components/QueryDrawer.tsx

@@ -1,41 +1,38 @@
 import { Drawer } from "@mui/material";
-import { head } from "lodash-es";
 import { useEffect, useState } from "react";
-import DataTable from "react-data-table-component";
 import { toast } from "react-hot-toast";
 import { useTranslation } from "react-i18next";
 import TextareaAutosize from "react-textarea-autosize";
 import { useQueryStore } from "@/store";
-import { ResponseObject } from "@/types";
+import { ExecutionResult, ResponseObject } from "@/types";
+import { checkStatementIsSelect, getMessageFromExecutionResult } from "@/utils";
 import Icon from "./Icon";
 import EngineIcon from "./EngineIcon";
-
-type RawQueryResult = {
-  [key: string]: any;
-};
+import DataTableView from "./ExecutionView/DataTableView";
+import NotificationView from "./ExecutionView/NotificationView";
+import ExecutionWarningBanner from "./ExecutionView/ExecutionWarningBanner";
 
 const QueryDrawer = () => {
   const { t } = useTranslation();
   const queryStore = useQueryStore();
-  const [rawResults, setRawResults] = useState<RawQueryResult[]>([]);
+  const [executionResult, setExecutionResult] = useState<ExecutionResult | undefined>(undefined);
+  const [statement, setStatement] = useState<string>("");
+  const [isLoading, setIsLoading] = useState(false);
   const context = queryStore.context;
-  const [statement, setStatement] = useState<string>(context?.statement || "");
-  const [isLoading, setIsLoading] = useState(true);
-  const columns = Object.keys(head(rawResults) || {}).map((key) => {
-    return {
-      name: key,
-      sortable: true,
-      selector: (row: RawQueryResult) => row[key],
-    };
-  });
+  const executionMessage = executionResult ? getMessageFromExecutionResult(executionResult) : "";
+  const showExecutionWarningBanner = !checkStatementIsSelect(statement);
 
   useEffect(() => {
     if (!queryStore.showDrawer) {
       return;
     }
 
-    setStatement(context?.statement || "");
-    executeStatement(context?.statement || "");
+    const statement = context?.statement || "";
+    setStatement(statement);
+    if (statement !== "" && checkStatementIsSelect(statement)) {
+      executeStatement(statement);
+    }
+    setExecutionResult(undefined);
   }, [context, queryStore.showDrawer]);
 
   const executeStatement = async (statement: string) => {
@@ -47,12 +44,12 @@ const QueryDrawer = () => {
     if (!context) {
       toast.error("No execution context found.");
       setIsLoading(false);
-      setRawResults([]);
+      setExecutionResult(undefined);
       return;
     }
 
     setIsLoading(true);
-    setRawResults([]);
+    setExecutionResult(undefined);
     const { connection, database } = context;
     try {
       const response = await fetch("/api/connection/execute", {
@@ -66,11 +63,14 @@ const QueryDrawer = () => {
           statement,
         }),
       });
-      const result = (await response.json()) as ResponseObject<RawQueryResult[]>;
+      const result = (await response.json()) as ResponseObject<ExecutionResult>;
       if (result.message) {
-        toast.error(result.message);
+        setExecutionResult({
+          rawResult: [],
+          error: result.message,
+        });
       } else {
-        setRawResults(result.data);
+        setExecutionResult(result.data);
       }
     } catch (error) {
       console.error(error);
@@ -101,6 +101,7 @@ const QueryDrawer = () => {
               <EngineIcon className="w-6 h-auto" engine={context.connection.engineType} />
               <span>{context.database?.name}</span>
             </div>
+            {showExecutionWarningBanner && <ExecutionWarningBanner className="rounded-lg mt-4" />}
             <div className="w-full h-auto mt-4 px-2 flex flex-row justify-between items-end border dark:border-zinc-700 rounded-lg overflow-clip">
               <TextareaAutosize
                 className="w-full h-full outline-none border-none bg-transparent leading-6 pl-2 py-2 resize-none hide-scrollbar text-sm font-mono break-all"
@@ -124,22 +125,18 @@ const QueryDrawer = () => {
                   <Icon.BiLoaderAlt className="w-7 h-auto opacity-70 animate-spin" />
                   <span className="text-sm font-mono text-gray-500 mt-2">{t("execution.message.executing")}</span>
                 </div>
-              ) : rawResults.length === 0 ? (
-                <div className="w-full flex flex-col justify-center items-center py-6 pt-10">
-                  <Icon.BsBox2 className="w-7 h-auto opacity-70" />
-                  <span className="text-sm font-mono text-gray-500 mt-2">{t("execution.message.no-data")}</span>
-                </div>
               ) : (
-                <div className="w-full">
-                  <DataTable
-                    className="w-full border !rounded-lg dark:border-zinc-700"
-                    columns={columns}
-                    data={rawResults}
-                    fixedHeader
-                    pagination
-                    responsive
-                  />
-                </div>
+                <>
+                  {executionResult ? (
+                    executionMessage ? (
+                      <NotificationView message={executionMessage} style={executionResult?.error ? "error" : "info"} />
+                    ) : (
+                      <DataTableView rawResults={executionResult?.rawResult || []} />
+                    )
+                  ) : (
+                    <></>
+                  )}
+                </>
               )}
             </div>
           </>

src/lib/connectors/index.ts

@@ -1,11 +1,11 @@
-import { Connection, Engine } from "@/types";
+import { Connection, Engine, ExecutionResult } from "@/types";
 import mysql from "./mysql";
 import postgres from "./postgres";
 import mssql from "./mssql";
 
 export interface Connector {
   testConnection: () => Promise<boolean>;
-  execute: (databaseName: string, statement: string) => Promise<any>;
+  execute: (databaseName: string, statement: string) => Promise<ExecutionResult>;
   getDatabases: () => Promise<string[]>;
   getTables: (databaseName: string) => Promise<string[]>;
   getTableStructure: (databaseName: string, tableName: string) => Promise<string>;

src/lib/connectors/mssql/index.ts

@@ -1,5 +1,5 @@
 import { ConnectionPool } from "mssql";
-import { Connection } from "@/types";
+import { Connection, ExecutionResult } from "@/types";
 import { Connector } from "..";
 
 const systemDatabases = ["master", "tempdb", "model", "msdb"];
@@ -37,7 +37,12 @@ const execute = async (connection: Connection, databaseName: string, statement:
   const request = pool.request();
   const result = await request.query(`USE ${databaseName}; ${statement}`);
   await pool.close();
-  return result.recordset;
+
+  const executionResult: ExecutionResult = {
+    rawResult: result.recordset,
+    affectedRows: result.rowsAffected.length,
+  };
+  return executionResult;
 };
 
 const getDatabases = async (connection: Connection): Promise<string[]> => {

src/lib/connectors/mysql/index.ts

@@ -1,6 +1,6 @@
 import { ConnectionOptions } from "mysql2";
 import mysql, { RowDataPacket } from "mysql2/promise";
-import { Connection } from "@/types";
+import { Connection, ExecutionResult } from "@/types";
 import { Connector } from "..";
 
 const systemDatabases = ["information_schema", "mysql", "performance_schema", "sys"];
@@ -33,9 +33,19 @@ const testConnection = async (connection: Connection): Promise<boolean> => {
 const execute = async (connection: Connection, databaseName: string, statement: string): Promise<any> => {
   connection.database = databaseName;
   const conn = await getMySQLConnection(connection);
-  const [rows] = await conn.query<RowDataPacket[]>(statement);
+  const [rows] = await conn.execute(statement);
   conn.destroy();
-  return rows;
+
+  const executionResult: ExecutionResult = {
+    rawResult: [],
+    affectedRows: 0,
+  };
+  if (Array.isArray(rows)) {
+    executionResult.rawResult = rows;
+  } else {
+    executionResult.affectedRows = rows.affectedRows;
+  }
+  return executionResult;
 };
 
 const getDatabases = async (connection: Connection): Promise<string[]> => {

src/lib/connectors/postgres/index.ts

@@ -1,5 +1,5 @@
 import { Client, ClientConfig } from "pg";
-import { Connection } from "@/types";
+import { Connection, ExecutionResult } from "@/types";
 import { Connector } from "..";
 
 const newPostgresClient = (connection: Connection) => {
@@ -27,12 +27,18 @@ const testConnection = async (connection: Connection): Promise<boolean> => {
   return true;
 };
 
-const execute = async (connection: Connection, _: string, statement: string): Promise<any> => {
+const execute = async (connection: Connection, databaseName: string, statement: string): Promise<any> => {
+  connection.database = databaseName;
   const client = newPostgresClient(connection);
   await client.connect();
-  const { rows } = await client.query(statement);
+  const { rows, rowCount } = await client.query(statement);
   await client.end();
-  return rows;
+
+  const executionResult: ExecutionResult = {
+    rawResult: rows,
+    affectedRows: rowCount,
+  };
+  return executionResult;
 };
 
 const getDatabases = async (connection: Connection): Promise<string[]> => {

src/locales/en.json

@@ -57,6 +57,7 @@
     "join-wechat-group": "Join WeChat Group"
   },
   "banner": {
-    "data-storage": "Connection settings are stored in your local browser"
+    "data-storage": "Connection settings are stored in your local browser",
+    "non-select-sql-warning": "The current statement may be non-SELECT SQL, which will result in a database schema or data change. Make sure you know what you are doing."
   }
 }

src/locales/zh.json

@@ -57,6 +57,7 @@
     "join-wechat-group": "加入微信群"
   },
   "banner": {
-    "data-storage": "连接设置存储在您的本地浏览器中"
+    "data-storage": "连接设置存储在您的本地浏览器中",
+    "non-select-sql-warning": "当前语句可能是非 SELECT SQL，这将导致数据库模式或数据变化。"
   }
 }

src/pages/api/connection/execute.ts

@@ -1,7 +1,6 @@
 import { NextApiRequest, NextApiResponse } from "next";
 import { newConnector } from "@/lib/connectors";
 import { Connection } from "@/types";
-import { checkStatementIsSelect } from "@/utils";
 
 // POST /api/connection/execute
 // req body: { connection: Connection, db: string, statement: string }
@@ -13,10 +12,6 @@ const handler = async (req: NextApiRequest, res: NextApiResponse) => {
   const connection = req.body.connection as Connection;
   const db = req.body.db as string;
   const statement = req.body.statement as string;
-  // We only support SELECT statements for now.
-  if (!checkStatementIsSelect(statement)) {
-    return res.status(400).json([]);
-  }
 
   try {
     const connector = newConnector(connection);

src/types/connector.ts

@@ -0,0 +1,9 @@
+export type RawResult = {
+  [key: string]: any;
+};
+
+export interface ExecutionResult {
+  rawResult: RawResult[];
+  affectedRows?: number;
+  error?: string;
+}

src/types/index.ts

@@ -6,3 +6,4 @@ export * from "./conversation";
 export * from "./message";
 export * from "./setting";
 export * from "./api";
+export * from "./connector";

src/utils/execution.ts

@@ -0,0 +1,11 @@
+import { ExecutionResult } from "@/types";
+
+export const getMessageFromExecutionResult = (result: ExecutionResult): string => {
+  if (result.error) {
+    return result.error;
+  }
+  if (result.affectedRows) {
+    return `${result.affectedRows} rows affected.`;
+  }
+  return "";
+};

src/utils/index.ts

@@ -1,3 +1,4 @@
 export * from "./id";
 export * from "./openai";
 export * from "./sql";
+export * from "./execution";