query status indicator in notebook

outerbase · Jan 16, 2025 · e1357d2 · e1357d2
1 parent 1d87f81
commit e1357d2
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 34 deletions.
diff --git a/src/extensions/notebook/notebook-block-code.tsx b/src/extensions/notebook/notebook-block-code.tsx
@@ -3,12 +3,21 @@ import { NotebookEditorBlockValue } from "./notebook-editor";
 import { NotebookVM } from "./notebook-vm";
 import JavascriptEditor from "@/components/editor/javascript-editor";
 import { Button } from "@/components/ui/button";
-import { LucideShieldAlert, PlayIcon, Terminal } from "lucide-react";
+import {
+  LucideLoader,
+  LucideShieldAlert,
+  LucideTerminal,
+  PlayIcon,
+  Terminal,
+} from "lucide-react";
 import { produce } from "immer";
 import { cn } from "@/lib/utils";
 
-interface OutputFormat {
-  type: "log";
+export interface NotebookOutputFormat {
+  id?: string;
+  type: "log" | "error" | "query";
+  sql?: string;
+  queryStatus?: "running" | "success" | "error";
   args: unknown[];
 }
 
@@ -42,25 +51,53 @@ function OutputArgItem({ value }: { value: unknown }) {
   return <span className="mr-2">{content}</span>;
 }
 
-function OutputItem({ value }: { value: OutputFormat }) {
-  const color = value.type === "log" ? "" : "text-red-400 dark:text-red-300";
+function OutputItem({ value }: { value: NotebookOutputFormat }) {
+  if (value.type === "query") {
+    let icon = <LucideLoader className="w-5 h-5 animate-spin" />;
+    let color = "";
 
-  return (
-    <div className={cn(color, "flex")}>
-      {value.type === "log" ? (
-        <div className="w-7"></div>
-      ) : (
+    if (value.queryStatus === "success") {
+      icon = <LucideTerminal className="w-5 h-5" />;
+      color = "text-green-600 dark:text-green-400";
+    } else if (value.queryStatus === "error") {
+      icon = <LucideTerminal className="w-5 h-5" />;
+      color = "text-red-400 dark:text-red-300";
+    }
+
+    return (
+      <div className="flex">
+        <div className="w-7">{icon}</div>
+        <pre className={"flex-1"}>
+          <span className={cn(color, "mr-2")}>⬤</span>
+          {value.sql}
+        </pre>
+      </div>
+    );
+  } else if (value.type === "error") {
+    return (
+      <div className="flex text-red-400 dark:text-red-300">
         <div className="w-7">
           <LucideShieldAlert className="w-5 h-5" />
         </div>
-      )}
-      <pre className="flex-1">
-        {value.args.map((argValue, argIndex) => (
-          <OutputArgItem value={argValue} key={argIndex} />
-        ))}
-      </pre>
-    </div>
-  );
+        <pre className="flex-1">
+          {value.args.map((argValue, argIndex) => (
+            <OutputArgItem value={argValue} key={argIndex} />
+          ))}
+        </pre>
+      </div>
+    );
+  } else {
+    return (
+      <div className="flex">
+        <div className="w-7"></div>
+        <pre className="flex-1">
+          {value.args.map((argValue, argIndex) => (
+            <OutputArgItem value={argValue} key={argIndex} />
+          ))}
+        </pre>
+      </div>
+    );
+  }
 }
 
 export default function NotebookBlockCode({
@@ -72,19 +109,20 @@ export default function NotebookBlockCode({
   value: NotebookEditorBlockValue;
   onChange: (value: NotebookEditorBlockValue) => void;
 }) {
-  const [output, setOutput] = useState<OutputFormat[]>([]);
+  const [output, setOutput] = useState<NotebookOutputFormat[]>([]);
 
   const onRunClick = () => {
     setOutput([]);
     vm.run(value.value, {
       complete: () => {
         console.log("Complete");
       },
-      stdOut: (data: any) => {
-        setOutput((prev) => [...prev, data]);
-      },
-      stdErr: () => {
-        console.log("Error");
+      stdOut: (data) => {
+        if (data.id) {
+          setOutput((prev) => [...prev.filter((p) => p.id !== data.id), data]);
+        } else {
+          setOutput((prev) => [...prev, data]);
+        }
       },
     });
   };

diff --git a/src/extensions/notebook/notebook-editor.tsx b/src/extensions/notebook/notebook-editor.tsx
@@ -47,7 +47,7 @@ export default function NotebookEditor() {
       value: `for(let i = 0; i < 5; i++) {
   await sleep(1000);
   const age = Math.floor(Math.random() * 100));
-  const name = "name \${i}";
+  const name = \`name \${i}\`;
   await query(\`INSERT INTO testing(name, age) VALUES ('\${name}', \${age})\`);
   console.log("Inserting", name, age);
 }`,

diff --git a/src/extensions/notebook/notebook-vm.tsx b/src/extensions/notebook/notebook-vm.tsx
@@ -1,5 +1,6 @@
 import { BaseDriver } from "@/drivers/base-driver";
 import { useEffect, useMemo } from "react";
+import { NotebookOutputFormat } from "./notebook-block-code";
 
 const workerCode = `
   let scope = {};
@@ -52,16 +53,14 @@ const workerCode = `
 
 interface RunOpions {
   complete?: () => void;
-  stdOut?: <T = any>(data: T) => void;
-  stdErr?: () => void;
+  stdOut?: (data: NotebookOutputFormat) => void;
 }
 
 export class NotebookVM {
   protected vm: Worker;
   protected driver: BaseDriver;
   protected onComplete?: () => void;
-  protected onStdOut?: <T = any>(data: T) => void;
-  protected onStdErr?: () => void;
+  protected onStdOut?: (data: NotebookOutputFormat) => void;
 
   constructor(vm: Worker, driver: BaseDriver) {
     this.vm = vm;
@@ -71,21 +70,25 @@ export class NotebookVM {
       const { type } = e.data;
 
       if (type === "log" || type === "error") {
-        if (this.onStdOut) {
-          this.onStdOut(e.data);
-        }
+        if (this.onStdOut) this.onStdOut(e.data);
       } else if (type === "query") {
+        if (this.onStdOut) this.onStdOut({ ...e.data, queryStatus: "running" });
+
         this.driver
           .query(e.data.sql)
           .then((result) => {
-            console.log("Got it result", result);
             this.vm.postMessage({
               type: "query_result",
               id: e.data.id,
               result: result,
             });
+            if (this.onStdOut)
+              this.onStdOut({ ...e.data, queryStatus: "success" });
           })
           .catch((error) => {
+            if (this.onStdOut)
+              this.onStdOut({ ...e.data, queryStatus: "error" });
+
             if (error instanceof Error) {
               this.vm.postMessage({
                 type: "query_result",
@@ -111,7 +114,6 @@ export class NotebookVM {
   run(code: string, options: RunOpions): void {
     this.onComplete = options.complete;
     this.onStdOut = options.stdOut;
-    this.onStdErr = options.stdErr;
 
     this.vm.postMessage({
       type: "eval",