| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122 |
- 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
- )
|