Added Visualization

src/streamlit_app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import re
 from io import BytesIO
 import base64
+import graphviz
 
 # Page configuration
 st.set_page_config(
@@ -113,49 +114,122 @@ class DBMLParser:
             })
 
 
-class MermaidGenerator:
-    """Generate Mermaid diagram from parsed DBML"""
+class GraphvizGenerator:
+    """Generate Graphviz diagram from parsed DBML"""
+    
+    # Color scheme
+    COLORS = {
+        'table_bg': '#E8F4F8',
+        'table_border': '#2E86AB',
+        'pk_bg': '#FFE5B4',
+        'pk_text': '#8B4513',
+        'header_bg': '#2E86AB',
+        'header_text': '#FFFFFF',
+        'text': '#333333',
+        'arrow': '#555555'
+    }
     
     def __init__(self, tables, relationships):
         self.tables = tables
         self.relationships = relationships
     
     def generate(self):
-        """Generate Mermaid ER diagram"""
-        lines = ['erDiagram']
+        """Generate Graphviz diagram"""
+        dot = graphviz.Digraph(comment='Database Schema')
+        dot.attr(rankdir='LR', bgcolor='white', splines='ortho', nodesep='1', ranksep='1.5')
+        dot.attr('node', shape='plaintext')
+        dot.attr('edge', color=self.COLORS['arrow'], penwidth='2')
         
-        # Add tables with columns
+        # Add tables
         for table_name, columns in self.tables.items():
-            for col in columns:
-                col_def = f"{col['type']}"
-                
-                # Add constraints
-                constraints = []
-                if col['is_pk']:
-                    constraints.append('PK')
-                if col['is_unique']:
-                    constraints.append('UK')
-                if col['is_not_null']:
-                    constraints.append('NOT NULL')
-                
-                constraint_str = f" \"{','.join(constraints)}\"" if constraints else ""
-                
-                lines.append(f"    {table_name} {{")
-                lines.append(f"        {col_def} {col['name']}{constraint_str}")
-                lines.append(f"    }}")
+            html_label = self._create_table_html(table_name, columns)
+            dot.node(table_name, label=f'<{html_label}>')
         
-        # Add relationships
+        # Add relationships with proper arrows
         for rel in self.relationships:
             if rel['type'] == 'one-to-many':
-                rel_symbol = '||--o{'
+                # One to many: arrow from parent to child
+                dot.edge(
+                    rel['to_table'], 
+                    rel['from_table'],
+                    label=f" {rel['to_column']} → {rel['from_column']} ",
+                    arrowhead='crow',
+                    fontsize='10',
+                    fontcolor='#666666'
+                )
             elif rel['type'] == 'many-to-one':
-                rel_symbol = '}o--||'
+                # Many to one: arrow from child to parent
+                dot.edge(
+                    rel['from_table'],
+                    rel['to_table'],
+                    label=f" {rel['from_column']} → {rel['to_column']} ",
+                    arrowhead='normal',
+                    fontsize='10',
+                    fontcolor='#666666'
+                )
             else:
-                rel_symbol = '||--||'
+                # One to one
+                dot.edge(
+                    rel['from_table'],
+                    rel['to_table'],
+                    label=f" {rel['from_column']} ↔ {rel['to_column']} ",
+                    arrowhead='none',
+                    dir='both',
+                    fontsize='10',
+                    fontcolor='#666666'
+                )
+        
+        return dot
+    
+    def _create_table_html(self, table_name, columns):
+        """Create HTML table for Graphviz node"""
+        html = f'''<
+        <TABLE BORDER="2" CELLBORDER="1" CELLSPACING="0" CELLPADDING="8" BGCOLOR="{self.COLORS['table_bg']}" COLOR="{self.COLORS['table_border']}">
+            <TR>
+                <TD COLSPAN="2" BGCOLOR="{self.COLORS['header_bg']}" ALIGN="CENTER">
+                    <FONT COLOR="{self.COLORS['header_text']}" POINT-SIZE="14"><B>{table_name}</B></FONT>
+                </TD>
+            </TR>
+        '''
+        
+        for col in columns:
+            # Determine background color
+            bg_color = self.COLORS['pk_bg'] if col['is_pk'] else self.COLORS['table_bg']
+            
+            # Build column name with constraints
+            col_name = col['name']
+            if col['is_pk']:
+                col_name = f"🔑 {col_name}"
+            
+            # Build type with constraints
+            type_str = col['type']
+            constraints = []
+            if col['is_pk']:
+                constraints.append('PK')
+            if col['is_unique']:
+                constraints.append('UQ')
+            if col['is_not_null']:
+                constraints.append('NOT NULL')
             
-            lines.append(f"    {rel['from_table']} {rel_symbol} {rel['to_table']} : \"{rel['from_column']}-{rel['to_column']}\"")
+            if constraints:
+                type_str += f" [{', '.join(constraints)}]"
+            
+            html += f'''
+            <TR>
+                <TD ALIGN="LEFT" BGCOLOR="{bg_color}">
+                    <FONT COLOR="{self.COLORS['text']}" POINT-SIZE="11"><B>{col_name}</B></FONT>
+                </TD>
+                <TD ALIGN="LEFT" BGCOLOR="{bg_color}">
+                    <FONT COLOR="{self.COLORS['text']}" POINT-SIZE="10">{type_str}</FONT>
+                </TD>
+            </TR>
+            '''
+        
+        html += '''
+        </TABLE>
+        >'''
         
