From 2790ef9b42d29453d10985ac01d260679b353d55 Mon Sep 17 00:00:00 2001
From: junjun <junjie.xia@fit2cloud.com>
Date: Tue, 19 Aug 2025 15:45:38 +0800
Subject: [PATCH] refactor: add schema api

---
 backend/apps/datasource/api/datasource.py | 16 +++++++++
 backend/apps/db/db.py                     | 41 ++++++++++++++++++++---
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/backend/apps/datasource/api/datasource.py b/backend/apps/datasource/api/datasource.py
index 0eeba6de..83d26ff0 100644
--- a/backend/apps/datasource/api/datasource.py
+++ b/backend/apps/datasource/api/datasource.py
@@ -20,6 +20,7 @@ from ..crud.datasource import get_datasource_list, check_status, create_ds, upda
 from ..crud.field import get_fields_by_table_id
 from ..crud.table import get_tables_by_ds_id
 from ..models.datasource import CoreDatasource, CreateDatasource, TableObj, CoreTable, CoreField
+from apps.db.db import get_schema
 
 router = APIRouter(tags=["datasource"], prefix="/datasource")
 path = settings.EXCEL_PATH
@@ -107,6 +108,21 @@ async def get_tables_by_conf(session: SessionDep, trans: Trans, ds: CoreDatasour
             raise HTTPException(status_code=500, detail=f'Get table Failed: {e.args}')
 
 
+@router.post("/getSchemaByConf")
+async def get_schema_by_conf(session: SessionDep, trans: Trans, ds: CoreDatasource):
+    try:
+        return get_schema(ds)
+    except Exception as e:
+        # check ds status
+        def inner():
+            return check_status(session, trans, ds, True)
+
+        status = await asyncio.to_thread(inner)
+        if status:
+            SQLBotLogUtil.error(f"get table failed: {e}")
+            raise HTTPException(status_code=500, detail=f'Get table Failed: {e.args}')
+
+
 @router.post("/getFields/{id}/{table_name}")
 async def get_fields(session: SessionDep, id: int, table_name: str):
     return getFields(session, id, table_name)
diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py
index 3f6efe50..383f55f8 100644
--- a/backend/apps/db/db.py
+++ b/backend/apps/db/db.py
@@ -64,11 +64,16 @@ def get_engine(ds: CoreDatasource, timeout: int = 0) -> Engine:
         conf.timeout = timeout
     if timeout > 0:
         conf.timeout = timeout
-    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={urllib.parse.quote(conf.dbSchema)}",
-                                             "connect_timeout": conf.timeout},
-                               pool_timeout=conf.timeout)
+    if ds.type == "pg":
+        if conf.dbSchema is not None and conf.dbSchema != "":
+            engine = create_engine(get_uri(ds),
+                                   connect_args={"options": f"-c search_path={urllib.parse.quote(conf.dbSchema)}",
+                                                 "connect_timeout": conf.timeout},
+                                   pool_timeout=conf.timeout)
+        else:
+            engine = create_engine(get_uri(ds),
+                                   connect_args={"connect_timeout": conf.timeout},
+                                   pool_timeout=conf.timeout)
     elif ds.type == 'sqlServer':
         engine = create_engine(get_uri(ds), pool_timeout=conf.timeout)
     elif ds.type == 'oracle':
@@ -86,6 +91,32 @@ def get_session(ds: CoreDatasource | AssistantOutDsSchema):
     return session
 
 
+def get_schema(ds: CoreDatasource):
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
+    db = DB.get_db(ds.type)
+    if db.connect_type == ConnectType.sqlalchemy:
+        with get_session(ds) as session:
+            sql: str = ''
+            if ds.type == "sqlServer":
+                sql = f"""select name from sys.schemas"""
+            elif ds.type == "pg" or ds.type == "excel":
+                sql = """SELECT nspname FROM pg_namespace"""
+            elif ds.type == "oracle":
+                sql = f"""select * from all_users"""
+            with session.execute(text(sql)) as result:
+                res = result.fetchall()
+                res_list = [item[0] for item in res]
+                return res_list
+    else:
+        if ds.type == 'dm':
+            with dmPython.connect(user=conf.username, password=conf.password, server=conf.host,
+                                  port=conf.port) as conn, conn.cursor() as cursor:
+                cursor.execute(f"""select OBJECT_NAME from dba_objects where object_type='SCH'""")
+                res = cursor.fetchall()
+                res_list = [item[0] for item in res]
+                return res_list
+
+
 def get_tables(ds: CoreDatasource):
     conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     db = DB.get_db(ds.type)