mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 09:17:06 +08:00
feat: Ask to add database version information
This commit is contained in:
@@ -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
@@ -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)
|
||||
|
||||
@@ -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"""
|
||||
|
||||
Reference in New Issue
Block a user