mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
120 lines
3.8 KiB
Python
120 lines
3.8 KiB
Python
import hashlib
|
|
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 ..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")
|
|
|
|
|
|
@router.get("/list")
|
|
async def datasource_list(session: SessionDep):
|
|
return get_datasource_list(session=session)
|
|
|
|
|
|
@router.post("/check")
|
|
async def check(session: SessionDep, ds: CoreDatasource):
|
|
return check_status(session, ds)
|
|
|
|
|
|
@router.post("/add", response_model=CoreDatasource)
|
|
async def add(session: SessionDep, ds: CreateDatasource):
|
|
return create_ds(session, ds)
|
|
|
|
|
|
@router.post("/chooseTables/{id}")
|
|
async def choose_tables(session: SessionDep, id: int, tables: List[CoreTable]):
|
|
chooseTables(session, id, tables)
|
|
|
|
|
|
@router.post("/update", response_model=CoreDatasource)
|
|
async def update(session: SessionDep, ds: CoreDatasource):
|
|
return update_ds(session, ds)
|
|
|
|
|
|
@router.post("/delete/{id}", response_model=CoreDatasource)
|
|
async def delete(session: SessionDep, id: int):
|
|
return delete_ds(session, id)
|
|
|
|
|
|
@router.post("/getTables/{id}")
|
|
async def get_tables(session: SessionDep, id: int):
|
|
return getTables(session, id)
|
|
|
|
|
|
@router.post("/getTablesByConf")
|
|
async def get_tables_by_conf(session: SessionDep, ds: CoreDatasource):
|
|
return getTablesByDs(session, ds)
|
|
|
|
|
|
@router.post("/getFields/{id}/{table_name}")
|
|
async def get_tables(session: SessionDep, id: int, table_name: str):
|
|
return getFields(session, id, table_name)
|
|
|
|
|
|
@router.post("/execSql/{id}/{sql}")
|
|
async def exec_sql(session: SessionDep, id: int, sql: str):
|
|
return execSql(session, id, sql)
|
|
|
|
|
|
@router.post("/tableList/{id}")
|
|
async def table_list(session: SessionDep, id: int):
|
|
return get_tables_by_ds_id(session, id)
|
|
|
|
|
|
@router.post("/fieldList/{id}")
|
|
async def field_list(session: SessionDep, id: int):
|
|
return get_fields_by_table_id(session, id)
|
|
|
|
|
|
@router.post("/editLocalComment")
|
|
async def edit_local(session: SessionDep, data: EditObj):
|
|
update_table_and_fields(session, data)
|
|
|
|
|
|
@router.post("/previewData/{id}")
|
|
async def edit_local(session: SessionDep, id: int, data: EditObj):
|
|
return preview(session, id, data)
|
|
|
|
|
|
@router.post("/uploadExcel")
|
|
async def upload_excel(session: SessionDep, file: UploadFile = File(...)):
|
|
ALLOWED_EXTENSIONS = {"xlsx", "xls", "csv"}
|
|
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 = []
|
|
conn = get_data_engine()
|
|
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}
|