Graphify / entity_relationship_generator.py
ZahirJS's picture
Update entity_relationship_generator.py
a025aee verified
raw
history blame
19 kB
import graphviz
import json
from tempfile import NamedTemporaryFile
import os
def generate_entity_relationship_diagram(json_input: str, output_format: str) -> str:
"""
Generates an Entity Relationship (ER) diagram from JSON input.
Args:
json_input (str): A JSON string describing the ER diagram structure.
It must follow the Expected JSON Format Example below.
Expected JSON Format Example:
{
"entities": [
{
"name": "User",
"type": "strong",
"attributes": [
{
"name": "user_id",
"type": "primary_key"
},
{
"name": "username",
"type": "regular"
},
{
"name": "email",
"type": "regular"
},
{
"name": "password_hash",
"type": "regular"
},
{
"name": "full_name",
"type": "composite"
},
{
"name": "phone_numbers",
"type": "multivalued"
},
{
"name": "age",
"type": "derived"
}
]
},
{
"name": "Product",
"type": "strong",
"attributes": [
{
"name": "product_id",
"type": "primary_key"
},
{
"name": "name",
"type": "regular"
},
{
"name": "description",
"type": "regular"
},
{
"name": "price",
"type": "regular"
},
{
"name": "stock_quantity",
"type": "regular"
},
{
"name": "tags",
"type": "multivalued"
}
]
},
{
"name": "Category",
"type": "strong",
"attributes": [
{
"name": "category_id",
"type": "primary_key"
},
{
"name": "name",
"type": "regular"
},
{
"name": "description",
"type": "regular"
}
]
},
{
"name": "Order",
"type": "strong",
"attributes": [
{
"name": "order_id",
"type": "primary_key"
},
{
"name": "order_date",
"type": "regular"
},
{
"name": "status",
"type": "regular"
},
{
"name": "total_amount",
"type": "derived"
},
{
"name": "shipping_address",
"type": "composite"
}
]
},
{
"name": "OrderItem",
"type": "weak",
"attributes": [
{
"name": "line_number",
"type": "partial_key"
},
{
"name": "quantity",
"type": "regular"
},
{
"name": "unit_price",
"type": "regular"
},
{
"name": "subtotal",
"type": "derived"
}
]
},
{
"name": "Payment",
"type": "strong",
"attributes": [
{
"name": "payment_id",
"type": "primary_key"
},
{
"name": "amount",
"type": "regular"
},
{
"name": "payment_method",
"type": "regular"
},
{
"name": "payment_date",
"type": "regular"
},
{
"name": "status",
"type": "regular"
}
]
},
{
"name": "Review",
"type": "strong",
"attributes": [
{
"name": "review_id",
"type": "primary_key"
},
{
"name": "rating",
"type": "regular"
},
{
"name": "comment",
"type": "regular"
},
{
"name": "review_date",
"type": "regular"
}
]
},
{
"name": "Vendor",
"type": "strong",
"attributes": [
{
"name": "vendor_id",
"type": "primary_key"
},
{
"name": "company_name",
"type": "regular"
},
{
"name": "contact_person",
"type": "regular"
},
{
"name": "contact_emails",
"type": "multivalued"
},
{
"name": "business_address",
"type": "composite"
}
]
},
{
"name": "ShoppingCart",
"type": "strong",
"attributes": [
{
"name": "cart_id",
"type": "primary_key"
},
{
"name": "created_date",
"type": "regular"
},
{
"name": "last_updated",
"type": "regular"
},
{
"name": "total_items",
"type": "derived"
}
]
},
{
"name": "CartItem",
"type": "weak",
"attributes": [
{
"name": "item_position",
"type": "partial_key"
},
{
"name": "quantity",
"type": "regular"
},
{
"name": "added_date",
"type": "regular"
}
]
}
],
"relationships": [
{
"name": "PlacesOrder",
"type": "regular",
"entities": ["User", "Order"],
"cardinalities": {
"User": "1",
"Order": "M"
},
"attributes": []
},
{
"name": "Contains",
"type": "identifying",
"entities": ["Order", "OrderItem"],
"cardinalities": {
"Order": "1",
"OrderItem": "M"
},
"attributes": []
},
{
"name": "OrdersProduct",
"type": "regular",
"entities": ["OrderItem", "Product"],
"cardinalities": {
"OrderItem": "M",
"Product": "1"
},
"attributes": []
},
{
"name": "BelongsTo",
"type": "regular",
"entities": ["Product", "Category"],
"cardinalities": {
"Product": "M",
"Category": "1"
},
"attributes": []
},
{
"name": "ProcessesPayment",
"type": "regular",
"entities": ["Order", "Payment"],
"cardinalities": {
"Order": "1",
"Payment": "M"
},
"attributes": []
},
{
"name": "WritesReview",
"type": "regular",
"entities": ["User", "Review"],
"cardinalities": {
"User": "1",
"Review": "M"
},
"attributes": []
},
{
"name": "ReviewsProduct",
"type": "regular",
"entities": ["Review", "Product"],
"cardinalities": {
"Review": "M",
"Product": "1"
},
"attributes": []
},
{
"name": "Supplies",
"type": "regular",
"entities": ["Vendor", "Product"],
"cardinalities": {
"Vendor": "M",
"Product": "M"
},
"attributes": [
{
"name": "supply_price"
},
{
"name": "lead_time"
}
]
},
{
"name": "HasCart",
"type": "regular",
"entities": ["User", "ShoppingCart"],
"cardinalities": {
"User": "1",
"ShoppingCart": "1"
},
"attributes": []
},
{
"name": "CartContains",
"type": "identifying",
"entities": ["ShoppingCart", "CartItem"],
"cardinalities": {
"ShoppingCart": "1",
"CartItem": "M"
},
"attributes": []
},
{
"name": "CartHasProduct",
"type": "regular",
"entities": ["CartItem", "Product"],
"cardinalities": {
"CartItem": "M",
"Product": "1"
},
"attributes": []
}
]
}
Returns:
str: The filepath to the generated PNG image file.
"""
try:
if not json_input.strip():
return "Error: Empty input"
data = json.loads(json_input)
if 'entities' not in data:
raise ValueError("Missing required field: entities")
dot = graphviz.Graph(comment='ER Diagram', engine='neato')
dot.attr(
bgcolor='white',
pad='1.5',
overlap='false',
splines='true',
sep='+25',
esep='+15'
)
dot.attr('node', fontname='Arial', fontsize='10', color='#404040')
dot.attr('edge', fontname='Arial', fontsize='9', color='#4a4a4a')
entity_color = '#BEBEBE'
attribute_color = '#D4D4D4'
relationship_color = '#E8E8E8'
isa_color = '#F0F0F0'
font_color = 'black'
entities = data.get('entities', [])
relationships = data.get('relationships', [])
# Process entities with new styling
for entity in entities:
entity_name = entity.get('name')
entity_type = entity.get('type', 'strong')
attributes = entity.get('attributes', [])
if not entity_name:
continue
if entity_type == 'weak':
dot.node(
entity_name,
entity_name,
shape='box',
style='filled,rounded',
fillcolor=entity_color,
fontcolor=font_color,
color='#404040',
penwidth='3',
width='1.8',
height='0.8',
fontsize='12'
)
else:
dot.node(
entity_name,
entity_name,
shape='box',
style='filled,rounded',
fillcolor=entity_color,
fontcolor=font_color,
color='#404040',
penwidth='1',
width='1.8',
height='0.8',
fontsize='12'
)
for i, attr in enumerate(attributes):
attr_name = attr.get('name', '')
attr_type = attr.get('type', 'regular')
attr_id = f"{entity_name}_attr_{i}"
if attr_type == 'primary_key':
dot.node(
attr_id,
f'{attr_name} (PK)',
shape='ellipse',
style='filled,rounded',
fillcolor=attribute_color,
fontcolor=font_color,
color='#404040',
width='1.2',
height='0.6',
fontsize='10'
)
elif attr_type == 'partial_key':
dot.node(
attr_id,
f'{attr_name} (Partial)',
shape='ellipse',
style='filled,rounded,dashed',
fillcolor=attribute_color,
fontcolor=font_color,
color='#404040',
width='1.2',
height='0.6',
fontsize='10'
)
elif attr_type == 'multivalued':
dot.node(
attr_id,
attr_name,
shape='ellipse',
style='filled,rounded',
fillcolor=attribute_color,
fontcolor=font_color,
color='#404040',
penwidth='3',
width='1.2',
height='0.6',
fontsize='10'
)
elif attr_type == 'derived':
dot.node(
attr_id,
f'/{attr_name}/',
shape='ellipse',
style='filled,rounded,dashed',
fillcolor=attribute_color,
fontcolor=font_color,
color='#404040',
width='1.2',
height='0.6',
fontsize='10'
)
elif attr_type == 'composite':
dot.node(
attr_id,
attr_name,
shape='ellipse',
style='filled,rounded',
fillcolor=attribute_color,
fontcolor=font_color,
color='#404040',
width='1.2',
height='0.6',
fontsize='10'
)
else:
dot.node(
attr_id,
attr_name,
shape='ellipse',
style='filled,rounded',
fillcolor=attribute_color,
fontcolor=font_color,
color='#404040',
width='1.2',
height='0.6',
fontsize='10'
)
dot.edge(entity_name, attr_id, color='#4a4a4a', len='1.5')
for relationship in relationships:
rel_name = relationship.get('name')
rel_type = relationship.get('type', 'regular')
entities_involved = relationship.get('entities', [])
cardinalities = relationship.get('cardinalities', {})
rel_attributes = relationship.get('attributes', [])
if not rel_name:
continue
if rel_type == 'isa':
parent = relationship.get('parent')
children = relationship.get('children', [])
if parent and children:
isa_id = f"isa_{rel_name}"
dot.node(
isa_id,
'ISA',
shape='triangle',
style='filled,rounded',
fillcolor=isa_color,
fontcolor=font_color,
color='#404040',
penwidth='2',
width='1.0',
height='0.8',
fontsize='10'
)
dot.edge(parent, isa_id, color='#4a4a4a', len='2.0')
for child in children:
dot.edge(isa_id, child, color='#4a4a4a', len='2.0')
elif len(entities_involved) >= 2:
if rel_type == 'identifying':
dot.node(
rel_name,
rel_name,
shape='diamond',
style='filled,rounded',
fillcolor=relationship_color,
fontcolor=font_color,
color='#404040',
penwidth='3',
width='1.8',
height='1.0',
fontsize='11'
)
else:
dot.node(
rel_name,
rel_name,
shape='diamond',
style='filled,rounded',
fillcolor=relationship_color,
fontcolor=font_color,
color='#404040',
penwidth='1',
width='1.8',
height='1.0',
fontsize='11'
)
for j, attr in enumerate(rel_attributes):
attr_name = attr.get('name', '')
attr_id = f"{rel_name}_attr_{j}"
dot.node(
attr_id,
attr_name,
shape='ellipse',
style='filled,rounded',
fillcolor=attribute_color,
fontcolor=font_color,
color='#404040',
width='1.0',
height='0.5',
fontsize='9'
)
dot.edge(rel_name, attr_id, color='#4a4a4a', len='1.0')
for entity in entities_involved:
cardinality = cardinalities.get(entity, '1')
dot.edge(
entity,
rel_name,
label=f' {cardinality} ',
color='#4a4a4a',
len='2.5',
fontcolor='#4a4a4a',
fontsize='10'
)
with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
dot.render(tmp.name, format=output_format, cleanup=True)
return f"{tmp.name}.{output_format}"
except json.JSONDecodeError:
return "Error: Invalid JSON format"
except Exception as e:
return f"Error: {str(e)}"