feat(datasource): Support Excel

Author: XiaJunjie2020

backend/apps/datasource/api/datasource.py

@@ -1,14 +1,18 @@
+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 common.core.deps import SessionDep
-from ..models.datasource import CoreDatasource, CreateDatasource, EditObj, CoreTable
-from ..crud.table import get_tables_by_ds_id
 from ..crud.field import get_fields_by_table_id
-from typing import List
-import uuid
-import pandas as pd
-import hashlib
+from ..crud.table import get_tables_by_ds_id
+from ..models.datasource import CoreDatasource, CreateDatasource, EditObj, CoreTable
 
 router = APIRouter(tags=["datasource"], prefix="/datasource")
 
@@ -92,18 +96,24 @@ async def upload_excel(session: SessionDep, file: UploadFile = File(...)):
     contents = await file.read()
     df_sheets = pd.read_excel(contents, sheet_name=None)
     # build columns and data to insert db
-    # todo
     sheets = []
+    conn = get_data_engine()
     for sheet_name, df in df_sheets.items():
-        print("--------------------")
-        print(f"Sheet: {sheet_name}, {hashlib.sha256(uuid.uuid4().bytes).hexdigest()[:10]}")
-        sheets.append({"name": f"{sheet_name}",
-                       "uniqueName": f"{sheet_name}_{hashlib.sha256(uuid.uuid4().bytes).hexdigest()[:10]}"})
+        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):
-            print(f"{df.columns[i]} , {df.dtypes[i]}")
-
-        for row in df.values:
-            print(row)
-
+            # 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}

backend/apps/datasource/crud/datasource.py

@@ -2,11 +2,13 @@
 import json
 from typing import List
 
-from sqlalchemy import text, and_
+from sqlalchemy import and_, text
 from sqlmodel import select
 
 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.utils.utils import deepcopy_ignore_extra
 from ..crud.field import delete_field_by_ds_id, update_field
@@ -36,8 +38,8 @@ def create_ds(session: SessionDep, create_ds: CreateDatasource):
     ds = CoreDatasource()
     deepcopy_ignore_extra(create_ds, ds)
     ds.create_time = datetime.datetime.now()
-    status = check_status(session, ds)
-    ds.status = "Success" if status is True else "Fail"
+    # status = check_status(session, ds)
+    ds.status = "Success"
     record = CoreDatasource(**ds.model_dump())
     session.add(record)
     session.flush()
@@ -46,9 +48,7 @@ def create_ds(session: SessionDep, create_ds: CreateDatasource):
     session.commit()
 
     # save tables and fields
-    if status:
-        sync_table(session, ds, create_ds.tables)
-
+    sync_table(session, ds, create_ds.tables)
     return ds
 
 
@@ -72,6 +72,16 @@ def update_ds(session: SessionDep, ds: CoreDatasource):
 
 def delete_ds(session: SessionDep, id: int):
     term = session.exec(select(CoreDatasource).where(CoreDatasource.id == id)).first()
+    if term.type == "excel":
+        # drop all tables for current datasource
+        engine = get_engine_conn()
+        conf = DatasourceConf(**json.loads(aes_decrypt(term.configuration)))
+        with engine.connect() as conn:
+            for sheet in conf.sheets:
+                conn.execute(text(f'DROP TABLE IF EXISTS "{sheet["tableName"]}"'))
+            conn.commit()
+            conn.close()
+
     session.delete(term)
     session.commit()
     delete_table_by_ds_id(session, id)
@@ -98,6 +108,11 @@ def getFields(session: SessionDep, id: int, table_name: str):
     return fields
 
 
+def getFieldsByDs(session: SessionDep, ds: CoreDatasource, table_name: str):
+    fields = get_fields(ds, table_name)
+    return fields
+
+
 def execSql(session: SessionDep, id: int, sql: str):
     ds = session.exec(select(CoreDatasource).where(CoreDatasource.id == id)).first()
     return exec_sql(ds, sql)
@@ -128,7 +143,7 @@ def sync_table(session: SessionDep, ds: CoreDatasource, tables: List[CoreTable])
             session.commit()
 
         # sync field
-        fields = getFields(session, ds.id, item.table_name)
+        fields = getFieldsByDs(session, ds, item.table_name)
         sync_fields(session, ds, item, fields)
 
     if len(id_list) > 0:
@@ -177,14 +192,14 @@ def update_table_and_fields(session: SessionDep, data: EditObj):
 
 def preview(session: SessionDep, id: int, data: EditObj):
     ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
-    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     sql: str = ""
     if ds.type == "mysql":
         sql = f"""SELECT `{"`, `".join([f.field_name for f in data.fields if f.checked])}` FROM `{data.table.table_name}` LIMIT 100"""
     elif ds.type == "sqlServer":
         sql = f"""SELECT [{"], [".join([f.field_name for f in data.fields if f.checked])}] FROM [{conf.dbSchema}].[{data.table.table_name}]
             ORDER BY [{data.fields[0].field_name}]
             OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"""
-    elif ds.type == "pg":
+    elif ds.type == "pg" or ds.type == "excel":
         sql = f"""SELECT "{'", "'.join([f.field_name for f in data.fields if f.checked])}" FROM "{conf.dbSchema}"."{data.table.table_name}" LIMIT 100"""
     return exec_sql(ds, sql)

backend/apps/datasource/models/datasource.py

@@ -1,9 +1,10 @@
-from sqlmodel import SQLModel, Field
-from sqlalchemy import Column, Text, BigInteger, DateTime, Integer, Identity
 from datetime import datetime
-from pydantic import BaseModel
 from typing import List, Optional
 
