feat: get tables

Author: XiaJunjie2020

backend/apps/datasource/api/datasource.py

@@ -1,26 +1,36 @@
 from fastapi import APIRouter
-from ..crud.datasource import get_datasource_list, check_status, create_ds, update_ds, delete_ds
+from ..crud.datasource import get_datasource_list, check_status, create_ds, update_ds, delete_ds, getTables
 from common.core.deps import SessionDep
 from ..models.datasource import CoreDatasource
 
 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: CoreDatasource):
     return create_ds(session, ds)
 
+
 @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)
+    return delete_ds(session, id)
+
+
+@router.post("/getTables/{id}")
+async def get_tables(session: SessionDep, id: int):
+    return getTables(session, id)

backend/apps/datasource/crud/datasource.py

@@ -4,20 +4,29 @@
 from common.core.deps import SessionDep
 import json
 from ..utils.utils import aes_decrypt
-from apps.db.db import get_connection
+from apps.db.db import get_session, get_tables
+from sqlalchemy import text
+
 
 def get_datasource_list(session: SessionDep) -> CoreDatasource:
     statement = select(CoreDatasource).order_by(CoreDatasource.create_time.desc())
     datasource_list = session.exec(statement).fetchall()
     return datasource_list
 
+
 def check_status(session: SessionDep, ds: CoreDatasource):
     conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
-    conn = get_connection(conf)
-    if conn is not None:
+    conn = get_session(conf, ds)
+    try:
+        conn.execute(text("SELECT 1")).scalar()
+        print("success")
         return True
-    else:
-        return False
+    except Exception as e:
+        print("Fail:", e)
+        raise e
+    finally:
+        conn.close()
+
 
 def create_ds(session: SessionDep, ds: CoreDatasource):
     ds.create_time = datetime.datetime.now()
@@ -27,6 +36,7 @@ def create_ds(session: SessionDep, ds: CoreDatasource):
     session.commit()
     return ds
 
+
 def update_ds(session: SessionDep, ds: CoreDatasource):
     ds.id = int(ds.id)
     record = session.exec(select(CoreDatasource).where(CoreDatasource.id == ds.id)).first()
@@ -37,10 +47,18 @@ def update_ds(session: SessionDep, ds: CoreDatasource):
     session.commit()
     return ds
 
+
 def delete_ds(session: SessionDep, id: int):
     term = session.exec(select(CoreDatasource).where(CoreDatasource.id == id)).first()
     session.delete(term)
     session.commit()
     return {
         "message": f"Datasource with ID {id} deleted successfully."
-    }
+    }
+
+
+def getTables(session: SessionDep, id: int):
+    ds = session.exec(select(CoreDatasource).where(CoreDatasource.id == id)).first()
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
+    tables = get_tables(conf, ds)
+    return tables

backend/apps/datasource/models/datasource.py

@@ -3,6 +3,7 @@
 from datetime import datetime
 from pydantic import BaseModel
 
+
 class CoreDatasource(SQLModel, table=True):
     __tablename__ = "core_datasource"
     id: int = Field(sa_column=Column(Integer, Identity(always=True), nullable=False, primary_key=True))
@@ -14,10 +15,20 @@ class CoreDatasource(SQLModel, table=True):
     create_by: int = Field(sa_column=Column(BigInteger()))
     status: str = Field(max_length=64, nullable=True)
 
+
 class DatasourceConf(BaseModel):
     host: str = ''
     port: int = 0
     username: str = ''
     password: str = ''
     database: str = ''
-    driver: str = ''
+    driver: str = ''
+
+
+class TableSchema:
+    def __init__(self, attr1, attr2):
+        self.tableName = attr1
+        self.tableRemark = attr2
+
+    tableName: str
+    tableRemark: str

backend/apps/db/db.py