-        return '\n'.join(lines)
+        return html
 
 
 class PostgreSQLGenerator:
@@ -252,7 +326,7 @@ def main():
         st.write("""
         This tool allows you to:
         - 📊 Visualize DBML schemas
-        - 💾 Download diagram as image
+        - 💾 Download diagram as PNG/SVG
         - 🔧 Generate PostgreSQL DDL
         - 📋 Copy SQL to clipboard
         """)
@@ -260,6 +334,10 @@ def main():
         st.header("💡 Example DBML")
         if st.button("Load Example"):
             st.session_state.example_loaded = True
+        
+        st.header("⚙️ Diagram Settings")
+        image_format = st.selectbox("Download Format", ["PNG", "SVG"], index=0)
+        st.session_state.image_format = image_format.lower()
     
     # Example DBML
     example_dbml = """Table users {
@@ -306,44 +384,74 @@ Ref: comments.user_id > users.id"""
         if st.button("🎨 Generate Visualization & SQL", type="primary", use_container_width=True):
             if dbml_input.strip():
                 try:
-                    # Parse DBML
-                    parser = DBMLParser(dbml_input)
-                    tables, relationships = parser.parse()
-                    
-                    if not tables:
-                        st.error("❌ No tables found in DBML code. Please check your syntax.")
-                    else:
-                        # Generate Mermaid diagram
-                        mermaid_gen = MermaidGenerator(tables, relationships)
-                        mermaid_code = mermaid_gen.generate()
-                        
-                        # Generate PostgreSQL
-                        sql_gen = PostgreSQLGenerator(tables, relationships)
-                        sql_code = sql_gen.generate()
-                        
-                        # Store in session state
-                        st.session_state.mermaid_code = mermaid_code
-                        st.session_state.sql_code = sql_code
-                        st.session_state.tables = tables
-                        st.session_state.relationships = relationships
+                    with st.spinner("Parsing DBML and generating visualization..."):
+                        # Parse DBML
+                        parser = DBMLParser(dbml_input)
+                        tables, relationships = parser.parse()
                         
-                        st.success(f"✅ Successfully parsed {len(tables)} tables and {len(relationships)} relationships!")
+                        if not tables:
+                            st.error("❌ No tables found in DBML code. Please check your syntax.")
+                        else:
+                            # Generate Graphviz diagram
+                            graph_gen = GraphvizGenerator(tables, relationships)
+                            dot = graph_gen.generate()
+                            
+                            # Generate PostgreSQL
+                            sql_gen = PostgreSQLGenerator(tables, relationships)
+                            sql_code = sql_gen.generate()
+                            
+                            # Store in session state
+                            st.session_state.dot = dot
+                            st.session_state.sql_code = sql_code
+                            st.session_state.tables = tables
+                            st.session_state.relationships = relationships
+                            
+                            st.success(f"✅ Successfully parsed {len(tables)} tables and {len(relationships)} relationships!")
                 
                 except Exception as e:
                     st.error(f"❌ Error parsing DBML: {str(e)}")
+                    import traceback
+                    st.code(traceback.format_exc())
             else:
                 st.warning("⚠️ Please enter some DBML code first.")
     
     with col2:
-        if 'mermaid_code' in st.session_state:
+        if 'dot' in st.session_state:
             st.subheader("📊 Database Diagram")
             
-            # Display Mermaid diagram
-            st.markdown(f"""
-            ```mermaid
-            {st.session_state.mermaid_code}
-            ```
-            """)
+            # Display the diagram
+            st.graphviz_chart(st.session_state.dot)
+            
+            # Download buttons
+            col_btn1, col_btn2 = st.columns(2)
+            
+            with col_btn1:
+                # PNG download
+                try:
+                    png_data = st.session_state.dot.pipe(format='png')
+                    st.download_button(
+                        label="📥 Download PNG",
+                        data=png_data,
+                        file_name="database_schema.png",
+                        mime="image/png",
+                        use_container_width=True
+                    )
+                except:
+                    st.warning("PNG export requires Graphviz installation")
+            
+            with col_btn2:
+                # SVG download
+                try:
+                    svg_data = st.session_state.dot.pipe(format='svg')
+                    st.download_button(
+                        label="📥 Download SVG",
+                        data=svg_data,
+                        file_name="database_schema.svg",
+                        mime="image/svg+xml",
+                        use_container_width=True
+                    )
+                except:
+                    st.warning("SVG export requires Graphviz installation")
             
             # Info about tables
             st.info(f"📋 **{len(st.session_state.tables)}** tables, **{len(st.session_state.relationships)}** relationships")
@@ -370,8 +478,8 @@ Ref: comments.user_id > users.id"""
                 use_container_width=True
             )
             
-            if st.button("📋 Copy to Clipboard", use_container_width=True):
-                st.toast("SQL copied to clipboard! (Use Ctrl+C to copy the code block)", icon="✅")
+            if st.button("📋 Copy SQL", use_container_width=True):
+                st.toast("Copy the SQL code from the code block above!", icon="📋")
 
 
 if __name__ == "__main__":