schema_engine.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122
  1. import json, os
  2. from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
  3. from providers.base import MetadataProvider
  4. from .utils import examples_to_str
  5. from .m_schema import MSchema
  6. class SchemaEngine:
  7. def __init__(self, schema: Optional[str] = None, metadata_provider: MetadataProvider = None,
  8. ignore_tables: Optional[List[str]] = None, include_tables: Optional[List[str]] = None,
  9. sample_rows_in_table_info: int = 3, indexes_in_table_info: bool = False,
  10. custom_table_info: Optional[dict] = None, view_support: bool = False, max_string_length: int = 300,
  11. mschema: Optional[MSchema] = None, db_name: Optional[str] = ''):
  12. self._db_name = db_name
  13. self._inspector = metadata_provider
  14. self._dialect = metadata_provider.get_type()
  15. # Dictionary to store table names and their corresponding schema
  16. self._tables_schemas: Dict[str, str] = {}
  17. # If a schema is specified, filter by that schema and store that value for every table.
  18. if schema:
  19. # self._usable_tables = [
  20. # table_name for table_name in self._usable_tables
  21. # if self._inspector.has_table(table_name, schema)
  22. # ]
  23. # for table_name in self._usable_tables:
  24. # self._tables_schemas[table_name] = schema
  25. all_tables = []
  26. tables = self._inspector.get_tables(database=schema)
  27. all_tables.extend(tables)
  28. for table in tables:
  29. self._tables_schemas[table] = schema
  30. self._usable_tables = all_tables
  31. else:
  32. all_tables = []
  33. # Iterate through all available schemas
  34. for s in self.get_schema_names():
  35. tables = self._inspector.get_tables(database=s)
  36. all_tables.extend(tables)
  37. for table in tables:
  38. self._tables_schemas[table] = s
  39. self._usable_tables = all_tables
  40. if mschema is not None:
  41. self._mschema = mschema
  42. else:
  43. self._mschema = MSchema(db_id=db_name, schema=None)
  44. self.init_mschema()
  45. @property
  46. def mschema(self) -> MSchema:
  47. """Return M-Schema"""
  48. return self._mschema
  49. def get_pk_constraint(self, table_name: str) -> Dict:
  50. return self._inspector.get_pk_constraint(table_name, self._tables_schemas[table_name] )
  51. def get_table_comment(self, table_name: str):
  52. try:
  53. return self._inspector.get_table_comment(table_name, self._tables_schemas[table_name])
  54. except: # sqlite does not support comments
  55. return ''
  56. def default_schema_name(self) -> Optional[str]:
  57. return self._inspector.default_schema_name
  58. def get_schema_names(self) -> List[str]:
  59. return self._inspector.get_databases()
  60. def get_foreign_keys(self, table_name: str):
  61. return self._inspector.get_foreign_keys(table_name, self._tables_schemas[table_name])
  62. def get_unique_constraints(self, table_name: str):
  63. return self._inspector.get_unique_constraints(table_name, self._tables_schemas[table_name])
  64. def fectch_distinct_values(self, column_name: str, table_name: str, max_num: int = 5):
  65. return self._inspector.fectch_distinct_values(column_name, table_name, self._tables_schemas[table_name], max_num)
  66. def init_mschema(self):
  67. for table_name in self._usable_tables:
  68. table_comment = self.get_table_comment(table_name)
  69. table_comment = '' if table_comment is None else table_comment.strip()
  70. table_with_schema = self._tables_schemas[table_name] + '.' + table_name
  71. self._mschema.add_table(table_with_schema, fields={}, comment=table_comment)
  72. pks = self.get_pk_constraint(table_name)
  73. fks = self.get_foreign_keys(table_name)
  74. for fk in fks:
  75. referred_schema = fk['referred_schema']
  76. for c, r in zip(fk['constrained_columns'], fk['referred_columns']):
  77. self._mschema.add_foreign_key(table_with_schema, c, referred_schema, fk['referred_table'], r)
  78. fields = self._inspector.get_columns(table=table_name, database=self._tables_schemas[table_name])
  79. for field in fields:
  80. field_type = f"{field['type']!s}"
  81. field_name = field['name']
  82. primary_key = field_name in pks
  83. field_comment = field.get("comment", None)
  84. field_comment = "" if field_comment is None else field_comment.strip()
  85. autoincrement = field.get('autoincrement', False)
  86. default = field.get('default', None)
  87. nullable = field.get('nullable', False)
  88. if default is not None:
  89. default = f'{default}'
  90. try:
  91. examples = self.fectch_distinct_values(field_name, table_name, 5)
  92. except:
  93. examples = []
  94. examples = examples_to_str(examples)
  95. self._mschema.add_field(
  96. table_with_schema, field_name, field_type=field_type, primary_key=primary_key,
  97. nullable=nullable, default=default, autoincrement=autoincrement,
  98. comment=field_comment, examples=examples
  99. )