refactor: code optimization

XiaJunjie2020 · XiaJunjie2020 · commit db2fdfc4c241 · 2025-08-20T14:11:37.000+08:00
diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py
@@ -1,23 +1,19 @@
 import datetime
 import json
-import platform
 from typing import List, Optional
 
-if platform.system() != "Darwin":
-    import dmPython
 from fastapi import HTTPException
 from sqlalchemy import and_, text
 from sqlmodel import select
 
 from apps.datasource.crud.permission import get_column_permission_fields, get_row_permission_filters, is_normal_user
 from apps.datasource.utils.utils import aes_decrypt
-from apps.db.constant import ConnectType
 from apps.db.constant import DB
-from apps.db.db import get_engine, get_tables, get_fields, exec_sql
+from apps.db.db import get_tables, get_fields, exec_sql, check_connection
 from apps.db.engine import get_engine_config, get_engine_conn
 from apps.db.type import db_type_relation
 from common.core.deps import SessionDep, CurrentUser, Trans
-from common.utils.utils import SQLBotLogUtil, deepcopy_ignore_extra
+from common.utils.utils import deepcopy_ignore_extra
 from .table import get_tables_by_ds_id
 from ..crud.field import delete_field_by_ds_id, update_field
 from ..crud.table import delete_table_by_ds_id, update_table
@@ -49,32 +45,7 @@ def check_status_by_id(session: SessionDep, trans: Trans, ds_id: int, is_raise:
 
 
 def check_status(session: SessionDep, trans: Trans, ds: CoreDatasource, is_raise: bool = False):
-    db = DB.get_db(ds.type)
-    if db.connect_type == ConnectType.sqlalchemy:
-        conn = get_engine(ds, 10)
-        try:
-            with conn.connect() as connection:
-                SQLBotLogUtil.info("success")
-                return True
-        except Exception as e:
-            SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
-            if is_raise:
-                raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
-            return False
-    else:
-        conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
-        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:
-                try:
-                    cursor.execute('select 1', timeout=10).fetchall()
-                    SQLBotLogUtil.info("success")
-                    return True
-                except Exception as e:
-                    SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
-                    if is_raise:
-                        raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
-                    return False
+    return check_connection(trans, ds, is_raise)
 
 
 def check_name(session: SessionDep, trans: Trans, user: CurrentUser, ds: CoreDatasource):
diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py
@@ -4,6 +4,8 @@
 import urllib.parse
 from decimal import Decimal
 
+from apps.db.db_sql import get_table_sql, get_field_sql
+
 if platform.system() != "Darwin":
     import dmPython
 from sqlalchemy import create_engine, text, Engine
@@ -15,6 +17,9 @@
 from apps.db.engine import get_engine_config
 from apps.system.crud.assistant import get_ds_engine
 from apps.system.schemas.system_schema import AssistantOutDsSchema
+from common.core.deps import Trans
+from common.utils.utils import SQLBotLogUtil
+from fastapi import HTTPException
 
 
 def get_uri(ds: CoreDatasource) -> str:
@@ -93,6 +98,35 @@ def get_session(ds: CoreDatasource | AssistantOutDsSchema):
     return session
 
 
+def check_connection(trans: Trans, ds: CoreDatasource, is_raise: bool = False):
+    db = DB.get_db(ds.type)
+    if db.connect_type == ConnectType.sqlalchemy:
+        conn = get_engine(ds, 10)
+        try:
+            with conn.connect() as connection:
+                SQLBotLogUtil.info("success")
+                return True
+        except Exception as e:
+            SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
+            if is_raise:
+                raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
+            return False
+    else:
+        conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
+        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:
+                try:
+                    cursor.execute('select 1', timeout=10).fetchall()
+                    SQLBotLogUtil.info("success")
+                    return True
+                except Exception as e:
+                    SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
+                    if is_raise:
+                        raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
+                    return False
+
+
 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)
@@ -122,78 +156,9 @@ def get_schema(ds: CoreDatasource):
 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)
+    sql = get_table_sql(ds, conf)
     if db.connect_type == ConnectType.sqlalchemy:
         with get_session(ds) as session:
