schema_sync/db_visualizer.py

import logging
import json
from sqlalchemy import text

logger = logging.getLogger(__name__)

class DatabaseVisualizer:
    def __init__(self, db_connector, schema_inspector):
        self.db = db_connector
        self.inspector = schema_inspector
        
    def get_database_schema(self):
        """Get complete database schema with tables, columns, and relationships"""
        try:
            # Get all tables
            tables_query = """
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public' AND table_type = 'BASE TABLE'
                ORDER BY table_name
            """
            tables_result = self.db.execute_query(tables_query)
            tables = [row[0] for row in tables_result]
            
            schema_data = {
                "tables": {},
                "relationships": []
            }
            
            # Get detailed info for each table
            for table in tables:
                # Get columns
                columns = self.inspector.get_column_info(table)
                
                # Get foreign keys
                foreign_keys = self.inspector.get_foreign_keys(table)
                
                # Get primary key
                pk_query = f"""
                    SELECT column_name
                    FROM information_schema.key_column_usage
                    WHERE table_name = '{table}'
                    AND constraint_name IN (
                        SELECT constraint_name
                        FROM information_schema.table_constraints
                        WHERE table_name = '{table}'
                        AND constraint_type = 'PRIMARY KEY'
                    )
                """
                pk_result = self.db.execute_query(pk_query)
                primary_keys = [row[0] for row in pk_result]
                
                # Store table info
                schema_data["tables"][table] = {
                    "columns": columns,
                    "primary_keys": primary_keys,
                    "foreign_keys": foreign_keys
                }
                
                # Add relationships
                for fk in foreign_keys:
                    relationship = {
                        "from_table": table,
                        "from_column": fk["column_name"],
                        "to_table": fk["foreign_table_name"],
                        "to_column": fk["foreign_column_name"]
                    }
                    schema_data["relationships"].append(relationship)
            
            return schema_data
            
        except Exception as e:
            logger.error(f"Error getting database schema: {str(e)}")
            return {"tables": {}, "relationships": []}
    
    def generate_mermaid_diagram(self):
        """Generate a Mermaid ER diagram of the database"""
        schema = self.get_database_schema()
        
        if not schema["tables"]:
            return "No tables found in database"
        
        mermaid_lines = ["erDiagram"]
        
        # Add tables and columns
        for table_name, table_info in schema["tables"].items():
            mermaid_lines.append(f"    {table_name} {{")
            
            for col in table_info["columns"]:
                col_name = col["column_name"]
                data_type = col["data_type"]
                is_pk = col_name in table_info["primary_keys"]
                is_fk = any(fk["column_name"] == col_name for fk in table_info["foreign_keys"])
                
                # Add type indicators
                type_indicator = ""
                if is_pk:
                    type_indicator = " PK"
                elif is_fk:
                    type_indicator = " FK"
                
                mermaid_lines.append(f"        {data_type} {col_name}{type_indicator}")
            
            mermaid_lines.append("    }")
        
        # Add relationships
        for rel in schema["relationships"]:
            # Mermaid relationship syntax: TableA ||--o{ TableB : relationship
            mermaid_lines.append(
                f"    {rel['to_table']} ||--o{{ {rel['from_table']} : \"{rel['from_column']} -> {rel['to_column']}\""
            )
        
        return "\n".join(mermaid_lines)
    
    def generate_table_summary(self):
        """Generate a text summary of database structure"""
        schema = self.get_database_schema()
        
        if not schema["tables"]:
            return "No tables found in database"
        
        summary_lines = ["# Database Schema Summary\n"]
        
        # Table overview
        summary_lines.append(f"**Total Tables:** {len(schema['tables'])}")
        summary_lines.append(f"**Total Relationships:** {len(schema['relationships'])}\n")
        
        # Detailed table info
        summary_lines.append("## Tables\n")
        
        for table_name, table_info in schema["tables"].items():
            summary_lines.append(f"### {table_name}")
            summary_lines.append(f"- **Columns:** {len(table_info['columns'])}")
            summary_lines.append(f"- **Primary Keys:** {', '.join(table_info['primary_keys']) if table_info['primary_keys'] else 'None'}")
            summary_lines.append(f"- **Foreign Keys:** {len(table_info['foreign_keys'])}")
            
            # Column details
            summary_lines.append("\n**Columns:**")
            for col in table_info["columns"]:
                nullable = "NULL" if col["is_nullable"] == "YES" else "NOT NULL"
                summary_lines.append(f"- `{col['column_name']}` ({col['data_type']}) {nullable}")
            
            summary_lines.append("")
        
        # Relationships summary
        if schema["relationships"]:
            summary_lines.append("## Relationships\n")
            for rel in schema["relationships"]:
                summary_lines.append(
                    f"- `{rel['from_table']}.{rel['from_column']}` → `{rel['to_table']}.{rel['to_column']}`"
                )
        
        return "\n".join(summary_lines)