From dcd2733733de2066e72f206f4c8d3efafdb206be Mon Sep 17 00:00:00 2001
From: Joshen Lim <joshenlimek@gmail.com>
Date: Mon, 10 Apr 2023 12:24:59 +0700
Subject: [PATCH] Add SQL formatter

---
 package-lock.json                             | 25 +++++++++++++++++++
 packages/ui/package.json                      |  3 ++-
 .../ui/src/components/Command/GenerateSQL.tsx | 22 +++++++++++++++-
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 24455884d80..9879567fa79 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -38941,6 +38941,7 @@
       "devDependencies": {
         "@types/react": "^17.0.37",
         "@types/react-dom": "^17.0.11",
+        "sql-formatter": "^12.2.0",
         "typescript": "4.3.4"
       }
     },
@@ -38986,6 +38987,19 @@
         "node": ">=4"
       }
     },
+    "packages/ui/node_modules/sql-formatter": {
+      "version": "12.2.0",
+      "resolved": "https://registry.npmjs.org/sql-formatter/-/sql-formatter-12.2.0.tgz",
+      "integrity": "sha512-wNsPUdOD6nnN9RUgHlNprQtm+iLW5LTOy/T0/2DDr2UeWSP8mvlQHrx6TY+IG1nfu5Kipq9GaOtS9SVq8s0Vig==",
+      "dev": true,
+      "dependencies": {
+        "argparse": "^2.0.1",
+        "nearley": "^2.20.1"
+      },
+      "bin": {
+        "sql-formatter": "bin/sql-formatter-cli.cjs"
+      }
+    },
     "packages/ui/node_modules/typescript": {
       "version": "4.3.4",
       "dev": true,
@@ -67603,6 +67617,7 @@
         "react-markdown": "^8.0.3",
         "react-syntax-highlighter": "^15.5.0",
         "remark-gfm": "^3.0.1",
+        "sql-formatter": "*",
         "sse.js": "^0.6.1",
         "tailwind-merge": "^1.10.0",
         "tailwindcss": "^3.0.15",
@@ -67640,6 +67655,16 @@
             "util-deprecate": "^1.0.2"
           }
         },
+        "sql-formatter": {
+          "version": "12.2.0",
+          "resolved": "https://registry.npmjs.org/sql-formatter/-/sql-formatter-12.2.0.tgz",
+          "integrity": "sha512-wNsPUdOD6nnN9RUgHlNprQtm+iLW5LTOy/T0/2DDr2UeWSP8mvlQHrx6TY+IG1nfu5Kipq9GaOtS9SVq8s0Vig==",
+          "dev": true,
+          "requires": {
+            "argparse": "^2.0.1",
+            "nearley": "^2.20.1"
+          }
+        },
         "typescript": {
           "version": "4.3.4",
           "dev": true
diff --git a/packages/ui/package.json b/packages/ui/package.json
index f9654a7f834..d0e4f7a2fe2 100644
--- a/packages/ui/package.json
+++ b/packages/ui/package.json
@@ -33,8 +33,8 @@
     "prop-types": "^15.7.2",
     "react-copy-to-clipboard": "^5.1.0",
     "react-feather": "^2.0.10",
-    "react-syntax-highlighter": "^15.5.0",
     "react-markdown": "^8.0.3",
+    "react-syntax-highlighter": "^15.5.0",
     "remark-gfm": "^3.0.1",
     "sse.js": "^0.6.1",
     "tailwind-merge": "^1.10.0",
@@ -44,6 +44,7 @@
   "devDependencies": {
     "@types/react": "^17.0.37",
     "@types/react-dom": "^17.0.11",
+    "sql-formatter": "^12.2.0",
     "typescript": "4.3.4"
   }
 }
diff --git a/packages/ui/src/components/Command/GenerateSQL.tsx b/packages/ui/src/components/Command/GenerateSQL.tsx
index 2c05d5d4756..d5826b00f9f 100644
--- a/packages/ui/src/components/Command/GenerateSQL.tsx
+++ b/packages/ui/src/components/Command/GenerateSQL.tsx
@@ -1,4 +1,5 @@
 import { SSE } from 'sse.js'
+import { format } from 'sql-formatter'
 import type { CreateCompletionResponse } from 'openai'
 import { useCallback, useEffect, useReducer, useRef, useState } from 'react'
 import {
@@ -250,12 +251,31 @@ Postgres SQL query:
     }
   }, [])
 
+  const formatAnswer = (answer: string) => {
+    try {
+      return format(answer, {
+        language: 'postgresql',
+        keywordCase: 'lower',
+      })
+    } catch (error: any) {
+      return answer
+    }
+  }
+
   return (
     <div onClick={(e) => e.stopPropagation()}>
       <div className={cn('relative mb-[62px] py-4 max-h-[720px] overflow-auto')}>
         {promptData.map((prompt, i) => {
           if (!prompt.query) return <></>
 
+          const formattedPromptAnwswer = (prompt?.answer ?? '')
+            .replace(/`/g, '')
+            .replace(/sql\n/g, '')
+            .trim()
+          const promptAnswer = isResponding
+            ? formattedPromptAnwswer
+            : formatAnswer(formattedPromptAnwswer)
+
           return (
             <>
               {prompt.query && (
@@ -298,7 +318,7 @@ Postgres SQL query:
                             language="sql"
                             className="relative prose dark:prose-dark bg-scale-300 max-w-none"
                           >
-                            {(prompt?.answer ?? '').replace(/`/g, '').replace(/sql\n/g, '').trim()}
+                            {promptAnswer}
                           </CodeBlock>
                           {!isResponding && <SQLOutputActions answer={prompt.answer} />}
                         </div>