Merge branch 'main' of https://github.com/dataease/SQLBot

Author: dataeaseShu

backend/apps/datasource/crud/datasource.py

@@ -267,7 +267,7 @@ def preview(session: SessionDep, current_user: CurrentUser, id: int, data: Table
 
     conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     sql: str = ""
-    if ds.type == "mysql":
+    if ds.type == "mysql" or ds.type == "doris":
         sql = f"""SELECT `{"`, `".join(fields)}` FROM `{data.table.table_name}` 
             {where} 
             LIMIT 100"""

backend/apps/db/constant.py

@@ -20,6 +20,7 @@ class DB(Enum):
     oracle = ('oracle', '"', '"', ConnectType.sqlalchemy)
     ck = ('ck', '"', '"', ConnectType.sqlalchemy)
     dm = ('dm', '"', '"', ConnectType.py_driver)
+    doris = ('doris', '`', '`', ConnectType.py_driver)
 
     def __init__(self, type, prefix, suffix, connect_type: ConnectType):
         self.type = type

backend/apps/db/db.py

@@ -8,6 +8,7 @@
 
 if platform.system() != "Darwin":
     import dmPython
+import pymysql
 from sqlalchemy import create_engine, text, Engine
 from sqlalchemy.orm import sessionmaker
 
@@ -125,6 +126,19 @@ def check_connection(trans: Trans, ds: CoreDatasource, is_raise: bool = False):
                     if is_raise:
                         raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
                     return False
+        elif ds.type == 'doris':
+            with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
+                                 port=conf.port, db=conf.database, connect_timeout=10,
+                                 read_timeout=10) as conn, conn.cursor() as cursor:
+                try:
+                    cursor.execute('select 1')
+                    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):
@@ -171,6 +185,14 @@ def get_tables(ds: CoreDatasource):
                 res = cursor.fetchall()
                 res_list = [TableSchema(*item) for item in res]
                 return res_list
+        elif ds.type == 'doris':
+            with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
+                                 port=conf.port, db=conf.database, connect_timeout=conf.timeout,
+                                 read_timeout=conf.timeout) as conn, conn.cursor() as cursor:
+                cursor.execute(sql)
+                res = cursor.fetchall()
+                res_list = [TableSchema(*item) for item in res]
+                return res_list
 
 
 def get_fields(ds: CoreDatasource, table_name: str = None):
@@ -191,6 +213,14 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
                 res = cursor.fetchall()
                 res_list = [ColumnSchema(*item) for item in res]
                 return res_list
+        elif ds.type == 'doris':
+            with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
+                                 port=conf.port, db=conf.database, connect_timeout=conf.timeout,
+                                 read_timeout=conf.timeout) as conn, conn.cursor() as cursor:
+                cursor.execute(sql)
+                res = cursor.fetchall()
+                res_list = [ColumnSchema(*item) for item in res]
+                return res_list
 
 
 def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=False):
@@ -233,3 +263,22 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
                             "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
                 except Exception as ex:
                     raise ex
+        elif ds.type == 'doris':
+            with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
+                                 port=conf.port, db=conf.database, connect_timeout=conf.timeout,
+                                 read_timeout=conf.timeout) as conn, conn.cursor() as cursor:
+                try:
+                    cursor.execute(sql)
+                    res = cursor.fetchall()
+                    columns = [field[0] for field in cursor.description] if origin_column else [field[0].lower() for
+                                                                                                field in
+                                                                                                cursor.description]
+                    result_list = [
+                        {str(columns[i]): float(value) if isinstance(value, Decimal) else value for i, value in
+                         enumerate(tuple_item)}
+                        for tuple_item in res
+                    ]
+                    return {"fields": columns, "data": result_list,
+                            "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
+                except Exception as ex:
+                    raise ex

backend/apps/db/db_sql.py

@@ -4,7 +4,7 @@
 
 
 def get_table_sql(ds: CoreDatasource, conf: DatasourceConf):
-    if ds.type == "mysql":
+    if ds.type == "mysql" or ds.type == "doris":
         return f"""
                 SELECT 
                     TABLE_NAME, 
@@ -31,15 +31,15 @@ def get_table_sql(ds: CoreDatasource, conf: DatasourceConf):
                     AND t.TABLE_SCHEMA = '{conf.dbSchema}'
                 """
     elif ds.type == "pg" or ds.type == "excel":
-        return """
+        return f"""
               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()
+              WHERE n.nspname = '{conf.dbSchema}'
                 AND c.relkind IN ('r', 'v', 'p', 'm')
                 AND c.relname NOT LIKE 'pg_%'
                 AND c.relname NOT LIKE 'sql_%'
@@ -83,7 +83,7 @@ def get_table_sql(ds: CoreDatasource, conf: DatasourceConf):
 
 
 def get_field_sql(ds: CoreDatasource, conf: DatasourceConf, table_name: str = None):
-    if ds.type == "mysql":
+    if ds.type == "mysql" or ds.type == "doris":
         sql1 = f"""
                 SELECT 
                     COLUMN_NAME,

backend/apps/db/type.py

@@ -11,5 +11,6 @@ def db_type_relation() -> Dict:
         "excel": "Excel/CSV",
         "oracle": "Oracle",
         "ck": "ClickHouse",
-        "dm": "达梦"
+        "dm": "达梦",
+        "doris": "Apache Doris"
     }

backend/apps/system/api/assistant.py

@@ -90,7 +90,8 @@ async def ui(session: SessionDep, data: str = Form(), files: List[UploadFile] =
             file.filename = file_name
             if flag_name == 'logo' or flag_name == 'float_icon':
                 SQLBotFileUtils.check_file(file=file, file_types=[".jpg", ".jpeg", ".png", ".svg"], limit_file_size=(10 * 1024 * 1024))
-                SQLBotFileUtils.detete_file(config_obj.get(flag_name))
+                if config_obj.get(flag_name):
+                    SQLBotFileUtils.detete_file(config_obj.get(flag_name))
                 file_id = await SQLBotFileUtils.upload(file)
                 ui_schema_dict[flag_name] = file_id
             else:

frontend/src/assets/datasource/icon_doris.png

@@ -5,6 +5,7 @@ import pg from '@/assets/datasource/icon_PostgreSQL.png'
 import sqlServer from '@/assets/datasource/icon_SQL_Server.png'
 import ck from '@/assets/datasource/icon_ck.png'
 import dm from '@/assets/datasource/icon_dm.png'
+import doris from '@/assets/datasource/icon_doris.png'
 import { i18n } from '@/i18n'
 
 const t = i18n.global.t
@@ -16,6 +17,7 @@ export const dsType = [
   { label: 'SQL Server', value: 'sqlServer' },
   { label: 'ClickHouse', value: 'ck' },
   { label: '达梦', value: 'dm' },
+  { label: 'Apache Doris', value: 'doris' },
 ]
 
 export const dsTypeWithImg = [
@@ -26,6 +28,7 @@ export const dsTypeWithImg = [
   { name: 'SQL Server', type: 'sqlServer', img: sqlServer },
   { name: 'ClickHouse', type: 'ck', img: ck },
   { name: '达梦', type: 'dm', img: dm },
+  { name: 'Apache Doris', type: 'doris', img: doris },
 ]
 
 export const haveSchema = ['sqlServer', 'pg', 'oracle', 'dm']