From 30f891a41e442a01757802c449ecf1afc1c8d391 Mon Sep 17 00:00:00 2001
From: "l1xnan@qq.com" <l1xnan@qq.com>
Date: Mon, 16 Oct 2023 21:06:24 +0800
Subject: [PATCH] feat: show tables

---
 src-tauri/src/main.rs       | 31 +++++++++++-
 src/Home.tsx                | 62 ++++++++++++++++++------
 src/components/FileTree.tsx | 23 +++++----
 src/components/Table.tsx    | 94 +++++++++++++++++++++++++++++++++++++
 src/stores/store.ts         | 10 ++++
 5 files changed, 193 insertions(+), 27 deletions(-)
 create mode 100644 src/components/Table.tsx

diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs
index 6c026dc..0962660 100644
--- a/src-tauri/src/main.rs
+++ b/src-tauri/src/main.rs
@@ -68,6 +68,30 @@ fn directory_tree(path: &str) -> FileNode {
   node
 }
 
+fn _show_tables(path: String) -> Result<ValidationResponse> {
+  let db = Connection::open(path)?;
+  let mut stmt = db.prepare("SHOW TABLES").unwrap();
+  let frames = stmt.query_arrow(duckdb::params![]).unwrap();
+  let schema = frames.get_schema();
+  let records: Vec<RecordBatch> = frames.collect();
+  let record_batch = arrow::compute::concat_batches(&schema, &records).unwrap();
+  Ok(ValidationResponse {
+    row_count: records.len(),
+    total_count: records.len(),
+    preview: serialize_preview(&record_batch).unwrap(),
+  })
+}
+
+#[tauri::command]
+async fn show_tables(path: String) -> ValidationResponse {
+  let res = _show_tables(path);
+  if let Ok(data) = res {
+    data
+  } else {
+    ValidationResponse::default()
+  }
+}
+
 #[tauri::command]
 fn get_folder_tree(name: &str) -> FileNode {
   directory_tree(name)
@@ -91,7 +115,6 @@ fn _query(sql: String, limit: i32, offset: i32) -> anyhow::Result<ValidationResp
     .query_row(count_sql.as_str(), [], |row| row.get(0))
     .unwrap();
 
-
   // query
   let sql = format!("{sql} limit {limit} offset {offset}");
   let mut stmt = db.prepare(sql.as_str()).unwrap();
@@ -161,7 +184,11 @@ fn main() {
       });
       Ok(())
     })
-    .invoke_handler(tauri::generate_handler![get_folder_tree, query])
+    .invoke_handler(tauri::generate_handler![
+      get_folder_tree,
+      query,
+      show_tables
+    ])
     .run(tauri::generate_context!())
     .expect("error while running tauri application");
 }
diff --git a/src/Home.tsx b/src/Home.tsx
index 98d5693..7662422 100644
--- a/src/Home.tsx
+++ b/src/Home.tsx
@@ -11,8 +11,22 @@ import RemoveIcon from "@mui/icons-material/Remove";
 import { isDarkTheme } from "@/utils";
 import { useLocalStorageState } from "ahooks";
 import { useEffect, useState } from "react";
