Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,378 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import re
+from io import BytesIO
+import base64
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# Page configuration
+st.set_page_config(
+    page_title="DBML Visualizer",
+    page_icon="🗄️",
+    layout="wide"
+)
+
+# Custom CSS for better styling
+st.markdown("""
+<style>
+    .main-header {
+        font-size: 2.5rem;
+        font-weight: bold;
+        margin-bottom: 0.5rem;
+        color: #1f77b4;
+    }
+    .sub-header {
+        font-size: 1.2rem;
+        color: #666;
+        margin-bottom: 2rem;
+    }
+    .stTextArea textarea {
+        font-family: 'Courier New', monospace;
+    }
+</style>
+""", unsafe_allow_html=True)
+
+
+class DBMLParser:
+    """Parse DBML code into structured data"""
+    
+    def __init__(self, dbml_code):
+        self.dbml_code = dbml_code
+        self.tables = {}
+        self.relationships = []
+        
+    def parse(self):
+        """Parse DBML code"""
+        self._parse_tables()
+        self._parse_relationships()
+        return self.tables, self.relationships
+    
+    def _parse_tables(self):
+        """Extract tables and columns from DBML"""
+        table_pattern = r'Table\s+(\w+)\s*{([^}]*)}'
+        
+        for match in re.finditer(table_pattern, self.dbml_code, re.DOTALL):
+            table_name = match.group(1)
+            table_body = match.group(2)
+            
+            columns = []
+            for line in table_body.strip().split('\n'):
+                line = line.strip()
+                if not line or line.startswith('//'):
+                    continue
+                
+                # Parse column definition
+                col_match = re.match(r'(\w+)\s+(\w+)(\s+\[[^\]]*\])?', line)
+                if col_match:
+                    col_name = col_match.group(1)
+                    col_type = col_match.group(2)
+                    constraints = col_match.group(3) or ''
+                    
+                    is_pk = 'pk' in constraints.lower() or 'primary key' in constraints.lower()
+                    is_unique = 'unique' in constraints.lower()
+                    is_null = 'null' in constraints.lower() and 'not null' not in constraints.lower()
+                    is_not_null = 'not null' in constraints.lower()
+                    
+                    columns.append({
+                        'name': col_name,
+                        'type': col_type,
+                        'is_pk': is_pk,
+                        'is_unique': is_unique,
+                        'is_null': is_null,
+                        'is_not_null': is_not_null
+                    })
+            
+            self.tables[table_name] = columns
+    
+    def _parse_relationships(self):
+        """Extract relationships from DBML"""
+        # Pattern: Ref: table1.column1 > table2.column2
+        ref_pattern = r'Ref:\s*(\w+)\.(\w+)\s*([<>-]+)\s*(\w+)\.(\w+)'
+        
+        for match in re.finditer(ref_pattern, self.dbml_code):
+            from_table = match.group(1)
+            from_column = match.group(2)
+            rel_type = match.group(3)
+            to_table = match.group(4)
+            to_column = match.group(5)
+            
+            # Determine relationship type
+            if '>' in rel_type:
+                rel_name = 'one-to-many'
+            elif '<' in rel_type:
+                rel_name = 'many-to-one'
+                from_table, to_table = to_table, from_table
+                from_column, to_column = to_column, from_column
+            else:
+                rel_name = 'one-to-one'
+            
+            self.relationships.append({
+                'from_table': from_table,
+                'from_column': from_column,
+                'to_table': to_table,
+                'to_column': to_column,
+                'type': rel_name
+            })
+
+
+class MermaidGenerator:
+    """Generate Mermaid diagram from parsed DBML"""
+    
+    def __init__(self, tables, relationships):
+        self.tables = tables
+        self.relationships = relationships
+    
+    def generate(self):
+        """Generate Mermaid ER diagram"""
+        lines = ['erDiagram']
+        
+        # Add tables with columns
+        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"    }}")
+        
+        # Add relationships
+        for rel in self.relationships:
+            if rel['type'] == 'one-to-many':
+                rel_symbol = '||--o{'
+            elif rel['type'] == 'many-to-one':
+                rel_symbol = '}o--||'
+            else:
+                rel_symbol = '||--||'
+            
+            lines.append(f"    {rel['from_table']} {rel_symbol} {rel['to_table']} : \"{rel['from_column']}-{rel['to_column']}\"")
+        
+        return '\n'.join(lines)
+
+
+class PostgreSQLGenerator:
+    """Generate PostgreSQL DDL from parsed DBML"""
+    
+    def __init__(self, tables, relationships):
+        self.tables = tables
+        self.relationships = relationships
+    
+    def generate(self):
+        """Generate PostgreSQL CREATE TABLE statements"""
+        sql_statements = []
+        
+        # Create tables
+        for table_name, columns in self.tables.items():
+            lines = [f"CREATE TABLE {table_name} ("]
+            
+            col_defs = []
+            pk_columns = []
+            
+            for col in columns:
+                col_def = f"    {col['name']} {self._map_type(col['type'])}"
+                
+                if col['is_pk']:
+                    pk_columns.append(col['name'])
+                
+                if col['is_not_null'] and not col['is_pk']:
+                    col_def += " NOT NULL"
+                elif col['is_null']:
+                    col_def += " NULL"
+                
+                if col['is_unique'] and not col['is_pk']:
+                    col_def += " UNIQUE"
+                
+                col_defs.append(col_def)
+            
+            # Add primary key constraint
+            if pk_columns:
+                col_defs.append(f"    PRIMARY KEY ({', '.join(pk_columns)})")
+            
+            lines.append(',\n'.join(col_defs))
+            lines.append(");")
+            
+            sql_statements.append('\n'.join(lines))
+        
+        # Add foreign key constraints
+        fk_statements = []
+        for rel in self.relationships:
+            fk_name = f"fk_{rel['from_table']}_{rel['to_table']}"
+            fk_stmt = f"""
+ALTER TABLE {rel['from_table']}
+ADD CONSTRAINT {fk_name}
+FOREIGN KEY ({rel['from_column']})
+REFERENCES {rel['to_table']}({rel['to_column']});"""
+            fk_statements.append(fk_stmt)
+        
+        all_sql = '\n\n'.join(sql_statements)
+        if fk_statements:
+            all_sql += '\n\n-- Foreign Key Constraints\n' + '\n'.join(fk_statements)
+        
+        return all_sql
+    
+    def _map_type(self, dbml_type):
+        """Map DBML types to PostgreSQL types"""
+        type_mapping = {
+            'int': 'INTEGER',
+            'integer': 'INTEGER',
+            'bigint': 'BIGINT',
+            'smallint': 'SMALLINT',
+            'varchar': 'VARCHAR(255)',
+            'text': 'TEXT',
+            'bool': 'BOOLEAN',
+            'boolean': 'BOOLEAN',
+            'date': 'DATE',
+            'datetime': 'TIMESTAMP',
+            'timestamp': 'TIMESTAMP',
+            'decimal': 'DECIMAL',
+            'float': 'REAL',
+            'double': 'DOUBLE PRECISION',
+            'uuid': 'UUID',
+            'json': 'JSONB'
+        }
+        
+        return type_mapping.get(dbml_type.lower(), dbml_type.upper())
+
+
+def main():
+    st.markdown('<div class="main-header">🗄️ DBML Visualizer & SQL Generator</div>', unsafe_allow_html=True)
+    st.markdown('<div class="sub-header">Visualize database schemas and generate PostgreSQL code</div>', unsafe_allow_html=True)
+    
+    # Sidebar
+    with st.sidebar:
+        st.header("📖 About")
+        st.write("""
+        This tool allows you to:
+        - 📊 Visualize DBML schemas
+        - 💾 Download diagram as image
+        - 🔧 Generate PostgreSQL DDL
+        - 📋 Copy SQL to clipboard
+        """)
+        
+        st.header("💡 Example DBML")
+        if st.button("Load Example"):
+            st.session_state.example_loaded = True
+    
+    # Example DBML
+    example_dbml = """Table users {
+  id int [pk, increment]
+  username varchar [unique, not null]
+  email varchar [unique, not null]
+  created_at timestamp
+}
+
+Table posts {
+  id int [pk, increment]
+  title varchar [not null]
+  content text
+  user_id int [not null]
+  created_at timestamp
+}
+
+Table comments {
+  id int [pk, increment]
+  post_id int [not null]
+  user_id int [not null]
+  content text [not null]
+  created_at timestamp
+}
+
+Ref: posts.user_id > users.id
+Ref: comments.post_id > posts.id
+Ref: comments.user_id > users.id"""
+    
+    # Main content
+    col1, col2 = st.columns([1, 1])
+    
+    with col1:
+        st.subheader("📝 DBML Input")
+        
+        default_value = example_dbml if st.session_state.get('example_loaded', False) else ""
+        dbml_input = st.text_area(
+            "Enter your DBML code:",
+            value=default_value,
+            height=400,
+            placeholder="Table users {\n  id int [pk]\n  name varchar\n}\n\nTable posts {\n  id int [pk]\n  user_id int\n}\n\nRef: posts.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
+                        
+                        st.success(f"✅ Successfully parsed {len(tables)} tables and {len(relationships)} relationships!")
+                
+                except Exception as e:
+                    st.error(f"❌ Error parsing DBML: {str(e)}")
+            else:
+                st.warning("⚠️ Please enter some DBML code first.")
+    
+    with col2:
+        if 'mermaid_code' in st.session_state:
+            st.subheader("📊 Database Diagram")
+            
+            # Display Mermaid diagram
+            st.markdown(f"""
+            ```mermaid
+            {st.session_state.mermaid_code}
+            ```
+            """)
+            
+            # Info about tables
+            st.info(f"📋 **{len(st.session_state.tables)}** tables, **{len(st.session_state.relationships)}** relationships")
+            
+        else:
+            st.info("👈 Enter DBML code and click 'Generate' to see the visualization")
+    
+    # SQL Output Section
+    if 'sql_code' in st.session_state:
+        st.markdown("---")
+        st.subheader("🔧 Generated PostgreSQL Code")
+        
+        col_sql1, col_sql2 = st.columns([4, 1])
+        
+        with col_sql1:
+            st.code(st.session_state.sql_code, language="sql", line_numbers=True)
+        
+        with col_sql2:
+            st.download_button(
+                label="📥 Download SQL",
+                data=st.session_state.sql_code,
+                file_name="schema.sql",
+                mime="text/sql",
+                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 __name__ == "__main__":
+    main()