From 55a2f99e2daa5e6585f47094bb2085a9a08140ab Mon Sep 17 00:00:00 2001
From: steven <stevenlgtm@gmail.com>
Date: Tue, 28 Mar 2023 17:10:43 +0800
Subject: [PATCH] chore: update util functions

---
 components/CodeBlock.tsx        | 5 +----
 pages/api/chat.ts               | 2 +-
 pages/api/connection/execute.ts | 6 ++++++
 utils/id.ts                     | 5 +++++
 utils/index.ts                  | 8 +++-----
 utils/sql.ts                    | 3 +++
 6 files changed, 19 insertions(+), 10 deletions(-)
 create mode 100644 utils/id.ts
 create mode 100644 utils/sql.ts

diff --git a/components/CodeBlock.tsx b/components/CodeBlock.tsx
index 67a4c4a..194f2b0 100644
--- a/components/CodeBlock.tsx
+++ b/components/CodeBlock.tsx
@@ -2,6 +2,7 @@ import { toast } from "react-hot-toast";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { oneDark } from "react-syntax-highlighter/dist/cjs/styles/prism";
 import { useConnectionStore, useQueryStore } from "@/store";
+import { checkStatementIsSelect } from "@/utils";
 import Icon from "./Icon";
 
 interface Props {
@@ -9,10 +10,6 @@ interface Props {
   value: string;
 }
 
-const checkStatementIsSelect = (statement: string) => {
-  return statement.toUpperCase().trim().startsWith("SELECT");
-};
-
 export const CodeBlock = (props: Props) => {
   const { language, value } = props;
   const connectionStore = useConnectionStore();
diff --git a/pages/api/chat.ts b/pages/api/chat.ts
index 5ad5d86..dfc4e6c 100644
--- a/pages/api/chat.ts
+++ b/pages/api/chat.ts
@@ -1,4 +1,4 @@
-import { openAIApiKey } from "@/utils/openai";
+import { openAIApiKey } from "@/utils";
 import { createParser, ParsedEvent, ReconnectInterval } from "eventsource-parser";
 import { NextRequest } from "next/server";
 
diff --git a/pages/api/connection/execute.ts b/pages/api/connection/execute.ts
index 37a364e..efd3c97 100644
--- a/pages/api/connection/execute.ts
+++ b/pages/api/connection/execute.ts
@@ -1,6 +1,7 @@
 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 }
@@ -12,6 +13,11 @@ 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);
     const result = await connector.execute(db, statement);
diff --git a/utils/id.ts b/utils/id.ts
new file mode 100644
index 0000000..5e43db9
--- /dev/null
+++ b/utils/id.ts
@@ -0,0 +1,5 @@
+import { v4 as uuidv4 } from "uuid";
+
+export const generateUUID = () => {
+  return uuidv4();
+};
diff --git a/utils/index.ts b/utils/index.ts
index 5e43db9..0a36f17 100644
--- a/utils/index.ts
+++ b/utils/index.ts
@@ -1,5 +1,3 @@
-import { v4 as uuidv4 } from "uuid";
-
-export const generateUUID = () => {
-  return uuidv4();
-};
+export * from "./id";
+export * from "./openai";
+export * from "./sql";
diff --git a/utils/sql.ts b/utils/sql.ts
new file mode 100644
index 0000000..4508855
--- /dev/null
+++ b/utils/sql.ts
@@ -0,0 +1,3 @@
+export const checkStatementIsSelect = (statement: string) => {
+  return statement.toUpperCase().trim().startsWith("SELECT");
+};