refactor: code optimization

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):

backend/apps/db/db.py

@@ -4,6 +4,8 @@ import platform
 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.constant import DB, ConnectType
 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

backend/apps/db/db_sql.py

@@ -0,0 +1,186 @@
+# Author: Junjun
+# Date: 2025/8/20
+from apps.datasource.models.datasource import CoreDatasource, DatasourceConf
+
+
+def get_table_sql(ds: CoreDatasource, conf: DatasourceConf):
+    if ds.type == "mysql":
+        return f"""
+                SELECT 
+                    TABLE_NAME, 
+                    TABLE_COMMENT
+                FROM 
+                    information_schema.TABLES
+                WHERE 
+                    TABLE_SCHEMA = '{conf.database}'
+                """
+    elif ds.type == "sqlServer":
+        return 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":
+        return """
+              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":
+        return 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":
+        return f"""
+                SELECT name, comment
+                FROM system.tables
+                WHERE database = '{conf.database}'
+                  AND engine NOT IN ('Dictionary')
+                ORDER BY name
+                """
+    elif ds.type == 'dm':
+        return f"""
+                select table_name, comments 
+                from all_tab_comments 
+                where owner='{conf.dbSchema}'
+                AND (table_type = 'TABLE' or table_type = 'VIEW')
+                """
+
+
+def get_field_sql(ds: CoreDatasource, conf: DatasourceConf, table_name: str = None):
+    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 ""
+        return 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 ""
+        return 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 ""
+        return 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 ""
+        return 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 ""
+        return sql1 + sql2
+    elif ds.type == 'dm':
+        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 ""
+        return sql1 + sql2