feat: Ask to add database version information

Author: XiaJunjie2020

backend/apps/chat/task/llm.py

@@ -28,7 +28,7 @@
 from apps.datasource.crud.datasource import get_table_schema
 from apps.datasource.crud.permission import get_row_permission_filters, is_normal_user
 from apps.datasource.models.datasource import CoreDatasource
-from apps.db.db import exec_sql
+from apps.db.db import exec_sql, get_version
 from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_ds
 from apps.system.schemas.system_schema import AssistantOutDsSchema
 from common.core.config import settings
@@ -91,13 +91,13 @@ def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
                 ds = self.out_ds_instance.get_ds(chat.datasource)
                 if not ds:
                     raise SingleMessageError("No available datasource configuration found")
-                chat_question.engine = ds.type
+                chat_question.engine = ds.type + get_version(ds)
                 chat_question.db_schema = self.out_ds_instance.get_db_schema(ds.id)
             else:
                 ds = self.session.get(CoreDatasource, chat.datasource)
                 if not ds:
                     raise SingleMessageError("No available datasource configuration found")
-                chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL'
+                chat_question.engine = (ds.type_name if ds.type != 'excel' else 'PostgreSQL') + get_version(ds)
                 chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds)
 
         self.generate_sql_logs = list_generate_sql_logs(session=self.session, chart_id=chat_id)
@@ -451,7 +451,7 @@ def select_datasource(self):
                 if self.current_assistant and self.current_assistant.type in dynamic_ds_types:
                     _ds = self.out_ds_instance.get_ds(data['id'])
                     self.ds = _ds
-                    self.chat_question.engine = _ds.type
+                    self.chat_question.engine = _ds.type + get_version(self.ds)
                     self.chat_question.db_schema = self.out_ds_instance.get_db_schema(self.ds.id)
                     _engine_type = self.chat_question.engine
                     _chat.engine_type = _ds.type
@@ -461,7 +461,8 @@ def select_datasource(self):
                         _datasource = None
                         raise SingleMessageError(f"Datasource configuration with id {_datasource} not found")
                     self.ds = CoreDatasource(**_ds.model_dump())
-                    self.chat_question.engine = _ds.type_name if _ds.type != 'excel' else 'PostgreSQL'
+                    self.chat_question.engine = (_ds.type_name if _ds.type != 'excel' else 'PostgreSQL') + get_version(
+                        self.ds)
                     self.chat_question.db_schema = get_table_schema(session=self.session,
                                                                     current_user=self.current_user, ds=self.ds)
                     _engine_type = self.chat_question.engine

backend/apps/db/db.py

@@ -4,7 +4,7 @@
 import urllib.parse
 from decimal import Decimal
 
-from apps.db.db_sql import get_table_sql, get_field_sql
+from apps.db.db_sql import get_table_sql, get_field_sql, get_version_sql
 
 if platform.system() != "Darwin":
     import dmPython
@@ -141,6 +141,35 @@ def check_connection(trans: Trans, ds: CoreDatasource, is_raise: bool = False):
                     return False
 
 
+def get_version(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_version_sql(ds, conf)
+    try:
+        if db.connect_type == ConnectType.sqlalchemy:
+            with get_session(ds) as session:
+                with session.execute(text(sql)) as result:
+                    res = result.fetchall()
+                    return res[0][0]
+        else:
+            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(sql, timeout=10)
+                    res = cursor.fetchall()
+                    return res[0][0]
+            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:
+                    cursor.execute(sql)
+                    res = cursor.fetchall()
+                    return res[0][0]
+    except Exception as e:
+        print(e)
+        return ''
+
+
 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)

backend/apps/db/db_sql.py

@@ -3,6 +3,32 @@
 from apps.datasource.models.datasource import CoreDatasource, DatasourceConf
 
 
+def get_version_sql(ds: CoreDatasource, conf: DatasourceConf):
+    if ds.type == "mysql" or ds.type == "doris":
+        return f"""
+                SELECT VERSION()
+                """
+    elif ds.type == "sqlServer":
+        return f"""
+                select SERVERPROPERTY('ProductVersion')
+                """
+    elif ds.type == "pg" or ds.type == "excel":
+        return f"""
+              select  version()
+              """
+    elif ds.type == "oracle":
+        return f"""
+                SELECT version FROM v$instance
+                """
+    elif ds.type == "ck":
+        return f"""
+                select  version()
+                """
+    elif ds.type == 'dm':
+        return f"""
+                SELECT * FROM v$version
+                """
+
 def get_table_sql(ds: CoreDatasource, conf: DatasourceConf):
     if ds.type == "mysql" or ds.type == "doris":
         return f"""