+from pydantic import BaseModel
+from sqlalchemy import Column, Text, BigInteger, DateTime, Integer, Identity
+from sqlmodel import SQLModel, Field
+
 
 class CoreDatasource(SQLModel, table=True):
     __tablename__ = "core_datasource"
@@ -68,6 +69,22 @@ class DatasourceConf(BaseModel):
     driver: str = ''
     extraJdbc: str = ''
     dbSchema: str = ''
+    filename: str = ''
+    sheets: List = ''
+
+    def to_dict(self):
+        return {
+            "host": self.host,
+            "port": self.port,
+            "username": self.username,
+            "password": self.password,
+            "database": self.database,
+            "driver": self.driver,
+            "extraJdbc": self.extraJdbc,
+            "dbSchema": self.dbSchema,
+            "filename": self.filename,
+            "sheets": self.sheets
+        }
 
 
 class TableSchema:

backend/apps/db/db.py

@@ -8,10 +8,11 @@
 
 from apps.datasource.models.datasource import DatasourceConf, CoreDatasource, TableSchema, ColumnSchema
 from apps.datasource.utils.utils import aes_decrypt
+from apps.db.engine import get_engine_config
 
 
 def get_uri(ds: CoreDatasource):
-    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     db_url: str
     if ds.type == "mysql":
         if conf.extraJdbc is not None and conf.extraJdbc != '':
@@ -23,7 +24,7 @@ def get_uri(ds: CoreDatasource):
             db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{urllib.parse.quote(conf.extraJdbc)}"
         else:
             db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
-    elif ds.type == "pg":
+    elif ds.type == "pg" or ds.type == "excel":
         if conf.extraJdbc is not None and conf.extraJdbc != '':
             db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{urllib.parse.quote(conf.extraJdbc)}"
         else:
@@ -34,7 +35,7 @@ def get_uri(ds: CoreDatasource):
 
 
 def get_engine(ds: CoreDatasource) -> Engine:
-    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     if ds.type == "pg" and (conf.dbSchema is not None and conf.dbSchema != ""):
         engine = create_engine(get_uri(ds), connect_args={"options": f"-c search_path={conf.dbSchema}"},
                                pool_timeout=30, pool_size=20, max_overflow=10)
@@ -51,7 +52,7 @@ def get_session(ds: CoreDatasource):
 
 
 def get_tables(ds: CoreDatasource):
-    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     session = get_session(ds)
     result: Result[Any]
     sql: str = ''
@@ -82,7 +83,7 @@ def get_tables(ds: CoreDatasource):
                         t.TABLE_TYPE IN ('BASE TABLE', 'VIEW')
                         AND t.TABLE_SCHEMA = '{conf.dbSchema}';
                     """
-        elif ds.type == "pg":
+        elif ds.type == "pg" or ds.type == "excel":
             sql = """
                     SELECT 
                         c.relname AS TABLE_NAME,
@@ -113,7 +114,7 @@ def get_tables(ds: CoreDatasource):
 
 
 def get_fields(ds: CoreDatasource, table_name: str = None):
-    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     session = get_session(ds)
     result: Result[Any]
     sql: str = ''
@@ -149,7 +150,7 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
                     """
             sql2 = f" AND C.TABLE_NAME = '{table_name}';" if table_name is not None and table_name != "" else ";"
             sql = sql1 + sql2
-        elif ds.type == "pg":
+        elif ds.type == "pg" or ds.type == "excel":
             sql1 = """
                     SELECT 
                         a.attname AS COLUMN_NAME,

backend/apps/db/engine.py

@@ -1,2 +1,68 @@
 # Author: Junjun
 # Date: 2025/5/19
+from typing import List
+
+from sqlalchemy import create_engine, text, MetaData, Table
+from sqlalchemy.orm import sessionmaker
+
+from apps.datasource.models.datasource import DatasourceConf
+
+
+def get_engine_config():
+    return DatasourceConf(username="root", password="123456", host="127.0.0.1", port=5432, database="sqlbot",
+                          dbSchema="public")
+
+
+def get_engine_uri(conf: DatasourceConf):
+    return f"postgresql+psycopg2://{conf.username}:{conf.password}@{conf.host}:{conf.port}/{conf.database}"
+
+
+def get_engine_conn():
+    conf = get_engine_config()
+    db_url = get_engine_uri(conf)
+    engine = create_engine(db_url, connect_args={"options": f"-c search_path={conf.dbSchema}"}, pool_timeout=300,
+                           pool_size=20,
+                           max_overflow=10)
+    return engine
+
+
+def get_data_engine():
+    engine = get_engine_conn()
+    session_maker = sessionmaker(bind=engine)
+    session = session_maker()
+    return session
+
+
+def create_table(session, table_name: str, fields: List[any]):
+    # field type relation
+    list = []
+    for f in fields:
+        if "object" in f["type"]:
+            f["relType"] = "text"
+        elif "int" in f["type"]:
+            f["relType"] = "bigint"
+        elif "float" in f["type"]:
+            f["relType"] = "numeric"
+        elif "datetime" in f["type"]:
+            f["relType"] = "timestamp"
+        else:
+            f["relType"] = "text"
+        list.append(f'"{f["name"]}" {f["relType"]}')
+
+    sql = f"""
+            CREATE TABLE "{table_name}" (
+                {", ".join(list)}
+            );
+            """
+    session.execute(text(sql))
+    session.commit()
+
+
+def insert_data(session, table_name: str, fields: List[any], data: List[any]):
+    engine = get_engine_conn()
+    metadata = MetaData()
+    table = Table(table_name, metadata, autoload_with=engine)
+    with engine.connect() as conn:
+        stmt = table.insert().values(data)
+        conn.execute(stmt)
+        conn.commit()