feat: support DM datasource

Author: XiaJunjie2020

backend/apps/datasource/crud/datasource.py

@@ -2,12 +2,14 @@
 import json
 from typing import List, Optional
 
+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.engine import get_engine_config, get_engine_conn
@@ -45,16 +47,32 @@ 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):
-    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
+    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').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 check_name(session: SessionDep, trans: Trans, user: CurrentUser, ds: CoreDatasource):
@@ -159,7 +177,7 @@ def getFieldsByDs(session: SessionDep, ds: CoreDatasource, table_name: str):
 
 def execSql(session: SessionDep, id: int, sql: str):
     ds = session.exec(select(CoreDatasource).where(CoreDatasource.id == id)).first()
-    return exec_sql(ds, sql)
+    return exec_sql(ds, sql, True)
 
 
 def sync_table(session: SessionDep, ds: CoreDatasource, tables: List[CoreTable]):
@@ -297,7 +315,11 @@ def preview(session: SessionDep, current_user: CurrentUser, id: int, data: Table
         sql = f"""SELECT "{'", "'.join(fields)}" FROM "{data.table.table_name}" 
             {where} 
             LIMIT 100"""
-    return exec_sql(ds, sql)
+    elif ds.type == "dm":
+        sql = f"""SELECT "{'", "'.join(fields)}" FROM "{conf.dbSchema}"."{data.table.table_name}" 
+            {where} 
+            LIMIT 100"""
+    return exec_sql(ds, sql, True)
 
 
 def fieldEnum(session: SessionDep, id: int):
@@ -313,7 +335,7 @@ def fieldEnum(session: SessionDep, id: int):
 
     db = DB.get_db(ds.type)
     sql = f"""SELECT DISTINCT {db.prefix}{field.field_name}{db.suffix} FROM {db.prefix}{table.table_name}{db.suffix}"""
-    res = exec_sql(ds, sql)
+    res = exec_sql(ds, sql, True)
     return [item.get(res.get('fields')[0]) for item in res.get('data')]
 
 

backend/apps/db/constant.py

@@ -4,18 +4,28 @@
 from enum import Enum
 
 
+class ConnectType(Enum):
+    sqlalchemy = ('sqlalchemy')
+    py_driver = ('py_driver')
+
+    def __init__(self, type_name):
+        self.type_name = type_name
+
+
 class DB(Enum):
-    mysql = ('mysql', '`', '`')
-    sqlServer = ('sqlServer', '[', ']')
-    pg = ('pg', '"', '"')
-    excel = ('excel', '"', '"')
-    oracle = ('oracle', '"', '"')
-    ck = ('ClickHouse', '"', '"')
-
-    def __init__(self, type, prefix, suffix):
+    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)
+
+    def __init__(self, type, prefix, suffix, connect_type: ConnectType):
         self.type = type
         self.prefix = prefix
         self.suffix = suffix
+        self.connect_type = connect_type
 
     @classmethod
     def get_db(cls, type):