github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-08-03 17:39:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix rdbms_db_summary (#1409)

Browse Source
This commit is contained in:
yyhhyy 2024-04-13 11:03:49 +08:00 committed by GitHub
parent 583009305a
commit f3ece627d0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 14 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
dbgpt/rag/summary/rdbms_db_summary.py
View File

@ -1,5 +1,5 @@
"""Summary for rdbms database.""" """Summary for rdbms database."""
import re
from typing import TYPE_CHECKING, List, Optional from typing import TYPE_CHECKING, List, Optional
from dbgpt._private.config import Config from dbgpt._private.config import Config
@ -102,10 +102,20 @@ def _parse_table_summary(
columns.append(f"{column['name']}") columns.append(f"{column['name']}")
column_str = ", ".join(columns) column_str = ", ".join(columns)
# Obtain index information
index_keys = [] index_keys = []
for index_key in conn.get_indexes(table_name): raw_indexes = conn.get_indexes(table_name)
key_str = ", ".join(index_key["column_names"]) for index in raw_indexes:
index_keys.append(f"{index_key['name']}(`{key_str}`) ") # noqa if isinstance(index, tuple): # Process tuple type index information
index_name, index_creation_command = index
# Extract column names using re
matched_columns = re.findall(r"\(([^)]+)\)", index_creation_command)
if matched_columns:
key_str = ", ".join(matched_columns)
index_keys.append(f"{index_name}(`{key_str}`) ")
else:
key_str = ", ".join(index["column_names"])
index_keys.append(f"{index['name']}(`{key_str}`) ")
table_str = summary_template.format(table_name=table_name, columns=column_str) table_str = summary_template.format(table_name=table_name, columns=column_str)
if len(index_keys) > 0: if len(index_keys) > 0:
index_key_str = ", ".join(index_keys) index_key_str = ", ".join(index_keys)