-            sql: str = ''
-            if ds.type == "mysql":
-                sql = f"""
-                        SELECT 
-                            TABLE_NAME, 
-                            TABLE_COMMENT
-                        FROM 
-                            information_schema.TABLES
-                        WHERE 
-                            TABLE_SCHEMA = '{conf.database}'
-                        """
-            elif ds.type == "sqlServer":
-                sql = f"""
-                        SELECT 
-                            TABLE_NAME AS [TABLE_NAME],
-                            ISNULL(ep.value, '') AS [TABLE_COMMENT]
-                        FROM 
-                            INFORMATION_SCHEMA.TABLES t
-                        LEFT JOIN 
-                            sys.extended_properties ep 
-                            ON ep.major_id = OBJECT_ID(t.TABLE_SCHEMA + '.' + t.TABLE_NAME)
-                            AND ep.minor_id = 0 
-                            AND ep.name = 'MS_Description' 
-                        WHERE 
-                            t.TABLE_TYPE IN ('BASE TABLE', 'VIEW')
-                            AND t.TABLE_SCHEMA = '{conf.dbSchema}'
-                        """
-            elif ds.type == "pg" or ds.type == "excel":
-                sql = """
-                      SELECT c.relname                                       AS TABLE_NAME,
-                             COALESCE(d.description, obj_description(c.oid)) AS TABLE_COMMENT
-                      FROM pg_class c
-                               LEFT JOIN
-                           pg_namespace n ON n.oid = c.relnamespace
-                               LEFT JOIN
-                           pg_description d ON d.objoid = c.oid AND d.objsubid = 0
-                      WHERE n.nspname = current_schema()
-                        AND c.relkind IN ('r', 'v', 'p', 'm')
-                        AND c.relname NOT LIKE 'pg_%'
-                        AND c.relname NOT LIKE 'sql_%'
-                      ORDER BY c.relname \
-                      """
-            elif ds.type == "oracle":
-                sql = f"""
-                        SELECT 
-                            t.TABLE_NAME AS "TABLE_NAME",
-                            NVL(c.COMMENTS, '') AS "TABLE_COMMENT"
-                        FROM (
-                            SELECT TABLE_NAME, 'TABLE' AS OBJECT_TYPE
-                            FROM DBA_TABLES
-                            WHERE OWNER = '{conf.dbSchema}'  
-                            UNION ALL
-                            SELECT VIEW_NAME AS TABLE_NAME, 'VIEW' AS OBJECT_TYPE
-                            FROM DBA_VIEWS
-                            WHERE OWNER = '{conf.dbSchema}'  
-                        ) t
-                        LEFT JOIN DBA_TAB_COMMENTS c 
-                            ON t.TABLE_NAME = c.TABLE_NAME 
-                            AND c.TABLE_TYPE = t.OBJECT_TYPE
-                            AND c.OWNER = '{conf.dbSchema}'   
-                        ORDER BY t.TABLE_NAME
-                        """
-            elif ds.type == "ck":
-                sql = f"""
-                        SELECT name, comment
-                        FROM system.tables
-                        WHERE database = '{conf.database}'
-                          AND engine NOT IN ('Dictionary')
-                        ORDER BY name
-                        """
             with session.execute(text(sql)) as result:
                 res = result.fetchall()
                 res_list = [TableSchema(*item) for item in res]
@@ -202,11 +167,7 @@ def get_tables(ds: CoreDatasource):
         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 table_name, comments 
