feat(datasource): Support csv

Author: XiaJunjie2020

backend/apps/datasource/api/datasource.py

@@ -1,10 +1,10 @@
 import hashlib
+import os
 import uuid
 from typing import List
 
 import pandas as pd
 from fastapi import APIRouter, File, UploadFile, HTTPException
-from numba.core.cgutils import if_zero
 
 from apps.db.engine import create_table, get_data_engine, insert_data
 from common.core.deps import SessionDep
@@ -15,6 +15,7 @@
 from ..models.datasource import CoreDatasource, CreateDatasource, EditObj, CoreTable
 
 router = APIRouter(tags=["datasource"], prefix="/datasource")
+path = "/opt/sqlbot/data/excel"
 
 
 @router.get("/list")
@@ -93,13 +94,17 @@ async def upload_excel(session: SessionDep, file: UploadFile = File(...)):
     if not file.filename.lower().endswith(tuple(ALLOWED_EXTENSIONS)):
         raise HTTPException(400, "Only support .xlsx/.xls/.csv")
 
-    contents = await file.read()
-    df_sheets = pd.read_excel(contents, sheet_name=None)
-    # build columns and data to insert db
-    sheets = []
+    os.makedirs(path, exist_ok=True)
+    filename = f"{file.filename.split('.')[0]}_{hashlib.sha256(uuid.uuid4().bytes).hexdigest()[:10]}.{file.filename.split('.')[1]}"
+    save_path = os.path.join(path, filename)
+    with open(save_path, "wb") as f:
+        f.write(await file.read())
+
     conn = get_data_engine()
-    for sheet_name, df in df_sheets.items():
-        tableName = f"{sheet_name}_{hashlib.sha256(uuid.uuid4().bytes).hexdigest()[:10]}"
+    sheets = []
+    if filename.endswith(".csv"):
+        df = pd.read_csv(save_path)
+        tableName = f"sheet1_{hashlib.sha256(uuid.uuid4().bytes).hexdigest()[:10]}"
         sheets.append({"tableName": tableName, "tableComment": ""})
         column_len = len(df.dtypes)
         fields = []
@@ -115,5 +120,27 @@ async def upload_excel(session: SessionDep, file: UploadFile = File(...)):
         ]
         # insert data
         insert_data(conn, tableName, fields, data)
+    else:
+        df_sheets = pd.read_excel(save_path, sheet_name=None)
+        # build columns and data to insert db
+        for sheet_name, df in df_sheets.items():
+            tableName = f"{sheet_name}_{hashlib.sha256(uuid.uuid4().bytes).hexdigest()[:10]}"
+            sheets.append({"tableName": tableName, "tableComment": ""})
+            column_len = len(df.dtypes)
+            fields = []
+            for i in range(column_len):
+                # build fields
+                fields.append({"name": df.columns[i], "type": str(df.dtypes[i]), "relType": ""})
+            # create table
+            create_table(conn, tableName, fields)
+
+            data = [
+                {df.columns[i]: int(row[i]) if "int" in str(df.dtypes[i]) else row[i] for i in range(len(row))}
+                for row in df.values
+            ]
+            # insert data
+            insert_data(conn, tableName, fields, data)
     conn.close()
-    return {"filename": file.filename, "sheets": sheets}
+
+    os.remove(save_path)
+    return {"filename": filename, "sheets": sheets}

backend/apps/datasource/crud/datasource.py

@@ -80,7 +80,6 @@ def delete_ds(session: SessionDep, id: int):
             for sheet in conf.sheets:
                 conn.execute(text(f'DROP TABLE IF EXISTS "{sheet["tableName"]}"'))
             conn.commit()
-            conn.close()
 
     session.delete(term)
     session.commit()

frontend/src/views/ds/form.vue

@@ -189,8 +189,9 @@ const open = (item: any, editTable: boolean = false) => {
         if(item.type === 'excel') {
           tableList.value = config.value.sheets
         } else {
-          datasourceApi.getTablesByConf(form.value).then((res) => {
-            tableList.value = res
+          datasourceApi.getTablesByConf(form.value).then((table) => {
+            tableList.value = table
+            checkList.value = checkList.value.filter((ele:string) => {return table.map((ele:any) => {return ele.tableName}).includes(ele)})
           })
         }
       })
@@ -285,15 +286,15 @@ const check = () => {
 }
 
 const next = async(formEl: FormInstance | undefined) => {
-  if (form.value.type === "excel") {
-    // next, show tables 
-    if (excelUploadSuccess.value) {
-      active.value++
-    }
-  } else {
-    if(!formEl) return
-    await formEl.validate((valid) => {
-      if (valid) {
+  if(!formEl) return
+  await formEl.validate((valid) => {
+    if (valid) {
+      if (form.value.type === "excel") {
+        // next, show tables 
+        if (excelUploadSuccess.value) {
+          active.value++
+        }
+      } else {
         // check status if success do next
         buildConf()
         datasourceApi.check(form.value).then((res: boolean) => {
@@ -306,8 +307,8 @@ const next = async(formEl: FormInstance | undefined) => {
           }
         })
       }
-    })
-  }
+    }
+  })
 }
 
 const preview = () => {
@@ -323,7 +324,6 @@ const beforeUpload = (rawFile: any) => {
 }
 
 const onSuccess = (response: any) => {
-  console.log(response)
   config.value.filename = response.data.filename
   config.value.sheets = response.data.sheets
   tableList.value = response.data.sheets