refactor: code optimization

backend/apps/datasource/crud/datasource.py

@@ -11,7 +11,6 @@ from apps.datasource.utils.utils import aes_decrypt
 from apps.db.constant import DB
 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 deepcopy_ignore_extra
 from .table import get_tables_by_ds_id
@@ -70,7 +69,7 @@ def create_ds(session: SessionDep, trans: Trans, user: CurrentUser, create_ds: C
     ds.create_by = user.id
     ds.oid = user.oid if user.oid is not None else 1
     ds.status = "Success"
-    ds.type_name = db_type_relation()[ds.type]
+    ds.type_name = DB.get_db(ds.type).db_name
     record = CoreDatasource(**ds.model_dump())
     session.add(record)
     session.flush()

backend/apps/db/constant.py

@@ -13,19 +13,20 @@ class ConnectType(Enum):
 
 
 class DB(Enum):
-    mysql = ('mysql', '`', '`', ConnectType.sqlalchemy)
-    sqlServer = ('sqlServer', '[', ']', ConnectType.sqlalchemy)
-    pg = ('pg', '"', '"', ConnectType.sqlalchemy)
-    excel = ('excel', '"', '"', ConnectType.sqlalchemy)
-    oracle = ('oracle', '"', '"', ConnectType.sqlalchemy)
-    ck = ('ck', '"', '"', ConnectType.sqlalchemy)
-    dm = ('dm', '"', '"', ConnectType.py_driver)
-    doris = ('doris', '`', '`', ConnectType.py_driver)
-    redshift = ('redshift', '"', '"', ConnectType.py_driver)
-    es = ('es', '"', '"', ConnectType.py_driver)
+    mysql = ('mysql', 'MySQL', '`', '`', ConnectType.sqlalchemy)
+    sqlServer = ('sqlServer', 'Microsoft SQL Server', '[', ']', ConnectType.sqlalchemy)
+    pg = ('pg', 'PostgreSQL', '"', '"', ConnectType.sqlalchemy)
+    excel = ('excel', 'Excel/CSV', '"', '"', ConnectType.sqlalchemy)
+    oracle = ('oracle', 'Oracle', '"', '"', ConnectType.sqlalchemy)
+    ck = ('ck', 'ClickHouse', '"', '"', ConnectType.sqlalchemy)
+    dm = ('dm', '达梦', '"', '"', ConnectType.py_driver)
+    doris = ('doris', 'Apache Doris', '`', '`', ConnectType.py_driver)
+    redshift = ('redshift', 'AWS Redshift', '"', '"', ConnectType.py_driver)
+    es = ('es', 'Elasticsearch', '"', '"', ConnectType.py_driver)
 
-    def __init__(self, type, prefix, suffix, connect_type: ConnectType):
+    def __init__(self, type, db_name, prefix, suffix, connect_type: ConnectType):
         self.type = type
+        self.db_name = db_name
         self.prefix = prefix
         self.suffix = suffix
         self.connect_type = connect_type

backend/apps/db/type.py

@@ -1,18 +0,0 @@
-# Author: Junjun
-# Date: 2025/5/22
-from typing import Dict
-
-
-def db_type_relation() -> Dict:
-    return {
-        "mysql": "MySQL",
-        "sqlServer": "Microsoft SQL Server",
-        "pg": "PostgreSQL",
-        "excel": "Excel/CSV",
-        "oracle": "Oracle",
-        "ck": "ClickHouse",
-        "dm": "达梦",
-        "doris": "Apache Doris",
-        "redshift": "AWS Redshift",
-        "es": "Elasticsearch"
-    }