-                                    from all_tab_comments 
-                                    where owner='{conf.dbSchema}'
-                                    AND (table_type = 'TABLE' or table_type = 'VIEW')
-                                """, timeout=conf.timeout)
+                cursor.execute(sql, timeout=conf.timeout)
                 res = cursor.fetchall()
                 res_list = [TableSchema(*item) for item in res]
                 return res_list
@@ -215,93 +176,9 @@ def get_tables(ds: CoreDatasource):
 def get_fields(ds: CoreDatasource, table_name: str = None):
     conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     db = DB.get_db(ds.type)
+    sql = get_field_sql(ds, conf, table_name)
     if db.connect_type == ConnectType.sqlalchemy:
         with get_session(ds) as session:
-            sql: str = ''
-            if ds.type == "mysql":
-                sql1 = f"""
-                        SELECT 
-                            COLUMN_NAME,
-                            DATA_TYPE,
-                            COLUMN_COMMENT
-                        FROM 
-                            INFORMATION_SCHEMA.COLUMNS
-                        WHERE 
-                            TABLE_SCHEMA = '{conf.database}'
-                        """
-                sql2 = f" AND TABLE_NAME = '{table_name}'" if table_name is not None and table_name != "" else ""
-                sql = sql1 + sql2
-            elif ds.type == "sqlServer":
-                sql1 = f"""
-                        SELECT 
-                            COLUMN_NAME AS [COLUMN_NAME],
-                            DATA_TYPE AS [DATA_TYPE],
-                            ISNULL(EP.value, '') AS [COLUMN_COMMENT]
-                        FROM 
-                            INFORMATION_SCHEMA.COLUMNS C
-                        LEFT JOIN 
-                            sys.extended_properties EP 
-                            ON EP.major_id = OBJECT_ID(C.TABLE_SCHEMA + '.' + C.TABLE_NAME)
-                            AND EP.minor_id = C.ORDINAL_POSITION
-                            AND EP.name = 'MS_Description'
-                        WHERE 
-                            C.TABLE_SCHEMA = '{conf.dbSchema}'
-                        """
-                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" or ds.type == "excel":
-                sql1 = """
-                       SELECT a.attname                                       AS COLUMN_NAME,
-                              pg_catalog.format_type(a.atttypid, a.atttypmod) AS DATA_TYPE,
-                              col_description(c.oid, a.attnum)                AS COLUMN_COMMENT
-                       FROM pg_catalog.pg_attribute a
-                                JOIN
-                            pg_catalog.pg_class c ON a.attrelid = c.oid
-                                JOIN
-                            pg_catalog.pg_namespace n ON n.oid = c.relnamespace
-                       WHERE n.nspname = current_schema()
-                         AND a.attnum > 0
-                         AND NOT a.attisdropped \
-                       """
-                sql2 = f" AND c.relname = '{table_name}'" if table_name is not None and table_name != "" else ""
-                sql = sql1 + sql2
-            elif ds.type == "oracle":
-                sql1 = f"""
-                        SELECT 
-                            col.COLUMN_NAME AS "COLUMN_NAME",
-                            (CASE 
-                                WHEN col.DATA_TYPE IN ('VARCHAR2', 'CHAR', 'NVARCHAR2', 'NCHAR') 
-                                    THEN col.DATA_TYPE || '(' || col.DATA_LENGTH || ')' 
-                                WHEN col.DATA_TYPE = 'NUMBER' AND col.DATA_PRECISION IS NOT NULL 
-                                    THEN col.DATA_TYPE || '(' || col.DATA_PRECISION || 
-                                         CASE WHEN col.DATA_SCALE > 0 THEN ',' || col.DATA_SCALE END || ')' 
-                                ELSE col.DATA_TYPE 
-                            END) AS "DATA_TYPE",
-                            NVL(com.COMMENTS, '') AS "COLUMN_COMMENT"
-                        FROM 
-                            DBA_TAB_COLUMNS col
-                        LEFT JOIN 
-                            DBA_COL_COMMENTS com 
-                            ON col.OWNER = com.OWNER 
-                            AND col.TABLE_NAME = com.TABLE_NAME 
-                            AND col.COLUMN_NAME = com.COLUMN_NAME
-                        WHERE 
-                            col.OWNER = '{conf.dbSchema}'
-                        """
-                sql2 = f" AND col.TABLE_NAME = '{table_name}'" if table_name is not None and table_name != "" else ""
-                sql = sql1 + sql2
-            elif ds.type == "ck":
-                sql1 = f"""
-                        SELECT 
-                            name AS COLUMN_NAME,
-                            type AS DATA_TYPE,
-                            comment AS COLUMN_COMMENT
-                        FROM system.columns
-                        WHERE database = '{conf.database}'
-                        """
-                sql2 = f" AND table = '{table_name}'" if table_name is not None and table_name != "" else ""
-                sql = sql1 + sql2
-
             with session.execute(text(sql)) as result:
                 res = result.fetchall()
                 res_list = [ColumnSchema(*item) for item in res]
@@ -310,23 +187,7 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
         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:
-                sql1 = f"""
-                                SELECT 
-                                    c.COLUMN_NAME    AS "COLUMN_NAME",
-                                    c.DATA_TYPE      AS "DATA_TYPE",
-                                    COALESCE(com.COMMENTS, '') AS "COMMENTS"
-                                FROM 
-                                    ALL_TAB_COLS c
-                                LEFT JOIN 
-                                    ALL_COL_COMMENTS com 
-                                    ON c.OWNER = com.OWNER 
-                                   AND c.TABLE_NAME = com.TABLE_NAME 
-                                   AND c.COLUMN_NAME = com.COLUMN_NAME
-                                WHERE 
-                                    c.OWNER = '{conf.dbSchema}'
-                                """
-                sql2 = f" AND c.TABLE_NAME = '{table_name}'" if table_name is not None and table_name != "" else ""
-                cursor.execute(sql1 + sql2, timeout=conf.timeout)
+                cursor.execute(sql, timeout=conf.timeout)
                 res = cursor.fetchall()
                 res_list = [ColumnSchema(*item) for item in res]
                 return res_list
diff --git a/backend/apps/db/db_sql.py b/backend/apps/db/db_sql.py
