Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 18 additions & 11 deletions backend/apps/chat/curd/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,8 @@ def get_last_execute_sql_error(session: SessionDep, chart_id: int):


def format_json_data(origin_data: dict):
result = {'fields': origin_data.get('fields') if origin_data.get('fields') else []}
result = {'fields': origin_data.get('fields') if origin_data.get('fields') else [],
'fields_info': origin_data.get('fields_info') if origin_data.get('fields_info') else None}
_list = origin_data.get('data') if origin_data.get('data') else []
data = format_json_list_data(_list)
result['data'] = data
Expand Down Expand Up @@ -238,21 +239,24 @@ def get_chart_data_with_user(session: SessionDep, current_user: CurrentUser, cha
pass
return {}


def get_chart_data_with_user_live(session: SessionDep, current_user: CurrentUser, chat_record_id: int):
stmt = select(ChatRecord.datasource,ChatRecord.sql).where(and_(ChatRecord.id == chat_record_id, ChatRecord.create_by == current_user.id))
stmt = select(ChatRecord.datasource, ChatRecord.sql).where(
and_(ChatRecord.id == chat_record_id, ChatRecord.create_by == current_user.id))
row = session.execute(stmt).first()
return get_chart_data_ds(session,row.datasource, row.sql)
return get_chart_data_ds(session, row.datasource, row.sql)

def get_chart_data_ds(session: SessionDep,ds_id,sql):
json_result: Dict[str, Any] = {'status': 'success','data':[],'message':''}

def get_chart_data_ds(session: SessionDep, ds_id, sql):
json_result: Dict[str, Any] = {'status': 'success', 'data': [], 'message': ''}
try:
datasource = get_ds(session,ds_id)
datasource = get_ds(session, ds_id)
if datasource is None:
json_result['status'] = 'failed'
json_result['message'] = 'Datasource not found'
return json_result
else:
result = exec_sql(ds=datasource,sql=sql, origin_column=False)
result = exec_sql(ds=datasource, sql=sql, origin_column=False)
_data = DataFormat.convert_large_numbers_in_object_array(result.get('data'))
_data = DataFormat.normalize_qualified_sql_column_keys_in_object_array(_data)
json_result['data'] = _data
Expand All @@ -264,6 +268,7 @@ def get_chart_data_ds(session: SessionDep,ds_id,sql):
pass
return json_result


def get_chat_chart_data(session: SessionDep, chat_record_id: int):
stmt = select(ChatRecord.data).where(and_(ChatRecord.id == chat_record_id))
res = session.execute(stmt)
Expand Down Expand Up @@ -336,7 +341,7 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
predict_alias_log = aliased(ChatLog)

stmt = (select(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time,
ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql,ChatRecord.datasource,
ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql, ChatRecord.datasource,
ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict,
ChatRecord.datasource_select_answer, ChatRecord.analysis_record_id, ChatRecord.predict_record_id,
ChatRecord.regenerate_record_id,
Expand All @@ -363,7 +368,7 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
ChatRecord.create_time))
if with_data:
stmt = select(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time,
ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql,ChatRecord.datasource,
ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql, ChatRecord.datasource,
ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict,
ChatRecord.datasource_select_answer, ChatRecord.analysis_record_id, ChatRecord.predict_record_id,
ChatRecord.regenerate_record_id,
Expand Down Expand Up @@ -429,7 +434,8 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
finish_time=row.finish_time,
duration=duration,
total_tokens=total_tokens,
question=row.question, sql_answer=row.sql_answer, sql=row.sql, datasource=row.datasource,
question=row.question, sql_answer=row.sql_answer, sql=row.sql,
datasource=row.datasource,
chart_answer=row.chart_answer, chart=row.chart,
analysis=row.analysis, predict=row.predict,
datasource_select_answer=row.datasource_select_answer,
Expand All @@ -448,7 +454,8 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
finish_time=row.finish_time,
duration=duration,
total_tokens=total_tokens,
question=row.question, sql_answer=row.sql_answer, sql=row.sql, datasource=row.datasource,
question=row.question, sql_answer=row.sql_answer, sql=row.sql,
datasource=row.datasource,
chart_answer=row.chart_answer, chart=row.chart,
analysis=row.analysis, predict=row.predict,
datasource_select_answer=row.datasource_select_answer,
Expand Down
160 changes: 147 additions & 13 deletions backend/apps/db/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
import dmPython
import pymysql
import redshift_connector
from sqlalchemy import create_engine, text, Engine
from sqlalchemy import create_engine, text, Engine, types
from sqlalchemy.orm import sessionmaker

from apps.datasource.models.datasource import DatasourceConf, CoreDatasource, TableSchema, ColumnSchema
Expand Down Expand Up @@ -596,12 +596,35 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
with session.execute(text(sql)) as result:
try:
columns = result.keys()._keys if origin_column else [item.lower() for item in result.keys()._keys]

fields_info = []
for col_info in result.cursor.description:
# col_info 是 (name, type_code, display_size, internal_size, precision, scale, null_ok)
col_name = col_info[0]

# 根据 type_code 判断是否为数值类型
# psycopg2 的类型 OID 常量
is_numeric = col_info[1] in (
20, # int8
21, # int2
23, # int4
700, # float4
701, # float8
1700, # numeric
16, # boolean
)

fields_info.append({
"name": col_name if origin_column else col_name.lower(),
"is_numeric": is_numeric
})

res = result.fetchall()
result_list = [
{str(columns[i]): convert_value(value) for i, value in enumerate(tuple_item)} for tuple_item in
res
]
return {"fields": columns, "data": result_list,
return {"fields": columns, "data": result_list, "fields_info": fields_info,
"sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
except Exception as ex:
raise ParseSQLResultError(str(ex))
Expand All @@ -617,11 +640,12 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
columns = [field[0] for field in cursor.description] if origin_column else [field[0].lower() for
field in
cursor.description]
fields_info = build_fields_info_from_cursor(cursor, origin_column, 'dm')
result_list = [
{str(columns[i]): convert_value(value) for i, value in enumerate(tuple_item)} for tuple_item in
res
]
return {"fields": columns, "data": result_list,
return {"fields": columns, "data": result_list, "fields_info": fields_info,
"sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
except Exception as ex:
raise ParseSQLResultError(str(ex))
Expand All @@ -637,11 +661,12 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
columns = [field[0] for field in cursor.description] if origin_column else [field[0].lower() for
field in
cursor.description]
fields_info = build_fields_info_from_cursor(cursor, origin_column, 'mysql')
result_list = [
{str(columns[i]): convert_value(value) for i, value in enumerate(tuple_item)} for tuple_item in
res
]
return {"fields": columns, "data": result_list,
return {"fields": columns, "data": result_list, "fields_info": fields_info,
"sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
except Exception as ex:
raise ParseSQLResultError(str(ex))
Expand All @@ -655,11 +680,12 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
columns = [field[0] for field in cursor.description] if origin_column else [field[0].lower() for
field in
cursor.description]
fields_info = build_fields_info_from_cursor(cursor, origin_column, 'postgresql')
result_list = [
{str(columns[i]): convert_value(value) for i, value in enumerate(tuple_item)} for tuple_item in
res
]
return {"fields": columns, "data": result_list,
return {"fields": columns, "data": result_list, "fields_info": fields_info,
"sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
except Exception as ex:
raise ParseSQLResultError(str(ex))
Expand All @@ -674,25 +700,28 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
columns = [field[0] for field in cursor.description] if origin_column else [field[0].lower() for
field in
cursor.description]
fields_info = build_fields_info_from_cursor(cursor, origin_column, 'postgresql')
result_list = [
{str(columns[i]): convert_value(value) for i, value in enumerate(tuple_item)} for tuple_item in
res
]
return {"fields": columns, "data": result_list,
return {"fields": columns, "data": result_list, "fields_info": fields_info,
"sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
except Exception as ex:
raise ParseSQLResultError(str(ex))
elif equals_ignore_case(ds.type, 'es'):
try:
res, columns = get_es_data_by_http(conf, sql)
columns = [field.get('name') for field in columns] if origin_column else [field.get('name').lower() for
field in
columns]
res, raw_columns = get_es_data_by_http(conf, sql)
columns = [field.get('name') for field in raw_columns] if origin_column else [field.get('name').lower()
for
field in
raw_columns]
fields_info = build_fields_info_from_es(raw_columns, origin_column)
result_list = [
{str(columns[i]): convert_value(value) for i, value in enumerate(tuple_item)} for tuple_item in
res
]
return {"fields": columns, "data": result_list,
return {"fields": columns, "data": result_list, "fields_info": fields_info,
"sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
except Exception as ex:
raise Exception(str(ex))
Expand All @@ -707,16 +736,120 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
columns = [field[0] for field in cursor.description] if origin_column else [field[0].lower() for
field in
cursor.description]
fields_info = build_fields_info_from_cursor(cursor, origin_column, 'hive')
result_list = [
{str(columns[i]): convert_value(value) for i, value in enumerate(tuple_item)} for tuple_item in
res
]
return {"fields": columns, "data": result_list,
return {"fields": columns, "data": result_list, "fields_info": fields_info,
"sql": bytes.decode(base64.b64encode(bytes(hive_sql, 'utf-8')))}
except Exception as ex:
raise ParseSQLResultError(str(ex))


def build_fields_info_from_cursor(cursor, origin_column, db_type='postgresql'):
"""
根据数据库游标的 description 构建字段信息列表

Args:
cursor: 数据库游标对象
origin_column: 是否保留原始列名大小写
db_type: 数据库类型,支持 'mysql', 'postgresql', 'redshift', 'kingbase', 'dm', 'hive'

Returns:
list: 包含字段名和是否数值类型的字典列表
"""
fields_info = []

for col_info in cursor.description:
col_name = col_info[0]

if db_type == 'mysql':
# MySQL/pymysql 类型码
is_numeric = col_info[1] in (
1, # TINYINT
2, # SMALLINT
3, # INT
4, # FLOAT
5, # DOUBLE
8, # BIGINT
9, # MEDIUMINT
16, # BIT
246, # DECIMAL
)
elif db_type in ('postgresql', 'redshift', 'kingbase'):
# PostgreSQL/psycopg2 类型 OID
is_numeric = col_info[1] in (
20, # int8
21, # int2
23, # int4
700, # float4
701, # float8
1700, # numeric
16, # boolean
)
elif db_type == 'dm':
# 达梦数据库类型码
is_numeric = col_info[1] in (
3, # DECIMAL/NUMERIC
2, # NUMBER
4, # INTEGER
5, # INT
6, # BIGINT
7, # TINYINT
8, # BYTE
9, # FLOAT
10, # DOUBLE
11, # REAL
12, # BOOLEAN
)
elif db_type == 'hive':
# Hive 类型对象转字符串判断
type_str = str(col_info[1]).lower()
NUMERIC_PREFIXES = ('tinyint', 'smallint', 'int', 'bigint', 'float', 'double', 'decimal', 'numeric')
is_numeric = type_str == 'boolean' or any(type_str.startswith(p) for p in NUMERIC_PREFIXES)
else:
is_numeric = False

fields_info.append({
"name": col_name if origin_column else col_name.lower(),
"is_numeric": is_numeric
})

return fields_info


def build_fields_info_from_es(raw_columns, origin_column):
"""
专门为 Elasticsearch 构建字段信息

Args:
raw_columns: ES 返回的列信息列表
origin_column: 是否保留原始列名大小写

Returns:
list: 包含字段名和是否数值类型的字典列表
"""
fields_info = []

for field in raw_columns:
field_name = field.get('name') if origin_column else field.get('name').lower()
field_type = field.get('type', '').lower()

is_numeric = field_type in (
'long', 'integer', 'short', 'byte',
'double', 'float', 'half_float', 'scaled_float',
'unsigned_long', 'boolean'
)

fields_info.append({
"name": field_name,
"is_numeric": is_numeric
})

return fields_info


def get_sqlglot_dialect(ds_type: str) -> str:
"""根据数据源类型获取 sqlglot dialect"""
if equals_ignore_case(ds_type, 'mysql', 'doris', 'starrocks'):
Expand Down Expand Up @@ -744,6 +877,7 @@ def get_sqlglot_dialect(ds_type: str) -> str:

# 危险模式正则表达式(用于检查特殊语法)
import re

DANGEROUS_PATTERNS = [
r'\bINTO\s+OUTFILE\b',
r'\bINTO\s+DUMPFILE\b',
Expand All @@ -765,7 +899,7 @@ def check_dangerous_functions(statements: list, ds_type: str) -> bool:
"""检查是否使用了危险函数,返回 True 表示安全"""
dangerous_functions = get_dangerous_functions(ds_type)
dangerous_functions_upper = {f.upper() for f in dangerous_functions}

for stmt in statements:
if stmt:
for func in stmt.find_all(exp.Anonymous):
Expand Down
11 changes: 9 additions & 2 deletions frontend/src/views/chat/chat-block/ChartBlock.vue
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import DisplayChartBlock from '@/views/chat/component/DisplayChartBlock.vue'
import ChartPopover from '@/views/chat/chat-block/ChartPopover.vue'
import { computed, ref, watch } from 'vue'
import { useClipboard } from '@vueuse/core'
import { concat } from 'lodash-es'
import { concat, filter, includes, map } from 'lodash-es'
import type { ChartTypes } from '@/views/chat/component/BaseChart.ts'
import ICON_BAR from '@/assets/svg/chart/icon_bar_outlined.svg'
import ICON_COLUMN from '@/assets/svg/chart/icon_dashboard_outlined.svg'
Expand Down Expand Up @@ -61,6 +61,7 @@ const emits = defineEmits(['exitFullScreen', 'update:thousandsSeparatorList', 'u

const dataObject = computed<{
fields: Array<string>
fields_info: Array<{ name: string; is_numeric: boolean }>
data: Array<{ [key: string]: any }>
limit: number | undefined
datasource: number | undefined
Expand Down Expand Up @@ -388,7 +389,13 @@ const enableThousandsSeparatorList = computed({

const optionList = ref<Array<{ name: string; value: string }>>([])
function getBaseAxis() {
optionList.value = chartRef.value?.getBaseAxis()
const _list = chartRef.value?.getBaseAxis()
if (dataObject.value.fields_info) {
const numberList = map(filter(dataObject.value.fields_info, { is_numeric: true }), 'name')
optionList.value = filter(_list, (obj) => includes(numberList, obj.value))
} else {
optionList.value = _list
}
}
</script>

Expand Down
Loading