feat: add export for whole query result (#95)

* feat: add design for export * feat: implementing export results as file * using LF as default for vscode * feat: export row to csv --------- Co-authored-by: Ave Aristov <[email protected]>
outerbase · Jun 20, 2024 · ebebcf2 · ebebcf2
1 parent bf21b1b
commit ebebcf2
Show file tree

Hide file tree

Showing 8 changed files with 312 additions and 146 deletions.
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -1,4 +1,5 @@
 {
     "editor.formatOnSave": true,
-    "typescript.tsdk": "node_modules\\typescript\\lib"
+    "typescript.tsdk": "node_modules\\typescript\\lib",
+    "files.eol": "\n"
 }
diff --git a/gui/src/components/export/export-result-button.tsx b/gui/src/components/export/export-result-button.tsx
@@ -0,0 +1,77 @@
+import { Button } from "../ui/button";
+import { Popover, PopoverContent, PopoverTrigger } from "../ui/popover";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "../ui/select";
+import OptimizeTableState from "../table-optimized/OptimizeTableState";
+import { useCallback, useState } from "react";
+import { getFormatHandlers } from "@gui/lib/export-helper";
+
+export default function ExportResultButton({
+  data,
+}: {
+  data: OptimizeTableState;
+}) {
+  const [format, setFormat] = useState<string | null>(null);
+
+  const onExportClicked = useCallback(() => {
+    if (!format) return;
+
+    let content = "";
+    const headers = data.getHeaders().map((header) => header.name);
+    const records = data
+      .getAllRows()
+      .map((row) => headers.map((header) => row.raw[header]));
+
+    const tableName = "UnknownTable"; //TODO: replace with actual table name
+
+    const formatHandlers = getFormatHandlers(records, headers, tableName);
+
+    const handler = formatHandlers[format];
+    if (handler) {
+      content = handler();
+    }
+
+    const blob = new Blob([content], { type: "text/plain;charset=utf-8" });
+    const url = URL.createObjectURL(blob);
+    const a = document.createElement("a");
+    a.href = url;
+    a.download = `export.${format}`;
+    a.click();
+    URL.revokeObjectURL(url);
+  }, [format, data]);
+
+  return (
+    <Popover>
+      <PopoverTrigger>
+        <Button variant="ghost" size={"sm"}>
+          Export
+        </Button>
+      </PopoverTrigger>
+      <PopoverContent className="p-0">
+        <div className="p-4">
+          <div className="mb-2 font-bold">Export</div>
+          <Select onValueChange={setFormat}>
+            <SelectTrigger>
+              <SelectValue placeholder="Select export format" />
+            </SelectTrigger>
+            <SelectContent>
+              <SelectItem value="csv">CSV</SelectItem>
+              <SelectItem value="json">JSON</SelectItem>
+              <SelectItem value="sql">SQL</SelectItem>
+            </SelectContent>
+          </Select>
+        </div>
+        <div className="p-2 pt-0 px-4">
+          <Button size="sm" onClick={onExportClicked}>
+            Export
+          </Button>
+        </div>
+      </PopoverContent>
+    </Popover>
+  );
+}
diff --git a/gui/src/components/table-optimized/OptimizeTableState.tsx b/gui/src/components/table-optimized/OptimizeTableState.tsx
@@ -335,6 +335,10 @@ export default class OptimizeTableState {
     return !!this.data[index]?.isRemoved;
   }
 
+  getAllRows() {
+    return this.data;
+  }
+
   // ------------------------------------------------
   // Handle focus logic
   // ------------------------------------------------

diff --git a/gui/src/components/tabs/query-tab.tsx b/gui/src/components/tabs/query-tab.tsx
@@ -21,6 +21,7 @@ import { MultipleQueryProgress, multipleQuery } from "@gui/lib/multiple-query";
 import { DatabaseResultStat } from "@gui/driver";
 import ResultStats from "../result-stat";
 import isEmptyResultStats from "@gui/lib/empty-stats";
+import ExportResultButton from "../export/export-result-button";
 
 export default function QueryWindow() {
   const { schema } = useAutoComplete();
@@ -130,7 +131,12 @@ export default function QueryWindow() {
             {stats && !isEmptyResultStats(stats) && (
               <div className="shrink-0">
                 <Separator />
-                <ResultStats stats={stats} />
+                <div className="flex p-1">
+                  <ResultStats stats={stats} />
+                  <div>
+                    <ExportResultButton data={data} />
+                  </div>
+                </div>
               </div>
             )}
           </div>

diff --git a/gui/src/lib/export-helper.ts b/gui/src/lib/export-helper.ts
@@ -1,4 +1,8 @@
-import { escapeIdentity, escapeSqlValue } from "@gui/sqlite/sql-helper";
+import {
+  escapeCsvValue,
+  escapeIdentity,
+  escapeSqlValue,
+} from "@gui/sqlite/sql-helper";
 
 export function selectArrayFromIndexList<T = unknown>(
   data: T[],
@@ -44,3 +48,59 @@ export function exportRowsToExcel(records: unknown[][]) {
 
   return result.join("\r\n");
 }
+
+export function exportRowsToJson(
+  headers: string[],
+  records: unknown[][]
+): string {
+  const recordsWithBigIntAsString = records.map((record) =>
+    record.map((value) =>
+      typeof value === "bigint" ? value.toString() : value
+    )
+  );
+
+  const recordsAsObjects = recordsWithBigIntAsString.map((record) =>
+    record.reduce<Record<string, unknown>>((obj, value, index) => {
+      const header = headers[index];
+      if (header !== undefined) {
+        obj[header] = value;
+      }
+      return obj;
+    }, {})
+  );
+
+  return JSON.stringify(recordsAsObjects, null, 2);
+}
+
+export function exportRowsToCsv(
+  headers: string[],
+  records: unknown[][]
+): string {
+  const result: string[] = [];
+
+  // Add headers
+  const escapedHeaders = headers.map(escapeCsvValue);
+  const headerLine = escapedHeaders.join(",");
+  result.push(headerLine);
+
+  // Add records
+  for (const record of records) {
+    const escapedRecord = record.map(escapeCsvValue);
+    const recordLine = escapedRecord.join(",");
+    result.push(recordLine);
+  }
+
+  return result.join("\n");
+}
+
+export function getFormatHandlers(
+  records: unknown[][],
+  headers: string[],
+  tableName: string
+): Record<string, (() => string) | undefined> {
+  return {
+    csv: () => exportRowsToCsv(headers, records),
+    json: () => exportRowsToJson(headers, records),
+    sql: () => exportRowsToSqlInsert(tableName, headers, records),
+  };
+}
diff --git a/gui/src/sqlite/sql-helper.ts b/gui/src/sqlite/sql-helper.ts
@@ -137,3 +137,21 @@ export function selectStatementFromPosition(
   }
   return undefined;
 }
+
+export function escapeCsvValue(value: unknown): string {
+  if (value === null || value === undefined) {
+    return "";
+  }
+
+  const stringValue = value.toString();
+  const needsEscaping =
+    stringValue.includes(",") ||
+    stringValue.includes('"') ||
+    stringValue.includes("\n");
+
+  if (needsEscaping) {
+    return `"${stringValue.replace(/"/g, '""')}"`;
+  }
+
+  return stringValue;
+}
diff --git a/studio/src/app/client/layout.tsx b/studio/src/app/client/layout.tsx
@@ -1,10 +1,10 @@
-import "@libsqlstudio/gui/css";
-import { Fragment } from "react";
-
-export default async function RootLayout({
-  children,
-}: {
-  children: React.ReactNode;
-}) {
-  return <Fragment>{children}</Fragment>;
-}
+import "@libsqlstudio/gui/css";
+import { Fragment } from "react";
+
+export default async function RootLayout({
+  children,
+}: {
+  children: React.ReactNode;
+}) {
+  return <Fragment>{children}</Fragment>;
+}