-import { useStore } from "@/stores/store";
-import { IconDatabasePlus, IconFolderPlus } from "@tabler/icons-react";
+import { showTables, useStore } from "@/stores/store";
+import {
+  IconDatabasePlus,
+  IconFolderPlus,
+  IconRefresh,
+} from "@tabler/icons-react";
+
+export const MuiIconButton = styled((props) => (
+  <IconButton color="inherit" {...props} />
+))<BoxProps>(({}) => ({
+  "& *": {
+    fontSize: 16,
+    height: 16,
+    width: 16,
+  },
+}));
 
 function Home() {
   const theme = useTheme();
@@ -49,8 +63,7 @@ function Home() {
           <Typography fontWeight={800}>Database Explorer</Typography>
 
           <Stack direction="row">
-            <IconButton
-              color="inherit"
+            <MuiIconButton
               onClick={async () => {
                 const res = await dialog.open({
                   directory: true,
@@ -68,9 +81,8 @@ function Home() {
               }}
             >
               <IconFolderPlus />
-            </IconButton>
-            <IconButton
-              color="inherit"
+            </MuiIconButton>
+            <MuiIconButton
               onClick={async () => {
                 const res = await dialog.open({
                   directory: false,
@@ -85,16 +97,9 @@ function Home() {
                   openDirectory(res.path);
                 }
               }}
-              sx={{
-                "& *": {
-                  fontSize: 16,
-                  height: 16,
-                  width: 16,
-                },
-              }}
             >
               <IconDatabasePlus />
-            </IconButton>
+            </MuiIconButton>
             <IconButton
               color="inherit"
               onClick={async () => {
@@ -107,6 +112,33 @@ function Home() {
             >
               <RemoveIcon />
             </IconButton>
+            <MuiIconButton
+              color="inherit"
+              onClick={async () => {
+                console.log(selectedPath);
+                if (selectedPath && selectedPath.endsWith(".duckdb")) {
+                  const res = await showTables(selectedPath);
+                  console.log(res);
+                  setFolders(
+                    folders?.map((item) => {
+                      if (item.path == selectedPath) {
+                        item.children = res.data.map(({ name }) => ({
+                          name,
+                          path: name,
+                          type: "table",
+                          children: [],
+                        }));
+                        return item;
+                      }
+                      return item;
+                    })
+                  );
+                  console.log(folders);
+                }
+              }}
+            >
+              <IconRefresh />
+            </MuiIconButton>
             <ToggleColorMode />
           </Stack>
         </ToolbarBox>
diff --git a/src/components/FileTree.tsx b/src/components/FileTree.tsx
index 2a21023..0f9a6fa 100644
--- a/src/components/FileTree.tsx
+++ b/src/components/FileTree.tsx
@@ -4,18 +4,20 @@ import { Typography } from "@mui/material";
 import { TreeItem } from "@mui/x-tree-view/TreeItem";
 import { TreeView, TreeViewProps } from "@mui/x-tree-view/TreeView";
 import {
+  IconDatabase,
   IconFile,
-  IconFileDatabase,
   IconFilePower,
   IconFileTypeCsv,
   IconFileTypeXls,
   IconFolder,
   IconFolderOpen,
+  IconTable,
 } from "@tabler/icons-react";
 import * as React from "react";
 
 export interface FileNode {
   name: string;
+  type?: string;
   path: string;
   is_dir: boolean;
   children: FileNode[];
@@ -25,19 +27,20 @@ export interface FileTreeProps extends TreeViewProps<undefined> {
   data: FileNode;
 }
 
-const getFileTypeIcon = (path: string) => {
-  const ext = path.split(".")[1];
-
-  if (ext == "duckdb") {
-    return <IconFileDatabase />;
+const getFileTypeIcon = (type: string) => {
+  if (type == "duckdb") {
+    return <IconDatabase />;
+  }
+  if (type == "table") {
+    return <IconTable />;
   }
-  if (ext == "csv") {
+  if (type == "csv") {
     return <IconFileTypeCsv />;
   }
-  if (ext == "xlsx") {
+  if (type == "xlsx") {
     return <IconFileTypeXls />;
   }
-  if (ext == "parquet") {
+  if (type == "parquet") {
     return <IconFilePower />;
   }
   return <IconFile />;
@@ -66,7 +69,7 @@ export default function FileTree({
       }
       icon={
         !node.is_dir ? (
-          getFileTypeIcon(node.path)
+          getFileTypeIcon(node?.type ?? node.path.split(".")[1])
         ) : expanded.includes(node.path) ? (
           <IconFolderOpen />
         ) : (
diff --git a/src/components/Table.tsx b/src/components/Table.tsx
new file mode 100644
index 0000000..5fca5a1
--- /dev/null
+++ b/src/components/Table.tsx
@@ -0,0 +1,94 @@
+import React from "react";
+import "./index.css";
+
+import {
+  getCoreRowModel,
+  ColumnDef,
+  flexRender,
+  useReactTable,
+} from "@tanstack/react-table";
+import { makeData, Person } from "./makeData";
+
+const columns: ColumnDef<Person>[] = [
+  {
+    header: "Name",
+    footer: (props) => props.column.id,
+  },
+  {
+    header: "Info",
+    footer: (props) => props.column.id,
+  },
+];
+
+function App() {
+  const data = React.useMemo(() => makeData(20), []);
+
+  const table = useReactTable({
+    data,
+    columns,
+    enableColumnResizing: true,
+    columnResizeMode: "onChange",
+    getCoreRowModel: getCoreRowModel(),
+    debugTable: true,
+    debugHeaders: true,
+    debugColumns: true,
+  });
+
+  return (
+    <div className="p-2 block max-w-full overflow-x-scroll overflow-y-hidden">
+      <div className="h-2" />
+      <table className="w-full ">
+        <thead>
+          {table.getHeaderGroups().map((headerGroup) => (
+            <tr key={headerGroup.id}>
+              {headerGroup.headers.map((header) => {
+                return (
+                  <th
+                    key={header.id}
+                    colSpan={header.colSpan}
+                    style={{ position: "relative", width: header.getSize() }}
+                  >
+                    {header.isPlaceholder
+                      ? null
+                      : flexRender(
+                          header.column.columnDef.header,
+                          header.getContext()
+                        )}
+                    {header.column.getCanResize() && (
+                      <div
+                        onMouseDown={header.getResizeHandler()}
+                        onTouchStart={header.getResizeHandler()}
+                        className={`resizer ${
+                          header.column.getIsResizing() ? "isResizing" : ""
+                        }`}
+                      ></div>
+                    )}
+                  </th>
+                );
+              })}
+            </tr>
+          ))}
+        </thead>
+        <tbody>
+          {table.getRowModel().rows.map((row) => {
+            return (
+              <tr key={row.id}>
+                {row.getVisibleCells().map((cell) => {
+                  return (
+                    <td key={cell.id} style={{ width: cell.column.getSize() }}>
+                      {flexRender(
+                        cell.column.columnDef.cell,
+                        cell.getContext()
+                      )}
+                    </td>
+                  );
+                })}
+              </tr>
+            );
+          })}
+        </tbody>
+      </table>
+      <div className="h-4" />
+    </div>
+  );
+}
diff --git a/src/stores/store.ts b/src/stores/store.ts
index 2597a97..276cf7a 100644
--- a/src/stores/store.ts
+++ b/src/stores/store.ts
@@ -172,6 +172,16 @@ export function convert(preview: Array<number>, totalCount: number) {
   };
 }
 
+export async function showTables(path?: string) {
+  const { preview, total_count }: ValidationResponse = await invoke(
+    "show_tables",
+    {
+      path,
+    }
+  );
+  return convert(preview, total_count);
+}
+
 export async function query(
   sql: string,
   limit: number,