(self, ds_id: int, question: str = '', embedding: bool = True,
table_list: list[str] = None)
| 181 | raise Exception("Datasource list is not found.") |
| 182 | |
| 183 | def get_db_schema(self, ds_id: int, question: str = '', embedding: bool = True, |
| 184 | table_list: list[str] = None) -> tuple[str, list]: |
| 185 | ds = self.get_ds(ds_id) |
| 186 | schema_str = "" |
| 187 | db_name = ds.db_schema if ds.db_schema is not None and ds.db_schema != "" else ds.dataBase |
| 188 | schema_str += f"【DB_ID】 {db_name}\n【Schema】\n" |
| 189 | tables = [] |
| 190 | i = 0 |
| 191 | for table in ds.tables: |
| 192 | # 如果传入了 table_list，则只处理在列表中的表 |
| 193 | if table_list is not None and table.name not in table_list: |
| 194 | continue |
| 195 | |
| 196 | i += 1 |
| 197 | schema_table = '' |
| 198 | schema_table += f"# Table: {db_name}.{table.name}" if ds.type != "mysql" and ds.type != "es" else f"# Table: {table.name}" |
| 199 | table_comment = table.comment |
| 200 | if table_comment == '': |
| 201 | schema_table += '\n[\n' |
| 202 | else: |
| 203 | schema_table += f", {table_comment}\n[\n" |
| 204 | |
| 205 | field_list = [] |
| 206 | for field in table.fields: |
| 207 | field_comment = field.comment |
| 208 | if field_comment == '': |
| 209 | field_list.append(f"({field.name}:{field.type})") |
| 210 | else: |
| 211 | field_list.append(f"({field.name}:{field.type}, {field_comment})") |
| 212 | schema_table += ",\n".join(field_list) |
| 213 | schema_table += '\n]\n' |
| 214 | t_obj = {"id": i, "schema_table": schema_table} |
| 215 | tables.append(t_obj) |
| 216 | |
| 217 | # do table embedding |
| 218 | # if embedding and tables and settings.TABLE_EMBEDDING_ENABLED: |
| 219 | # tables = get_table_embedding(tables, question) |
| 220 | |
| 221 | if tables: |
| 222 | for s in tables: |
| 223 | schema_str += s.get('schema_table') |
| 224 | |
| 225 | return schema_str, [] |
| 226 | |
| 227 | def get_ds(self, ds_id: int, trans: Trans = None): |
| 228 | if self.ds_list: |
no test coverage detected