m_schema.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186
  1. from .utils import examples_to_str, read_json, write_json
  2. from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
  3. class MSchema:
  4. def __init__(self, db_id: str = 'Anonymous', schema: Optional[str] = None):
  5. self.db_id = db_id
  6. self.schema = schema
  7. self.tables = {}
  8. self.foreign_keys = []
  9. def add_table(self, name, fields={}, comment=None):
  10. self.tables[name] = {"fields": fields.copy(), 'examples': [], 'comment': comment}
  11. def add_field(self, table_name: str, field_name: str, field_type: str = "",
  12. primary_key: bool = False, nullable: bool = True, default: Any = None,
  13. autoincrement: bool = False, comment: str = "", examples: list = [], **kwargs):
  14. self.tables[table_name]["fields"][field_name] = {
  15. "type": field_type,
  16. "primary_key": primary_key,
  17. "nullable": nullable,
  18. "default": default if default is None else f'{default}',
  19. "autoincrement": autoincrement,
  20. "comment": comment,
  21. "examples": examples.copy(),
  22. **kwargs}
  23. def add_foreign_key(self, table_name, field_name, ref_schema, ref_table_name, ref_field_name):
  24. self.foreign_keys.append([table_name, field_name, ref_schema, ref_table_name, ref_field_name])
  25. def get_field_type(self, field_type, simple_mode=True)->str:
  26. if not simple_mode:
  27. return field_type
  28. else:
  29. return field_type.split("(")[0]
  30. def has_table(self, table_name: str) -> bool:
  31. if table_name in self.tables.keys():
  32. return True
  33. else:
  34. return False
  35. def has_column(self, table_name: str, field_name: str) -> bool:
  36. if self.has_table(table_name):
  37. if field_name in self.tables[table_name]["fields"].keys():
  38. return True
  39. else:
  40. return False
  41. else:
  42. return False
  43. def get_field_info(self, table_name: str, field_name: str) -> Dict:
  44. try:
  45. return self.tables[table_name]['fields'][field_name]
  46. except:
  47. return {}
  48. def single_table_mschema(self, table_name: str, selected_columns: List = None,
  49. example_num=3, show_type_detail=False) -> str:
  50. table_info = self.tables.get(table_name, {})
  51. output = []
  52. table_comment = table_info.get('comment', '')
  53. if table_comment is not None and table_comment != 'None' and len(table_comment) > 0:
  54. if self.schema is not None and len(self.schema) > 0:
  55. output.append(f"# Table: {self.schema}.{table_name}, {table_comment}")
  56. else:
  57. output.append(f"# Table: {table_name}, {table_comment}")
  58. else:
  59. if self.schema is not None and len(self.schema) > 0:
  60. output.append(f"# Table: {self.schema}.{table_name}")
  61. else:
  62. output.append(f"# Table: {table_name}")
  63. field_lines = []
  64. # 处理表中的每一个字段
  65. for field_name, field_info in table_info['fields'].items():
  66. if selected_columns is not None and field_name.lower() not in selected_columns:
  67. continue
  68. raw_type = self.get_field_type(field_info['type'], not show_type_detail)
  69. field_line = f"({field_name}:{raw_type.upper()}"
  70. if field_info['comment'] != '':
  71. field_line += f", {field_info['comment'].strip()}"
  72. else:
  73. pass
  74. ## 打上主键标识
  75. is_primary_key = field_info.get('primary_key', False)
  76. if is_primary_key:
  77. field_line += f", Primary Key"
  78. # 如果有示例,添加上
  79. if len(field_info.get('examples', [])) > 0 and example_num > 0:
  80. examples = field_info['examples']
  81. examples = [s for s in examples if s is not None]
  82. examples = examples_to_str(examples)
  83. if len(examples) > example_num:
  84. examples = examples[:example_num]
  85. if raw_type in ['DATE', 'TIME', 'DATETIME', 'TIMESTAMP']:
  86. examples = [examples[0]]
  87. elif len(examples) > 0 and max([len(s) for s in examples]) > 20:
  88. if max([len(s) for s in examples]) > 50:
  89. examples = []
  90. else:
  91. examples = [examples[0]]
  92. else:
  93. pass
  94. if len(examples) > 0:
  95. example_str = ', '.join([str(example) for example in examples])
  96. field_line += f", Examples: [{example_str}]"
  97. else:
  98. pass
  99. else:
  100. field_line += ""
  101. field_line += ")"
  102. field_lines.append(field_line)
  103. output.append('[')
  104. output.append(',\n'.join(field_lines))
  105. output.append(']')
  106. return '\n'.join(output)
  107. def to_mschema(self, selected_tables: List = None, selected_columns: List = None,
  108. example_num=3, show_type_detail=False) -> str:
  109. """
  110. convert to a MSchema string.
  111. selected_tables: 默认为None,表示选择所有的表
  112. selected_columns: 默认为None,表示所有列全选,格式['table_name.column_name']
  113. """
  114. output = []
  115. output.append(f"【DB_ID】 {self.db_id}")
  116. output.append(f"【Schema】")
  117. if selected_tables is not None:
  118. selected_tables = [s.lower() for s in selected_tables]
  119. if selected_columns is not None:
  120. selected_columns = [s.lower() for s in selected_columns]
  121. selected_tables = [s.split('.')[0].lower() for s in selected_columns]
  122. # 依次处理每一个表
  123. for table_name, table_info in self.tables.items():
  124. if selected_tables is None or table_name.lower() in selected_tables:
  125. cur_table_type = table_info.get('type', 'table')
  126. column_names = list(table_info['fields'].keys())
  127. if selected_columns is not None:
  128. cur_selected_columns = [c.lower() for c in column_names if f"{table_name}.{c}".lower() in selected_columns]
  129. else:
  130. cur_selected_columns = selected_columns
  131. output.append(self.single_table_mschema(table_name, cur_selected_columns, example_num, show_type_detail))
  132. # 添加外键信息,选择table_type为view时不展示外键
  133. if self.foreign_keys:
  134. output.append("【Foreign keys】")
  135. for fk in self.foreign_keys:
  136. ref_schema = fk[2]
  137. table1, column1, _, table2, column2 = fk
  138. if selected_tables is None or \
  139. (table1.lower() in selected_tables and table2.lower() in selected_tables):
  140. if ref_schema == self.schema:
  141. output.append(f"{fk[0]}.{fk[1]}={fk[3]}.{fk[4]}")
  142. return '\n'.join(output)
  143. def dump(self):
  144. schema_dict = {
  145. "db_id": self.db_id,
  146. "schema": self.schema,
  147. "tables": self.tables,
  148. "foreign_keys": self.foreign_keys
  149. }
  150. return schema_dict
  151. def save(self, file_path: str):
  152. schema_dict = self.dump()
  153. write_json(file_path, schema_dict)
  154. def load(self, file_path: str):
  155. data = read_json(file_path)
  156. self.db_id = data.get("db_id", "Anonymous")
  157. self.schema = data.get("schema", None)
  158. self.tables = data.get("tables", {})
  159. self.foreign_keys = data.get("foreign_keys", [])