feat: download tables query results as CSV (#3320)

Co-authored-by: Henry Fontanier <henry@dust.tt>
dust-tt · Jan 19, 2024 · 80ad06b · 80ad06b
1 parent c69519d
commit 80ad06b
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 3 deletions.
diff --git a/core/src/sqlite_workers/sqlite_database.rs b/core/src/sqlite_workers/sqlite_database.rs
@@ -54,7 +54,7 @@ impl SqliteDatabase {
             let time_query_start = utils::now();
 
             // Execute the query and collect results
-            let mut stmt = conn.prepare(&query).unwrap();
+            let mut stmt = conn.prepare(&query)?;
             let column_names = stmt
                 .column_names()
                 .into_iter()

diff --git a/front/components/assistant/conversation/TablesQueryAction.tsx b/front/components/assistant/conversation/TablesQueryAction.tsx
@@ -1,16 +1,21 @@
 import {
+  Button,
   ChevronDownIcon,
   ChevronRightIcon,
   Chip,
   Icon,
   Spinner,
   Tooltip,
 } from "@dust-tt/sparkle";
+import { CloudArrowDownIcon } from "@dust-tt/sparkle";
 import type { TablesQueryActionType } from "@dust-tt/types";
+import { stringify } from "csv-stringify";
 import dynamic from "next/dynamic";
-import { useState } from "react";
+import { useContext, useState } from "react";
 import { amber, emerald, slate } from "tailwindcss/colors";
 
+import { SendNotificationsContext } from "@app/components/sparkle/Notification";
+
 const SyntaxHighlighter = dynamic(
   () => import("react-syntax-highlighter").then((mod) => mod.Light),
   { ssr: false }
@@ -21,6 +26,7 @@ export default function TablesQueryAction({
 }: {
   tablesQueryAction: TablesQueryActionType;
 }) {
+  const sendNotification = useContext(SendNotificationsContext);
   const [isOutputExpanded, setIsOutputExpanded] = useState(false);
 
   // Extracting question from the params
@@ -42,6 +48,39 @@ export default function TablesQueryAction({
     return t.length > maxLength ? t.substring(0, maxLength) + "..." : t;
   };
 
+  const handleDownload = () => {
+    const results =
+      output &&
+      "results" in output &&
+      Array.isArray(output?.results) &&
+      output?.results;
+
+    if (!results) {
+      return;
+    }
+
+    const queryTitle =
+      (output && "query_title" in output && output?.query_title) ??
+      "query_results";
+
+    stringify(results, { header: true }, (err, output) => {
+      if (err) {
+        sendNotification({
+          title: "Error Downloading CSV",
+          type: "error",
+          description: `An error occurred while downloading the CSV: ${err}`,
+        });
+        return;
+      }
+      const blob = new Blob([output], { type: "text/csv" });
+      const url = URL.createObjectURL(blob);
+      const a = document.createElement("a");
+      a.href = url;
+      a.download = `${queryTitle}.csv`;
+      a.click();
+    });
+  };
+
   return (
     <>
       {question && (
@@ -52,6 +91,20 @@ export default function TablesQueryAction({
           <Tooltip label={question}>
             <Chip color="slate" label={trimText(question)} />
           </Tooltip>
+          {isOutputStepCompleted &&
+            "results" in output &&
+            Array.isArray(output?.results) &&
+            output?.results?.length > 0 && (
+              <div>
+                <Button
+                  label="Download results as CSV"
+                  variant="secondary"
+                  icon={CloudArrowDownIcon}
+                  size="xs"
+                  onClick={handleDownload}
+                />
+              </div>
+            )}
         </div>
       )}
 

diff --git a/front/package-lock.json b/front/package-lock.json
diff --git a/front/package.json b/front/package.json
@@ -40,6 +40,7 @@
     "ajv": "^8.12.0",
     "blake3": "^2.1.7",
     "csv-parse": "^5.5.2",
+    "csv-stringify": "^6.4.5",
     "dd-trace": "^3.16.0",
     "emoji-mart": "^5.5.2",
     "eventsource-parser": "^1.0.0",

diff --git a/types/src/front/lib/actions/registry.ts b/types/src/front/lib/actions/registry.ts
@@ -138,7 +138,7 @@ export const DustProdActionRegistry = createActionRegistry({
       workspaceId: PRODUCTION_DUST_APPS_WORKSPACE_ID,
       appId: "b4f205e453",
       appHash:
-        "5f814f091a270cc56e2120a21dbf5fba7b64360868cc2dea1da53af3a2b9cc3b",
+        "3a4dc976412180b7aa596de606bc1b3d7c1642c91775e56db14d6ec1c2aa27bd",
     },
     config: {
       MODEL: {