|
1 | 1 | """Summary for rdbms database."""
|
2 |
| - |
| 2 | +import re |
3 | 3 | from typing import TYPE_CHECKING, List, Optional
|
4 | 4 |
|
5 | 5 | from dbgpt._private.config import Config
|
@@ -102,10 +102,20 @@ def _parse_table_summary(
|
102 | 102 | columns.append(f"{column['name']}")
|
103 | 103 |
|
104 | 104 | column_str = ", ".join(columns)
|
| 105 | + # Obtain index information |
105 | 106 | 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}`) ") |
109 | 119 | table_str = summary_template.format(table_name=table_name, columns=column_str)
|
110 | 120 | if len(index_keys) > 0:
|
111 | 121 | index_key_str = ", ".join(index_keys)
|
|
0 commit comments