From 9969c7fa2a9dcdf5489930e2519ba929239f020d Mon Sep 17 00:00:00 2001
From: steven <stevenlgtm@gmail.com>
Date: Thu, 23 Mar 2023 17:51:31 +0800
Subject: [PATCH] feat: implement postgres connector

---
 components/ConnectionSidebar.tsx     |   5 +-
 components/CreateConnectionModal.tsx |   4 +-
 lib/connectors/index.ts              |   7 +-
 lib/connectors/mysql/index.ts        |  18 +++--
 lib/connectors/postgres/index.ts     |  81 +++++++++++++++++++++
 package.json                         |   2 +
 pnpm-lock.yaml                       | 104 +++++++++++++++++++++++++++
 store/connection.ts                  |  17 ++++-
 8 files changed, 222 insertions(+), 16 deletions(-)
 create mode 100644 lib/connectors/postgres/index.ts

diff --git a/components/ConnectionSidebar.tsx b/components/ConnectionSidebar.tsx
index 1c1d08d..3234b74 100644
--- a/components/ConnectionSidebar.tsx
+++ b/components/ConnectionSidebar.tsx
@@ -77,7 +77,8 @@ const ConnectionSidebar = () => {
               </button>
             ))}
             <button
-              className="w-10 h-10 !mt-5 ml-2 p-2 bg-gray-50 rounded-full text-gray-500 cursor-pointer hover:opacity-100"
+              className="tooltip tooltip-right w-10 h-10 !mt-5 ml-2 p-2 bg-gray-50 rounded-full text-gray-500 cursor-pointer"
+              data-tip="Create Connection"
               onClick={() => toggleCreateConnectionModal(true)}
             >
               <Icon.Ai.AiOutlinePlus className="w-auto h-full mx-auto" />
@@ -130,7 +131,7 @@ const ConnectionSidebar = () => {
               ) : (
                 <Icon.Io5.IoChatbubbleOutline className="w-5 h-auto mr-2 opacity-80 shrink-0" />
               )}
-              <span className="truncate">{chat.title}</span>
+              <span className="truncate">{chat.title || "SQL Chat"}</span>
             </div>
           ))}
           <button
