From db2fdfc4c241eb6795b2cd6a0b50b31146e94d8c Mon Sep 17 00:00:00 2001 From: junjun Date: Wed, 20 Aug 2025 14:11:37 +0800 Subject: [PATCH] refactor: code optimization --- backend/apps/datasource/crud/datasource.py | 35 +--- backend/apps/db/db.py | 215 ++++----------------- backend/apps/db/db_sql.py | 186 ++++++++++++++++++ 3 files changed, 227 insertions(+), 209 deletions(-) create mode 100644 backend/apps/db/db_sql.py diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index cd4f4e84..47e72b05 100644 --- 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 index e5411dce..69c44f8f 100644 --- a/backend/apps/db/db.py +++ b/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 diff --git a/backend/apps/db/db_sql.py b/backend/apps/db/db_sql.py new file mode 100644 index 00000000..d7ba83a8 --- /dev/null +++ b/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