@@ -1,23 +1,41 @@
-import pyodbc
-from apps.datasource.models.datasource import DatasourceConf
+from sqlalchemy import create_engine, text, Result
+from sqlalchemy.orm import sessionmaker
+from apps.datasource.models.datasource import DatasourceConf, CoreDatasource, TableSchema
+import urllib.parse
+from typing import Any
 
-def get_connection(conf: DatasourceConf):
-    conn_str = (
-        "DRIVER={MySQL ODBC 9.3 Unicode Driver};"  # todo driver config
-        f"SERVER={conf.host};"
-        f"DATABASE={conf.database};"
-        f"UID={conf.username};"
-        f"PWD={conf.password};"
-        f"PORT={conf.port};"
-    )
-    conn = None
-    try:
-        conn = pyodbc.connect(conn_str)
-        print("Connect Success")
-        return conn
-    except pyodbc.Error as e:
-        print(f"Connect Fail：{e}")
-        raise e
-    finally:
-        if conn is not None:
-            conn.close()
+
+def get_session(conf: DatasourceConf, ds: CoreDatasource):
+    db_url: str
+    if ds.type == "mysql":
+        db_url = f"mysql+pymysql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
+    else:
+        raise 'The datasource type not support.'
+    engine = create_engine(db_url)
+    session_maker = sessionmaker(bind=engine)
+    session = session_maker()
+    return session
+
+
+def get_tables(conf: DatasourceConf, ds: CoreDatasource):
+    session = get_session(conf, ds)
+    result: Result[Any]
+    if ds.type == "mysql":
+        sql = f"""SELECT 
+                        TABLE_NAME AS `Table Name`, 
+                        TABLE_COMMENT AS `Table Comment`
+                    FROM 
+                        information_schema.TABLES
+                    WHERE 
+                        TABLE_SCHEMA = '{conf.database}';"""
+        result = session.execute(text(sql))
+
+    res = result.fetchall()
+    res_list = [TableSchema(*item) for item in res]
+    return res_list
+
+
+def get_fields(conf: DatasourceConf, ds: CoreDatasource):
+    session = get_session(conf, ds)
+    result: Result[Any]
+    pass

backend/pyproject.toml

@@ -18,8 +18,8 @@ dependencies = [
     "pydantic-settings<3.0.0,>=2.2.1",
     "sentry-sdk[fastapi]<2.0.0,>=1.40.6",
     "pyjwt<3.0.0,>=2.8.0",
-    "pyodbc (>=5.2.0,<5.2.1)",
     "pycryptodome (>=3.22.0,<4.0.0)",
+    "pymysql (>=1.1.1,<2.0.0)",
 ]
 [[tool.uv.index]]
 url = "https://pypi.tuna.tsinghua.edu.cn/simple"

frontend/src/api/datasource.ts

@@ -5,5 +5,6 @@ export const datasourceApi = {
   add: (data: any) => request.post('/datasource/add', data),
   list: () => request.get('/datasource/list'),
   update: (data: any) => request.post('/datasource/update', data),
-  delete: (id: Number) => request.post(`/datasource/delete/${id}`)
+  delete: (id: Number) => request.post(`/datasource/delete/${id}`),
+  getTables: (id: Number) => request.post(`/datasource/getTables/${id}`),
 }

frontend/src/views/ds/index.vue

@@ -35,6 +35,7 @@
         </div>
         <div class="connection-status" :class="`${getStatus(ds.status)}`">{{ ds.status }}</div>
         <div class="connection-actions">
+          <el-button class="action-btn" circle @click="getTables(ds.id)"/>
           <el-button type="primary" class="action-btn" circle @click="editDs(ds)" :icon="IconOpeEdit"/>
           <el-button type="danger" class="action-btn" circle @click="deleteDs(ds)" :icon="IconOpeDelete"/>
         </div>
@@ -114,6 +115,12 @@ const deleteDs = (item: any) => {
   })
 }
 
+const getTables = (id: Number) => {
+  datasourceApi.getTables(id).then((res) => {
+    console.log(res)
+  })
+}
+
 onMounted(() => {
   list()
 })