| 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 diagram from JSON input. |
| |
| Args: |
| json_input (str): A JSON string describing the entity relationship diagram structure. |
| It must follow the Expected JSON Format Example below. |
| |
| Expected JSON Format Example: |
| { |
| "entities": [ |
| { |
| "name": "Customer", |
| "type": "strong", |
| "attributes": [ |
| {"name": "customer_id", "type": "primary_key"}, |
| {"name": "first_name", "type": "regular"}, |
| {"name": "last_name", "type": "regular"}, |
| {"name": "email", "type": "regular"}, |
| {"name": "phone", "type": "multivalued", "multivalued": true}, |
| {"name": "age", "type": "derived", "derived": true}, |
| {"name": "address", "type": "composite", "composite": true} |
| ] |
| }, |
| { |
| "name": "Order", |
| "type": "strong", |
| "attributes": [ |
| {"name": "order_id", "type": "primary_key"}, |
| {"name": "customer_id", "type": "foreign_key"}, |
| {"name": "order_date", "type": "regular"}, |
| {"name": "total_amount", "type": "regular"}, |
| {"name": "status", "type": "regular"}, |
| {"name": "shipping_address", "type": "composite", "composite": true} |
| ] |
| }, |
| { |
| "name": "Product", |
| "type": "strong", |
| "attributes": [ |
| {"name": "product_id", "type": "primary_key"}, |
| {"name": "name", "type": "regular"}, |
| {"name": "description", "type": "regular"}, |
| {"name": "price", "type": "regular"}, |
| {"name": "category_id", "type": "foreign_key"}, |
| {"name": "stock_quantity", "type": "regular"} |
| ] |
| }, |
| { |
| "name": "Category", |
| "type": "strong", |
| "attributes": [ |
| {"name": "category_id", "type": "primary_key"}, |
| {"name": "name", "type": "regular"}, |
| {"name": "description", "type": "regular"}, |
| {"name": "parent_category_id", "type": "foreign_key"} |
| ] |
| }, |
| { |
| "name": "OrderItem", |
| "type": "weak", |
| "attributes": [ |
| {"name": "order_id", "type": "primary_key"}, |
| {"name": "product_id", "type": "primary_key"}, |
| {"name": "quantity", "type": "regular"}, |
| {"name": "unit_price", "type": "regular"}, |
| {"name": "discount", "type": "regular"} |
| ] |
| }, |
| { |
| "name": "Payment", |
| "type": "strong", |
| "attributes": [ |
| {"name": "payment_id", "type": "primary_key"}, |
| {"name": "order_id", "type": "foreign_key"}, |
| {"name": "payment_method", "type": "regular"}, |
| {"name": "amount", "type": "regular"}, |
| {"name": "payment_date", "type": "regular"}, |
| {"name": "transaction_id", "type": "regular"} |
| ] |
| }, |
| { |
| "name": "Supplier", |
| "type": "strong", |
| "attributes": [ |
| {"name": "supplier_id", "type": "primary_key"}, |
| {"name": "company_name", "type": "regular"}, |
| {"name": "contact_person", "type": "regular"}, |
| {"name": "email", "type": "regular"}, |
| {"name": "phone", "type": "multivalued", "multivalued": true}, |
| {"name": "address", "type": "composite", "composite": true} |
| ] |
| }, |
| { |
| "name": "Review", |
| "type": "weak", |
| "attributes": [ |
| {"name": "customer_id", "type": "primary_key"}, |
| {"name": "product_id", "type": "primary_key"}, |
| {"name": "rating", "type": "regular"}, |
| {"name": "comment", "type": "regular"}, |
| {"name": "review_date", "type": "regular"} |
| ] |
| } |
| ], |
| "relationships": [ |
| { |
| "name": "Places", |
| "type": "regular", |
| "entities": ["Customer", "Order"], |
| "cardinalities": { |
| "Customer": "1", |
| "Order": "N" |
| } |
| }, |
| { |
| "name": "Contains", |
| "type": "identifying", |
| "entities": ["Order", "OrderItem"], |
| "cardinalities": { |
| "Order": "1", |
| "OrderItem": "N" |
| } |
| }, |
| { |
| "name": "Includes", |
| "type": "regular", |
| "entities": ["OrderItem", "Product"], |
| "cardinalities": { |
| "OrderItem": "N", |
| "Product": "1" |
| } |
| }, |
| { |
| "name": "BelongsTo", |
| "type": "regular", |
| "entities": ["Product", "Category"], |
| "cardinalities": { |
| "Product": "N", |
| "Category": "1" |
| } |
| }, |
| { |
| "name": "ProcessedBy", |
| "type": "regular", |
| "entities": ["Order", "Payment"], |
| "cardinalities": { |
| "Order": "1", |
| "Payment": "1" |
| } |
| }, |
| { |
| "name": "Supplies", |
| "type": "regular", |
| "entities": ["Supplier", "Product"], |
| "cardinalities": { |
| "Supplier": "N", |
| "Product": "M" |
| } |
| }, |
| { |
| "name": "Writes", |
| "type": "weak", |
| "entities": ["Customer", "Review"], |
| "cardinalities": { |
| "Customer": "1", |
| "Review": "N" |
| } |
| }, |
| { |
| "name": "ReviewsProduct", |
| "type": "weak", |
| "entities": ["Review", "Product"], |
| "cardinalities": { |
| "Review": "N", |
| "Product": "1" |
| } |
| }, |
| { |
| "name": "Subcategory", |
| "type": "regular", |
| "entities": ["Category", "Category"], |
| "cardinalities": { |
| "Category": "1", |
| "Category": "N" |
| } |
| } |
| ] |
| } |
| |
| 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='dot') |
| dot.attr(rankdir='TB', bgcolor='white', pad='0.5') |
| dot.attr('node', fontname='Arial', fontsize='10', color='black') |
| dot.attr('edge', fontname='Arial', fontsize='9', 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 |
| |
| label_parts = [entity_name] |
| |
| if attributes: |
| attr_lines = [] |
| for attr in attributes: |
| attr_name = attr.get('name', '') |
| attr_type = attr.get('type', 'regular') |
| is_multivalued = attr.get('multivalued', False) |
| is_derived = attr.get('derived', False) |
| is_composite = attr.get('composite', False) |
| |
| if attr_type == 'primary_key': |
| if is_multivalued: |
| attr_display = f"PK: {{ {attr_name} }}" |
| else: |
| attr_display = f"PK: {attr_name}" |
| elif attr_type == 'foreign_key': |
| attr_display = f"FK: {attr_name}" |
| else: |
| attr_display = attr_name |
| if is_derived: |
| attr_display = f"/ {attr_display} /" |
| if is_multivalued: |
| attr_display = f"{{ {attr_display} }}" |
| if is_composite: |
| attr_display = f"( {attr_display} )" |
| |
| attr_lines.append(attr_display) |
| |
| if attr_lines: |
| label_parts.extend(attr_lines) |
| |
| label = "\\n".join(label_parts) |
| |
| if entity_type == 'weak': |
| style = 'filled' |
| fillcolor = '#e8e8e8' |
| penwidth = '2' |
| else: |
| style = 'filled' |
| fillcolor = '#d0d0d0' |
| penwidth = '1' |
| |
| dot.node(entity_name, label, shape='box', style=style, fillcolor=fillcolor, penwidth=penwidth) |
| |
| for relationship in relationships: |
| rel_name = relationship.get('name') |
| rel_type = relationship.get('type', 'regular') |
| entities_involved = relationship.get('entities', []) |
| cardinalities = relationship.get('cardinalities', {}) |
| |
| if not rel_name or len(entities_involved) < 2: |
| continue |
| |
| if rel_type == 'identifying': |
| style = 'filled' |
| fillcolor = '#c0c0c0' |
| penwidth = '2' |
| else: |
| style = 'filled' |
| fillcolor = '#c0c0c0' |
| penwidth = '1' |
| |
| dot.node(rel_name, rel_name, shape='diamond', style=style, fillcolor=fillcolor, penwidth=penwidth) |
| |
| for entity in entities_involved: |
| cardinality = cardinalities.get(entity, '1') |
| dot.edge(entity, rel_name, label=cardinality) |
|
|
| 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)}" |