feat: Ask to add database version information

This commit is contained in:
junjun
2025-08-26 10:42:26 +08:00
parent 030392c5ce
commit a9f4b5c64e
3 changed files with 62 additions and 6 deletions
+6 -5
View File
@@ -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
+30 -1
View File
@@ -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)
+26
View File
@@ -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"""