Merge remote-tracking branch 'origin/main'

Author: ziyujiahao

backend/apps/datasource/api/datasource.py

@@ -1,20 +1,21 @@
 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
+from common.core.deps import SessionDep, CurrentUser
 from ..crud.datasource import get_datasource_list, check_status, create_ds, update_ds, delete_ds, getTables, getFields, \
     execSql, update_table_and_fields, getTablesByDs, chooseTables, preview
 from ..crud.field import get_fields_by_table_id
 from ..crud.table import get_tables_by_ds_id
 from ..models.datasource import CoreDatasource, CreateDatasource, EditObj, CoreTable
 
 router = APIRouter(tags=["datasource"], prefix="/datasource")
+path = "/opt/sqlbot/data/excel"
 
 
 @router.get("/list")
@@ -28,8 +29,8 @@ async def check(session: SessionDep, ds: CoreDatasource):
 
 
 @router.post("/add", response_model=CoreDatasource)
-async def add(session: SessionDep, ds: CreateDatasource):
-    return create_ds(session, ds)
+async def add(session: SessionDep, user: CurrentUser, ds: CreateDatasource):
+    return create_ds(session, user, ds)
 
 
 @router.post("/chooseTables/{id}")
@@ -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

@@ -9,15 +9,15 @@
 from apps.db.db import get_engine, get_tables, get_fields, exec_sql
 from apps.db.engine import get_engine_config
 from apps.db.engine import get_engine_conn
-from common.core.deps import SessionDep
+from common.core.deps import SessionDep, CurrentUser
 from common.utils.utils import deepcopy_ignore_extra
 from ..crud.field import delete_field_by_ds_id, update_field
 from ..crud.table import delete_table_by_ds_id, update_table
 from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema, EditObj, \
     DatasourceConf
 
 
-def get_datasource_list(session: SessionDep) -> CoreDatasource:
+def get_datasource_list(session: SessionDep):
     statement = select(CoreDatasource).order_by(CoreDatasource.create_time.desc())
     datasource_list = session.exec(statement).fetchall()
     return datasource_list
@@ -34,11 +34,12 @@ def check_status(session: SessionDep, ds: CoreDatasource):
         return False
 
 
-def create_ds(session: SessionDep, create_ds: CreateDatasource):
+def create_ds(session: SessionDep, user: CurrentUser, create_ds: CreateDatasource):
     ds = CoreDatasource()
     deepcopy_ignore_extra(create_ds, ds)
     ds.create_time = datetime.datetime.now()
     # status = check_status(session, ds)
+    ds.create_by = user.id
     ds.status = "Success"
     record = CoreDatasource(**ds.model_dump())
     session.add(record)
@@ -80,7 +81,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)})
           })
         }
       })
@@ -280,20 +281,26 @@ const check = () => {
         type: 'success',
         showClose: true
       })
+    } else {
+      ElMessage({
+        message: 'Connect failed',
+        type: 'error',
+        showClose: true
+      })
     }
   })
 }
 
 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) => {
@@ -303,11 +310,17 @@ const next = async(formEl: FormInstance | undefined) => {
             datasourceApi.getTablesByConf(form.value).then((res) => {
               tableList.value = res
             })
+          } else {
+            ElMessage({
+              message: 'Connect failed',
+              type: 'error',
+              showClose: true
+            })
           }
         })
       }
-    })
-  }
+    }
+  })
 }
 
 const preview = () => {
@@ -323,7 +336,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

frontend/src/views/ds/index.vue

@@ -1,5 +1,5 @@
 <template>
-  <div>
+  <div v-loading="loading">
     <div class="header">
       <div class="mt-4">
         <el-input
@@ -72,6 +72,7 @@ const dsForm = ref()
 const dsList = ref<any>([])// show ds list
 const allDsList = ref<any>([])// all ds list
 const router = useRouter()
+const loading = ref(false)
 
 const getStatus = (status: string) => {
   if (status === 'Success') {
@@ -98,9 +99,11 @@ const refresh = () => {
 }
 
 const list = () => {
+  loading.value = true
   datasourceApi.list().then((res) => {
     allDsList.value = res
     dsList.value = JSON.parse(JSON.stringify(allDsList.value))
+    loading.value = false
   })
 }