Graphify / entity_relationship_generator.py
ZahirJS's picture
Update entity_relationship_generator.py
c5a9b0f verified
raw
history blame
34.3 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)}"
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"
}
]
}
],
"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": []
}
]
}
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 = '#B8D4F1'
relationship_color = '#FFF9C4'
isa_color = '#A8E6CF'
font_color = 'black'
entities = data.get('entities', [])
relationships = data.get('relationships', [])
for entity in entities:
entity_name = entity.get('name')
entity_type = entity.get('type', 'strong')
attributes = entity.get('attributes', [])
if not entity_name:
continue
entity_color = '#BEBEBE'
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}"
attr_color = attribute_color
if attr_type == 'primary_key':
dot.node(
attr_id,
f'{attr_name} (PK)',
shape='ellipse',
style='filled,rounded',
fillcolor=attr_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=attr_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=attr_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=attr_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=attr_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=attr_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}"
isa_color = isa_color
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:
rel_color = relationship_color
if rel_type == 'identifying':
dot.node(
rel_name,
rel_name,
shape='diamond',
style='filled,rounded',
fillcolor=rel_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=rel_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}"
attr_color = attribute_color
dot.node(
attr_id,
attr_name,
shape='ellipse',
style='filled,rounded',
fillcolor=attr_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)}"