mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
feat(datasource): Support csv
This commit is contained in:
@@ -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 ..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")
|
||||
@@ -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}
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user