Source

grrdrr / src / grrdrr / database.py



from model import build, WithNodeLabel, WithArcWeight
from htsql import HTSQL
from htsql.introspect import introspect


def build_schema(db):
    with HTSQL(db):
        catalog = introspect()

    graph = build(WithNodeLabel, WithArcWeight)

    node_by_table = {}
    for schema in catalog.schemas:
        for table in schema.tables:
            node = graph.add_node()
            node.set_label(table.name)
            node_by_table[table] = node

    for schema in catalog.schemas:
        for table in schema.tables:
            origin = node_by_table[table]
            for foreign_key in table.foreign_keys:
                target = node_by_table[foreign_key.target]
                is_parental = (foreign_key.target.primary_key is not None and
                               all(column in table.columns
                                   for column in foreign_key.origin_columns))
                if is_parental:
                    weight = 1.0
                else:
                    weight = 0.5
                arc = graph.add_arc(origin, target)
                arc.set_weight(weight)

    graph.freeze()
    return graph