From a9f4b5c64e1e3cc84493311cf88b87d10458a013 Mon Sep 17 00:00:00 2001 From: junjun Date: Tue, 26 Aug 2025 10:42:26 +0800 Subject: [PATCH] feat: Ask to add database version information --- backend/apps/chat/task/llm.py | 11 ++++++----- backend/apps/db/db.py | 31 ++++++++++++++++++++++++++++++- backend/apps/db/db_sql.py | 26 ++++++++++++++++++++++++++ 3 files changed, 62 insertions(+), 6 deletions(-) diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 0f47af7f..55279269 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -28,7 +28,7 @@ from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameCh 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: _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 diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index 5081c2bb..73aea904 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -4,7 +4,7 @@ import platform 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) diff --git a/backend/apps/db/db_sql.py b/backend/apps/db/db_sql.py index 54be1617..30f06409 100644 --- a/backend/apps/db/db_sql.py +++ b/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"""