diff --git a/backend/apps/datasource/api/datasource.py b/backend/apps/datasource/api/datasource.py index fbc50924..d50b4061 100644 --- a/backend/apps/datasource/api/datasource.py +++ b/backend/apps/datasource/api/datasource.py @@ -1,13 +1,13 @@ 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 @@ -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") @@ -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} diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index 3b5f0f9d..6438aa8c 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -9,7 +9,7 @@ from apps.datasource.utils.utils import aes_decrypt 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 @@ -17,7 +17,7 @@ from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, Cor 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() diff --git a/frontend/src/views/ds/form.vue b/frontend/src/views/ds/form.vue index 8b35c6a6..05f33039 100644 --- a/frontend/src/views/ds/form.vue +++ b/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 diff --git a/frontend/src/views/ds/index.vue b/frontend/src/views/ds/index.vue index c47ce8a3..b73c7f8e 100644 --- a/frontend/src/views/ds/index.vue +++ b/frontend/src/views/ds/index.vue @@ -1,5 +1,5 @@