Skip to content

Commit f3ece62

Browse files
authored
fix rdbms_db_summary (#1409)
1 parent 5830093 commit f3ece62

File tree

1 file changed

+14
-4
lines changed

1 file changed

+14
-4
lines changed

dbgpt/rag/summary/rdbms_db_summary.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
"""Summary for rdbms database."""
2-
2+
import re
33
from typing import TYPE_CHECKING, List, Optional
44

55
from dbgpt._private.config import Config
@@ -102,10 +102,20 @@ def _parse_table_summary(
102102
columns.append(f"{column['name']}")
103103

104104
column_str = ", ".join(columns)
105+
# Obtain index information
105106
index_keys = []
106-
for index_key in conn.get_indexes(table_name):
107-
key_str = ", ".join(index_key["column_names"])
108-
index_keys.append(f"{index_key['name']}(`{key_str}`) ") # noqa
107+
raw_indexes = conn.get_indexes(table_name)
108+
for index in raw_indexes:
109+
if isinstance(index, tuple): # Process tuple type index information
110+
index_name, index_creation_command = index
111+
# Extract column names using re
112+
matched_columns = re.findall(r"\(([^)]+)\)", index_creation_command)
113+
if matched_columns:
114+
key_str = ", ".join(matched_columns)
115+
index_keys.append(f"{index_name}(`{key_str}`) ")
116+
else:
117+
key_str = ", ".join(index["column_names"])
118+
index_keys.append(f"{index['name']}(`{key_str}`) ")
109119
table_str = summary_template.format(table_name=table_name, columns=column_str)
110120
if len(index_keys) > 0:
111121
index_key_str = ", ".join(index_keys)

0 commit comments

Comments
 (0)