Add query input and persist it to / load from url

src/App.tsx

@@ -1,4 +1,4 @@
-import { useCallback, useState } from "react";
+import { useCallback, useEffect, useState } from "react";
 import { Table } from "./table";
 import { Button } from "~/components/ui/button";
 import { Textarea } from "~/components/ui/textarea";
@@ -24,16 +24,32 @@ const usePersistedTextfield = (fieldName: string) => {
 
 function App() {
   const [datasetUrl, setDatasetUrl] = usePersistedTextfield("datasetUrl");
+  const [datasetQuery, setDatasetQuery] = usePersistedTextfield("datasetQuery");
   const [nextDatasetUrl, setNextDatasetUrl] = useState(() => datasetUrl);
+  const [nextDatasetQuery, setNextDatasetQuery] = useState(() => datasetQuery);
   const db = useDuckDb();
-  const { loading, dataset, error, clearDataset } = useParquetTable(
-    db,
-    datasetUrl
+  const onQueryChanged = useCallback(
+    (query: string) => {
+      console.log("Query changed", query);
+      setNextDatasetQuery(query);
+    },
+    [setNextDatasetQuery]
   );
+  useEffect(() => {
+    setNextDatasetQuery(datasetQuery);
+  }, [datasetQuery]);
+  const { loading, dataset, error, clearDataset } = useParquetTable(db, {
+    datasetUrl,
+    datasetQuery,
+    setDatasetQuery,
+    onQueryChanged,
+  });
 
   const handleSubmit = (e: React.FormEvent<HTMLFormElement>) => {
     e.preventDefault();
+    setNextDatasetQuery("");
     setDatasetUrl(nextDatasetUrl);
+    setDatasetQuery("");
   };
 
   return (
@@ -67,6 +83,10 @@ function App() {
             variant={"secondary"}
             type="button"
             onClick={() => {
+              if (datasetUrl !== DEFAULT_DATASET_URL) {
+                setDatasetQuery("");
+                setNextDatasetQuery("");
+              }
               setNextDatasetUrl(DEFAULT_DATASET_URL);
               setDatasetUrl(DEFAULT_DATASET_URL);
             }}
@@ -86,12 +106,37 @@ function App() {
           </Button>
         </div>
       </form>
+      {dataset && (
+        <form
+          onSubmit={(e) => {
+            e.preventDefault();
+            setDatasetQuery(nextDatasetQuery);
+          }}
+          className="flex gap-2 w-full items-center"
+        >
+          <Textarea
+            rows={2}
+            name="textField"
+            onChange={(e) => setNextDatasetQuery(e.target.value)}
+            value={nextDatasetQuery}
+            className="w-full lg:w-full"
+            placeholder="SELECT * FROM dataset LIMIT 10;"
+          />
+          <Button type="submit" className="h-full">
+            Run Query
+          </Button>
+        </form>
+      )}
       <div className="flex flex-col gap-2 w-full max-w-full overflow-x-auto bg-secondary rounded">
         {loading && (
           <p className="text-lg font-medium self-center">Loading...</p>
         )}
         {error && <p style={{ color: "red" }}>{error}</p>}
-        {dataset && <Table data={dataset} />}
+        {dataset && (
+          <div className="flex flex-col gap-2">
+            <Table data={dataset} />
+          </div>
+        )}
       </div>
     </section>
   );

src/useParquetTable.ts

@@ -1,5 +1,5 @@
 import { AsyncDuckDB, AsyncDuckDBConnection } from "@duckdb/duckdb-wasm";
-import { useCallback, useEffect, useState } from "react";
+import { useCallback, useEffect, useRef, useState } from "react";
 
 // https://duckdb.org/docs/api/wasm/query#arrow-table-to-json
 // TODO: import the arrow lib and use the correct type
@@ -9,10 +9,40 @@ const arrowResultToJson = (arrowResult: any) => {
   return arrowResult.toArray().map((row: any) => row.toJSON());
 };
 
-export const useParquetTable = (db: AsyncDuckDB | null, datasetUrl: string) => {
+export const useParquetTable = (
+  db: AsyncDuckDB | null,
+  {
+    datasetUrl,
+    datasetQuery,
+    setDatasetQuery: _setDatasetQuery,
+    onQueryChanged,
+  }: {
+    datasetUrl: string;
+    datasetQuery: string;
+    setDatasetQuery: (query: string) => void;
+    onQueryChanged?: (query: string) => void;
+  }
+) => {
+  const [datasetLoaded, setDatasetLoaded] = useState(false);
   const [loading, setLoading] = useState(false);
+  const [querying, setQuerying] = useState(false);
   const [dataset, setDataset] = useState<null>(null);
   const [error, setError] = useState<string | null>(null);
+  const onQueryChangedRef = useRef(onQueryChanged);
+  useEffect(() => {
+    onQueryChangedRef.current = onQueryChanged;
+  }, [onQueryChanged]);
+  const datasetQueryRef = useRef(datasetQuery);
+  useEffect(() => {
+    datasetQueryRef.current = datasetQuery;
+  }, [datasetQuery]);
+  const setDatasetQuery = useCallback(
+    (query: string) => {
+      _setDatasetQuery(query);
+      onQueryChangedRef.current?.(query);
+    },
+    [_setDatasetQuery, onQueryChangedRef]
+  );
 
   const clearDataset = useCallback(() => {
     setDataset(null);
@@ -24,6 +54,7 @@ export const useParquetTable = (db: AsyncDuckDB | null, datasetUrl: string) => {
       setLoading(true);
       setDataset(null);
       setError(null);
+      setDatasetLoaded(false);
       console.log("Loading", datasetUrl);
       const status: { conn: AsyncDuckDBConnection | null; killed: boolean } = {
         conn: null,
@@ -41,13 +72,21 @@ export const useParquetTable = (db: AsyncDuckDB | null, datasetUrl: string) => {
           // https://duckdb.org/docs/api/wasm/extensions
           // https://duckdb.org/docs/api/wasm/data_ingestion#parquet
           .query(
-            `LOAD parquet;LOAD httpfs;SELECT * FROM '${datasetUrl}' LIMIT 10`
+            `LOAD parquet;LOAD httpfs;DROP TABLE IF EXISTS dataset;CREATE TABLE dataset AS SELECT * FROM '${datasetUrl}';`
           )
-          .then((result) => setDataset(arrowResultToJson(result)))
+          .then(() => {
+            setDatasetLoaded(true);
+            if (!datasetQueryRef.current) {
+              console.log("Setting default query");
+              setDatasetQuery(`SELECT * FROM dataset LIMIT 10;`);
+            }
+          })
           .catch((err) => {
             console.error(err);
             setError(err.message);
             setDataset(null);
+            setDatasetLoaded(false);
+            setDatasetQuery("");
           })
           .finally(() => {
             conn.close();
@@ -64,13 +103,48 @@ export const useParquetTable = (db: AsyncDuckDB | null, datasetUrl: string) => {
       };
     } else if (db) {
       console.log("Resetting db");
+      setDataset(null);
+      setError(null);
+      setDatasetLoaded(false);
+      setDatasetQuery("");
       db.reset();
     }
-  }, [db, datasetUrl]);
+  }, [db, datasetUrl, setDatasetQuery, datasetQueryRef]);
+
+  const performQuery = useCallback(
+    (query: string) => {
+      if (db && datasetLoaded) {
+        setQuerying(true);
+        db.connect()
+          .then((conn) => {
+            conn
+              .query(query)
+              .then((result) => setDataset(arrowResultToJson(result)))
+              .catch(setError);
+          })
+          .finally(() => setQuerying(false));
+      }
+    },
+    [db, datasetLoaded]
+  );
+
+  useEffect(() => {
+    if (datasetLoaded && datasetQuery) {
+      performQuery(datasetQuery);
+    }
+  }, [datasetLoaded, datasetQuery, performQuery]);
 
   useEffect(() => {
     loadDataset();
   }, [loadDataset]);
 
-  return { loading, dataset, error, clearDataset, loadDataset };
+  return {
+    loading,
+    dataset,
+    error,
+    clearDataset,
+    loadDataset,
+    querying,
+    performQuery,
+  };
 };