-
Notifications
You must be signed in to change notification settings - Fork 0
/
graph.py
60 lines (49 loc) · 2.05 KB
/
graph.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
from collections import namedtuple
from networkx import Graph
from sqlalchemy import inspect, MetaData
from sqlalchemy.ext.automap import automap_base
from sqlalchemy.orm import Session
Node = namedtuple("Node", ["table", "primary_key"])
def get_graph(engine, table, primary_key):
"""Construct the graph for a specified data-point
Args:
engine: An sql-alchemy engine instance, used to connect to the database.
table: Name of the table.
primary_key: The primary key for the row.
Returns:
A graph of relations for the row.
"""
metadata = MetaData()
metadata.reflect(engine)
Base = automap_base(metadata=metadata)
Base.prepare()
_table = Base.classes[table]
graph = Graph()
with Session(engine) as session:
row = session.query(_table).get(primary_key)
row_node = Node(
table=_get_table_name_from_row(row),
primary_key=_get_primary_key_from_row(row),
)
graph.add_node(row_node)
relationships = row.__mapper__.relationships
for relationship in relationships:
# This is a bit hacky - but they don't call it a hackathon for nothing.
relationship_name = str(relationship).split(".")[-1]
related_rows = getattr(row, relationship_name)
for related_row in related_rows:
related_row_node = Node(
table=_get_table_name_from_row(related_row),
primary_key=_get_primary_key_from_row(related_row),
)
graph.add_node(related_row_node)
graph.add_edge(row_node, related_row_node)
return graph
def _get_table_name_from_row(row):
return row.__table__.name
def _get_primary_key_from_row(row):
primary_key_columns = row.__mapper__.primary_key
primary_key_values = [getattr(row, column.name) for column in primary_key_columns]
if len(primary_key_values) != 1:
raise NotImplementedError("We just consider cases with single column pk for the time being")
return primary_key_values[0]