import json, os from typing import Any, Dict, Iterable, List, Optional, Tuple, Union from providers.base import MetadataProvider from .utils import examples_to_str from .m_schema import MSchema class SchemaEngine: def __init__(self, schema: Optional[str] = None, metadata_provider: MetadataProvider = None, ignore_tables: Optional[List[str]] = None, include_tables: Optional[List[str]] = None, sample_rows_in_table_info: int = 3, indexes_in_table_info: bool = False, custom_table_info: Optional[dict] = None, view_support: bool = False, max_string_length: int = 300, mschema: Optional[MSchema] = None, db_name: Optional[str] = ''): self._db_name = db_name self._inspector = metadata_provider self._dialect = metadata_provider.get_type() # Dictionary to store table names and their corresponding schema self._tables_schemas: Dict[str, str] = {} # If a schema is specified, filter by that schema and store that value for every table. if schema: # self._usable_tables = [ # table_name for table_name in self._usable_tables # if self._inspector.has_table(table_name, schema) # ] # for table_name in self._usable_tables: # self._tables_schemas[table_name] = schema all_tables = [] tables = self._inspector.get_tables(database=schema) all_tables.extend(tables) for table in tables: self._tables_schemas[table] = schema self._usable_tables = all_tables else: all_tables = [] # Iterate through all available schemas for s in self.get_schema_names(): tables = self._inspector.get_tables(database=s) all_tables.extend(tables) for table in tables: self._tables_schemas[table] = s self._usable_tables = all_tables if mschema is not None: self._mschema = mschema else: self._mschema = MSchema(db_id=db_name, schema=None) self.init_mschema() @property def mschema(self) -> MSchema: """Return M-Schema""" return self._mschema def get_pk_constraint(self, table_name: str) -> Dict: return self._inspector.get_pk_constraint(table_name, self._tables_schemas[table_name] ) def get_table_comment(self, table_name: str): try: return self._inspector.get_table_comment(table_name, self._tables_schemas[table_name]) except: # sqlite does not support comments return '' def default_schema_name(self) -> Optional[str]: return self._inspector.default_schema_name def get_schema_names(self) -> List[str]: return self._inspector.get_databases() def get_foreign_keys(self, table_name: str): return self._inspector.get_foreign_keys(table_name, self._tables_schemas[table_name]) def get_unique_constraints(self, table_name: str): return self._inspector.get_unique_constraints(table_name, self._tables_schemas[table_name]) def fectch_distinct_values(self, column_name: str, table_name: str, max_num: int = 5): return self._inspector.fectch_distinct_values(column_name, table_name, self._tables_schemas[table_name], max_num) def init_mschema(self): for table_name in self._usable_tables: table_comment = self.get_table_comment(table_name) table_comment = '' if table_comment is None else table_comment.strip() table_with_schema = self._tables_schemas[table_name] + '.' + table_name self._mschema.add_table(table_with_schema, fields={}, comment=table_comment) pks = self.get_pk_constraint(table_name) fks = self.get_foreign_keys(table_name) for fk in fks: referred_schema = fk['referred_schema'] for c, r in zip(fk['constrained_columns'], fk['referred_columns']): self._mschema.add_foreign_key(table_with_schema, c, referred_schema, fk['referred_table'], r) fields = self._inspector.get_columns(table=table_name, database=self._tables_schemas[table_name]) for field in fields: field_type = f"{field['type']!s}" field_name = field['name'] primary_key = field_name in pks field_comment = field.get("comment", None) field_comment = "" if field_comment is None else field_comment.strip() autoincrement = field.get('autoincrement', False) default = field.get('default', None) nullable = field.get('nullable', False) if default is not None: default = f'{default}' try: examples = self.fectch_distinct_values(field_name, table_name, 5) except: examples = [] examples = examples_to_str(examples) self._mschema.add_field( table_with_schema, field_name, field_type=field_type, primary_key=primary_key, nullable=nullable, default=default, autoincrement=autoincrement, comment=field_comment, examples=examples )