entity_relationship_generator.py

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(
            name='ERDiagram',
            format='png',
            graph_attr={
                'rankdir': 'TB',
                'splines': 'ortho',
                'bgcolor': 'white',
                'pad': '0.5',
                'nodesep': '2.0',
                'ranksep': '2.5'
            }
        )
        
        base_color = '#19191a'
        lightening_factor = 0.15
        
        entities = data.get('entities', [])
        relationships = data.get('relationships', [])
        
        for i, entity in enumerate(entities):
            entity_name = entity.get('name')
            entity_type = entity.get('type', 'strong')
            attributes = entity.get('attributes', [])
            
            if not entity_name:
                raise ValueError(f"Invalid entity: {entity}")
            
            current_depth = i % 6
            
            if not isinstance(base_color, str) or not base_color.startswith('#') or len(base_color) != 7:
                base_color_safe = '#19191a'
            else:
                base_color_safe = base_color
                
            base_r = int(base_color_safe[1:3], 16)
            base_g = int(base_color_safe[3:5], 16)
            base_b = int(base_color_safe[5:7], 16)
            
            current_r = base_r + int((255 - base_r) * current_depth * lightening_factor)
            current_g = base_g + int((255 - base_g) * current_depth * lightening_factor)
            current_b = base_b + int((255 - base_b) * current_depth * lightening_factor)
            
            current_r = min(255, current_r)
            current_g = min(255, current_g)
            current_b = min(255, current_b)
            
            node_color = f'#{current_r:02x}{current_g:02x}{current_b:02x}'
            font_color = 'white' if current_depth * lightening_factor < 0.6 else 'black'
            
            entity_label = f"{entity_name}\\n"
            
            if attributes:
                primary_keys = []
                foreign_keys = []
                regular_attrs = []
                
                for attr in attributes:
                    attr_name = attr.get('name', '')
                    attr_type = attr.get('type', 'key')
                    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:
                            primary_keys.append(f"{{{{ {attr_name} }}}}")
                        else:
                            primary_keys.append(f"[PK] {attr_name}")
                    elif attr_type == 'foreign_key':
                        foreign_keys.append(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} )"
                        regular_attrs.append(attr_display)
                
                if primary_keys:
                    entity_label += "\\n".join(primary_keys) + "\\n"
                if foreign_keys:
                    entity_label += "\\n".join(foreign_keys) + "\\n"
                if regular_attrs:
                    entity_label += "\\n".join(regular_attrs)
            
            if entity_type == 'weak':
                shape = 'doubleoctagon'
                style = 'filled'
            else:
                shape = 'box'
                style = 'filled,rounded'
            
            dot.node(
                entity_name,
                entity_label,
                shape=shape,
                style=style,
                fillcolor=node_color,
                fontcolor=font_color,
                fontsize='10',
                fontname='Helvetica'
            )
        
        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:
                raise ValueError(f"Invalid relationship: {relationship}")
            
            rel_node_color = '#e6f3ff'
            
            if rel_type == 'identifying':
                rel_shape = 'diamond'
                rel_style = 'filled,bold'
                rel_color = '#4a90e2'
            elif rel_type == 'weak':
                rel_shape = 'diamond'
                rel_style = 'filled,dashed'
                rel_color = '#a0a0a0'
            else:
                rel_shape = 'diamond'
                rel_style = 'filled'
                rel_color = '#4a90e2'
            
            dot.node(
                rel_name,
                rel_name,
                shape=rel_shape,
                style=rel_style,
                fillcolor=rel_color,
                fontcolor='white',
                fontsize='10',
                fontname='Helvetica'
            )
            
            for entity in entities_involved:
                cardinality = cardinalities.get(entity, '1')
                
                edge_label = cardinality
                if cardinality in ['1:1', '1:N', 'M:N', '1', 'N', 'M']:
                    pass
                else:
                    edge_label = cardinality
                
                dot.edge(
                    entity,
                    rel_name,
                    label=edge_label,
                    color='#4a4a4a',
                    fontsize='9',
                    fontcolor='#4a4a4a'
                )

        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)}"