Spaces:
Paused
Paused
File size: 6,108 Bytes
3545eca |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
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)
|