diff --git a/components/CreateConnectionModal.tsx b/components/CreateConnectionModal.tsx
index b63f98d..9da56f8 100644
--- a/components/CreateConnectionModal.tsx
+++ b/components/CreateConnectionModal.tsx
@@ -60,9 +60,7 @@ const CreateConnectionModal = (props: Props) => {
               onChange={(e) => setPartialConnection({ engineType: e.target.value as Engine })}
             >
               <option value={Engine.MySQL}>MySQL</option>
-              <option className="hidden" value={Engine.PostgreSQL}>
-                PostgreSQL
-              </option>
+              <option value={Engine.PostgreSQL}>PostgreSQL</option>
             </select>
           </div>
           <div className="w-full flex flex-col">
diff --git a/lib/connectors/index.ts b/lib/connectors/index.ts
index 44c7311..3bd882c 100644
--- a/lib/connectors/index.ts
+++ b/lib/connectors/index.ts
@@ -1,5 +1,6 @@
-import { Connection } from "@/types";
+import { Connection, Engine } from "@/types";
 import mysql from "./mysql";
+import postgres from "./postgres";
 
 export interface Connector {
   testConnection: () => Promise<boolean>;
@@ -10,8 +11,10 @@ export interface Connector {
 
 export const newConnector = (connection: Connection): Connector => {
   switch (connection.engineType) {
-    case "MYSQL":
+    case Engine.MySQL:
       return mysql(connection);
+    case Engine.PostgreSQL:
+      return postgres(connection);
     default:
       throw new Error("Unsupported engine type.");
   }
diff --git a/lib/connectors/mysql/index.ts b/lib/connectors/mysql/index.ts
index 43b1115..d426f46 100644
--- a/lib/connectors/mysql/index.ts
+++ b/lib/connectors/mysql/index.ts
@@ -1,7 +1,9 @@
-import { Connection } from "@/types";
 import mysql, { RowDataPacket } from "mysql2/promise";
+import { Connection } from "@/types";
 import { Connector } from "..";
 
+const systemDatabases = ["information_schema", "mysql", "performance_schema", "sys"];
+
 const convertToConnectionUrl = (connection: Connection): string => {
   // Connection URL format: mysql://USER:PASSWORD@HOST:PORT/DATABASE
   return `mysql://${connection.username}:${connection.password}@${connection.host}:${connection.port}/${connection.database ?? ""}`;
@@ -21,7 +23,10 @@ const testConnection = async (connection: Connection): Promise<boolean> => {
 const getDatabases = async (connection: Connection): Promise<string[]> => {
   const connectionUrl = convertToConnectionUrl(connection);
   const conn = await mysql.createConnection(connectionUrl);
-  const [rows] = await conn.query<RowDataPacket[]>("SELECT schema_name as db_name FROM information_schema.schemata;");
+  const [rows] = await conn.query<RowDataPacket[]>(
+    `SELECT schema_name as db_name FROM information_schema.schemata WHERE schema_name NOT IN (?);`,
+    [systemDatabases]
+  );
   conn.destroy();
   const databaseList = [];
   for (const row of rows) {
@@ -36,16 +41,17 @@ const getTables = async (connection: Connection, databaseName: string): Promise<
   const connectionUrl = convertToConnectionUrl(connection);
   const conn = await mysql.createConnection(connectionUrl);
   const [rows] = await conn.query<RowDataPacket[]>(
-    `SELECT TABLE_NAME as table_name FROM information_schema.tables WHERE TABLE_SCHEMA='${databaseName}' AND TABLE_TYPE='BASE TABLE';`
+    `SELECT TABLE_NAME as table_name FROM information_schema.tables WHERE TABLE_SCHEMA=? AND TABLE_TYPE='BASE TABLE';`,
+    [databaseName]
   );
   conn.destroy();
-  const databaseList = [];
+  const tableList = [];
   for (const row of rows) {
     if (row["table_name"]) {
-      databaseList.push(row["table_name"]);
+      tableList.push(row["table_name"]);
     }
   }
-  return databaseList;
+  return tableList;
 };
 
 const getTableStructure = async (connection: Connection, databaseName: string, tableName: string): Promise<string> => {
diff --git a/lib/connectors/postgres/index.ts b/lib/connectors/postgres/index.ts
new file mode 100644
index 0000000..2712438
--- /dev/null
+++ b/lib/connectors/postgres/index.ts
@@ -0,0 +1,81 @@
+import { Client } from "pg";
+import { Connection } from "@/types";
+import { Connector } from "..";
+
+const newPostgresClient = (connection: Connection) => {
+  return new Client({
+    host: connection.host,
+    port: Number(connection.port),
+    user: connection.username,
+    password: connection.password,
+    database: connection.database,
+  });
+};
+
+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 getDatabases = async (connection: Connection): Promise<string[]> => {
+  const client = newPostgresClient(connection);
+  await client.connect();
+  await client.end();
+  return [connection.database!];
+};
+
+const getTables = async (connection: Connection, databaseName: string): Promise<string[]> => {
+  const client = newPostgresClient(connection);
+  await client.connect();
+  const { rows } = await client.query(
+    `SELECT table_name FROM information_schema.tables WHERE table_schema='public' AND table_type='BASE TABLE' AND table_catalog=$1;`,
+    [databaseName]
+  );
+  await client.end();
+  const tableList = [];
+  for (const row of rows) {
+    if (row["table_name"]) {
+      tableList.push(row["table_name"]);
+    }
+  }
+  return tableList;
+};
+
+const getTableStructure = async (connection: Connection, _: string, tableName: string): Promise<string> => {
+  const client = newPostgresClient(connection);
+  await client.connect();
+  const { rows } = await client.query(
+    `SELECT column_name, data_type, is_nullable, column_default FROM information_schema.columns WHERE table_schema='public' AND table_name=$1;`,
+    [tableName]
+  );
+  await client.end();
+  const columnList = [];
+  // TODO(steven): transform it to standard schema string.
+  for (const row of rows) {
+    columnList.push(`\`${row["column_name"]}\` ${row["data_type"]} ${String(row["is_nullable"]).toUpperCase()} ${row["column_default"]},`);
+  }
+  return `CREATE TABLE \`${tableName}\` (
+    ${columnList.join("\n")}
+  );`;
+};
+
+const newConnector = (connection: Connection): Connector => {
+  return {
+    testConnection: () => testConnection(connection),
+    getDatabases: () => getDatabases(connection),
+    getTables: (databaseName: string) => getTables(connection, databaseName),
+    getTableStructure: (databaseName: string, tableName: string) => getTableStructure(connection, databaseName, tableName),
+  };
+};
+
+export default newConnector;
diff --git a/package.json b/package.json
index 24a857f..30d92fb 100644
--- a/package.json
+++ b/package.json
@@ -31,6 +31,7 @@
     "@types/lodash-es": "^4.17.7",
     "@types/marked": "^4.0.8",
     "@types/node": "^18.11.18",
+    "@types/pg": "^8.6.6",
     "@types/react": "^18.0.26",
     "@types/react-dom": "^18.0.11",
     "@types/uuid": "^9.0.1",
@@ -38,6 +39,7 @@
     "eslint": "8.20.0",
     "eslint-config-next": "12.2.3",
     "mysql2": "^3.2.0",
+    "pg": "^8.10.0",
     "postcss": "^8.4.20",
     "tailwindcss": "^3.2.4",
     "typescript": "^4.9.4"
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 518d052..0b98e78 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -6,6 +6,7 @@ specifiers:
   '@types/lodash-es': ^4.17.7
   '@types/marked': ^4.0.8
   '@types/node': ^18.11.18
+  '@types/pg': ^8.6.6
   '@types/react': ^18.0.26
   '@types/react-dom': ^18.0.11
   '@types/uuid': ^9.0.1
@@ -23,6 +24,7 @@ specifiers:
   mysql2: ^3.2.0
   next: ^13.2.4
   openai: ^3.0.0
+  pg: ^8.10.0
   postcss: ^8.4.20
   react: ^18.2.0
   react-dom: ^18.2.0
@@ -59,6 +61,7 @@ devDependencies:
   '@types/lodash-es': 4.17.7
   '@types/marked': 4.0.8
   '@types/node': 18.15.3
+  '@types/pg': 8.6.6
   '@types/react': 18.0.28
   '@types/react-dom': 18.0.11
   '@types/uuid': 9.0.1
@@ -66,6 +69,7 @@ devDependencies:
   eslint: 8.20.0
   eslint-config-next: 12.2.3_bqegqxcnsisudkhpmmezgt6uoa
   mysql2: 3.2.0
+  pg: 8.10.0
   postcss: 8.4.21
   tailwindcss: 3.2.7_postcss@8.4.21
   typescript: 4.9.5
@@ -311,6 +315,14 @@ packages:
     resolution: {integrity: sha512-p6ua9zBxz5otCmbpb5D3U4B5Nanw6Pk3PPyX05xnxbB/fRv71N7CPmORg7uAD5P70T0xmx1pzAx/FUfa5X+3cw==}
     dev: true
 
+  /@types/pg/8.6.6:
+    resolution: {integrity: sha512-O2xNmXebtwVekJDD+02udOncjVcMZQuTEQEMpKJ0ZRf5E7/9JJX3izhKUcUifBkyKpljyUM6BTgy2trmviKlpw==}
+    dependencies:
+      '@types/node': 18.15.3
+      pg-protocol: 1.6.0
+      pg-types: 2.2.0
+    dev: true
+
   /@types/prop-types/15.7.5:
     resolution: {integrity: sha512-JCB8C6SnDoQf0cNycqd/35A7MjcnK+ZTqE7judS6o7utxUCg6imJg3QK2qzHKszlTjcj2cn+NwMB2i96ubpj7w==}
 
@@ -616,6 +628,11 @@ packages:
       node-releases: 2.0.10
       update-browserslist-db: 1.0.10_browserslist@4.21.5
 
+  /buffer-writer/2.0.0:
+    resolution: {integrity: sha512-a7ZpuTZU1TRtnwyCNW3I5dc0wWNC3VR9S++Ewyk2HHZdrO3CQJqSpd+95Us590V6AL7JqUAH2IwZ/398PmNFgw==}
+    engines: {node: '>=4'}
+    dev: true
+
   /call-bind/1.0.2:
     resolution: {integrity: sha512-7O+FbCihrB5WGbFYesctwmTKae6rOiIzmz1icreWJ+0aA7LJfuqhEso2T9ncpcFtzMQtzXf2QGGueWJGTYsqrA==}
     dependencies:
@@ -2006,6 +2023,10 @@ packages:
       word-wrap: 1.2.3
     dev: true
 
+  /packet-reader/1.0.0:
+    resolution: {integrity: sha512-HAKu/fG3HpHFO0AA8WE8q2g+gBJaZ9MG7fcKk+IJPLTGAD6Psw4443l+9DGRbOIh3/aXr7Phy0TjilYivJo5XQ==}
+    dev: true
+
   /parent-module/1.0.1:
     resolution: {integrity: sha512-GQ2EWRpQV8/o+Aw8YqtfZZPfNRWZYkbidE9k5rpl/hC3vtHHBfGm2Ifi6qWV+coDGkrUKZAxE3Lot5kcsRlh+g==}
     engines: {node: '>=6'}
@@ -2031,6 +2052,62 @@ packages:
     engines: {node: '>=8'}
     dev: true
 
+  /pg-connection-string/2.5.0:
+    resolution: {integrity: sha512-r5o/V/ORTA6TmUnyWZR9nCj1klXCO2CEKNRlVuJptZe85QuhFayC7WeMic7ndayT5IRIR0S0xFxFi2ousartlQ==}
+    dev: true
+
+  /pg-int8/1.0.1:
+    resolution: {integrity: sha512-WCtabS6t3c8SkpDBUlb1kjOs7l66xsGdKpIPZsg4wR+B3+u9UAum2odSsF9tnvxg80h4ZxLWMy4pRjOsFIqQpw==}
+    engines: {node: '>=4.0.0'}
+    dev: true
+
+  /pg-pool/3.6.0_pg@8.10.0:
+    resolution: {integrity: sha512-clFRf2ksqd+F497kWFyM21tMjeikn60oGDmqMT8UBrynEwVEX/5R5xd2sdvdo1cZCFlguORNpVuqxIj+aK4cfQ==}
+    peerDependencies:
+      pg: '>=8.0'
+    dependencies:
+      pg: 8.10.0
+    dev: true
+
+  /pg-protocol/1.6.0:
+    resolution: {integrity: sha512-M+PDm637OY5WM307051+bsDia5Xej6d9IR4GwJse1qA1DIhiKlksvrneZOYQq42OM+spubpcNYEo2FcKQrDk+Q==}
+    dev: true
+
+  /pg-types/2.2.0:
+    resolution: {integrity: sha512-qTAAlrEsl8s4OiEQY69wDvcMIdQN6wdz5ojQiOy6YRMuynxenON0O5oCpJI6lshc6scgAY8qvJ2On/p+CXY0GA==}
+    engines: {node: '>=4'}
+    dependencies:
+      pg-int8: 1.0.1
+      postgres-array: 2.0.0
+      postgres-bytea: 1.0.0
+      postgres-date: 1.0.7
+      postgres-interval: 1.2.0
+    dev: true
+
+  /pg/8.10.0:
+    resolution: {integrity: sha512-ke7o7qSTMb47iwzOSaZMfeR7xToFdkE71ifIipOAAaLIM0DYzfOAXlgFFmYUIE2BcJtvnVlGCID84ZzCegE8CQ==}
+    engines: {node: '>= 8.0.0'}
+    peerDependencies:
+      pg-native: '>=3.0.1'
+    peerDependenciesMeta:
+      pg-native:
+        optional: true
+    dependencies:
+      buffer-writer: 2.0.0
+      packet-reader: 1.0.0
+      pg-connection-string: 2.5.0
+      pg-pool: 3.6.0_pg@8.10.0
+      pg-protocol: 1.6.0
+      pg-types: 2.2.0
+      pgpass: 1.0.5
+    dev: true
+
+  /pgpass/1.0.5:
+    resolution: {integrity: sha512-FdW9r/jQZhSeohs1Z3sI1yxFQNFvMcnmfuj4WBMUTxOrAyLMaTcE1aAMBiTlbMNaXvBCQuVi0R7hd8udDSP7ug==}
+    dependencies:
+      split2: 4.1.0
+    dev: true
+
   /picocolors/1.0.0:
     resolution: {integrity: sha512-1fygroTLlHu66zi26VoTDv8yRgm0Fccecssto+MhsZ0D/DGW2sm8E8AjW7NU5VVTRt5GxbeZ5qBuJr+HyLYkjQ==}
 
@@ -2122,6 +2199,28 @@ packages:
       picocolors: 1.0.0
       source-map-js: 1.0.2
 
+  /postgres-array/2.0.0:
+    resolution: {integrity: sha512-VpZrUqU5A69eQyW2c5CA1jtLecCsN2U/bD6VilrFDWq5+5UIEVO7nazS3TEcHf1zuPYO/sqGvUvW62g86RXZuA==}
+    engines: {node: '>=4'}
+    dev: true
+
+  /postgres-bytea/1.0.0:
+    resolution: {integrity: sha512-xy3pmLuQqRBZBXDULy7KbaitYqLcmxigw14Q5sj8QBVLqEwXfeybIKVWiqAXTlcvdvb0+xkOtDbfQMOf4lST1w==}
+    engines: {node: '>=0.10.0'}
+    dev: true
+
+  /postgres-date/1.0.7:
+    resolution: {integrity: sha512-suDmjLVQg78nMK2UZ454hAG+OAW+HQPZ6n++TNDUX+L0+uUlLywnoxJKDou51Zm+zTCjrCl0Nq6J9C5hP9vK/Q==}
+    engines: {node: '>=0.10.0'}
+    dev: true
+
+  /postgres-interval/1.2.0:
+    resolution: {integrity: sha512-9ZhXKM/rw350N1ovuWHbGxnGh/SNJ4cnxHiM0rxE4VN41wsg8P8zWn9hv/buK00RP4WvlOyr/RBDiptyxVbkZQ==}
+    engines: {node: '>=0.10.0'}
+    dependencies:
+      xtend: 4.0.2
+    dev: true
+
   /prelude-ls/1.2.1:
     resolution: {integrity: sha512-vkcDPrRZo1QZLbn5RLGPpg/WmIQ65qoWWhcGKf/b5eplkkarX0m9z8ppCat4mlOqUsWpyNuYgO3VRyrYHSzX5g==}
     engines: {node: '>= 0.8.0'}
@@ -2344,6 +2443,11 @@ packages:
     resolution: {integrity: sha512-R0XvVJ9WusLiqTCEiGCmICCMplcCkIwwR11mOSD9CR5u+IXYdiseeEuXCVAjS54zqwkLcPNnmU4OeJ6tUrWhDw==}
     engines: {node: '>=0.10.0'}
 
+  /split2/4.1.0:
+    resolution: {integrity: sha512-VBiJxFkxiXRlUIeyMQi8s4hgvKCSjtknJv/LVYbrgALPwf5zSKmEwV9Lst25AkvMDnvxODugjdl6KZgwKM1WYQ==}
+    engines: {node: '>= 10.x'}
+    dev: true
+
   /sqlstring/2.3.3:
     resolution: {integrity: sha512-qC9iz2FlN7DQl3+wjwn3802RTyjCx7sDvfQEXchwa6CWOx07/WVfh91gBmQ9fahw8snwGEWU3xGzOt4tFyHLxg==}
     engines: {node: '>= 0.6'}
diff --git a/store/connection.ts b/store/connection.ts
index 90c90b2..eead5cb 100644
--- a/store/connection.ts
+++ b/store/connection.ts
@@ -1,11 +1,11 @@
+import axios from "axios";
+import { uniqBy } from "lodash-es";
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
 import { Connection, Database, Engine, Table, UNKNOWN_ID } from "@/types";
 import { generateUUID } from "@/utils";
-import axios from "axios";
-import { uniqBy } from "lodash-es";
 
-export const connectionSampleData: Connection = {
+export const connectionMySQLSampleData: Connection = {
   id: UNKNOWN_ID,
   title: "",
   engineType: Engine.MySQL,
@@ -15,6 +15,17 @@ export const connectionSampleData: Connection = {
   password: "",
 };
 
+export const connectionPostgreSQLSampleData: Connection = {
+  id: UNKNOWN_ID,
+  title: "",
+  engineType: Engine.PostgreSQL,
+  host: "127.0.0.1",
+  port: "5432",
+  username: "postgres",
+  password: "",
+  database: "test",
+};
+
 interface ConnectionContext {
   connection: Connection;
   database?: Database;