import graphviz import json from tempfile import NamedTemporaryFile import os def generate_entity_relationship_diagram(json_input: str, output_format: str) -> str: 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)}"