mcp_database_universal/schema_inspector.py

"""Schema inspector — discovers tables, columns, relationships."""

from dataclasses import dataclass, field
from mcp_database_universal.engines.base import BaseEngine


@dataclass
class Relationship:
    from_table: str
    from_column: str
    to_table: str
    to_column: str
    cardinality: str  # "one-to-one", "one-to-many", "many-to-many"
    is_explicit: bool  # from DB constraint or from naming convention


class SchemaInspector:
    def __init__(self, engine: BaseEngine):
        self.engine = engine

    async def discover_relationships(self) -> list[Relationship]:
        tables = await self.engine.get_tables()
        relationships: list[Relationship] = []
        explicit_pairs: set[tuple[str, str, str, str]] = set()

        for table_info in tables:
            detail = await self.engine.get_table_detail(table_info.name)
            for fk in detail.foreign_keys:
                relationships.append(Relationship(
                    from_table=table_info.name,
                    from_column=fk.column,
                    to_table=fk.references_table,
                    to_column=fk.references_column,
                    cardinality="many-to-one",
                    is_explicit=True,
                ))
                explicit_pairs.add((
                    table_info.name, fk.column,
                    fk.references_table, fk.references_column,
                ))

        inferred = self._infer_by_naming(tables, explicit_pairs)
        relationships.extend(inferred)

        return relationships

    def _infer_by_naming(
        self,
        tables: list,
        explicit_pairs: set[tuple[str, str, str, str]],
    ) -> list[Relationship]:
        table_names = {t.name.lower() for t in tables}
        inferred: list[Relationship] = []

        for table_info in tables:
            detail = None
            for col in (table_info.columns if hasattr(table_info, 'columns') else []):
                col_lower = col.name.lower()
                if col_lower.endswith("_id") and col_lower != "id":
                    candidate_table = col_lower[:-3]
                    if candidate_table in table_names:
                        pair = (table_info.name, col.name, candidate_table, "id")
                        if pair not in explicit_pairs:
                            inferred.append(Relationship(
                                from_table=table_info.name,
                                from_column=col.name,
                                to_table=candidate_table,
                                to_column="id",
                                cardinality="many-to-one",
                                is_explicit=False,
                            ))

        return inferred

    def detect_junction_tables(
        self, relationships: list[Relationship]
    ) -> list[str]:
        fk_count: dict[str, list[Relationship]] = {}
        for rel in relationships:
            fk_count.setdefault(rel.from_table, []).append(rel)

        junctions = []
        for table, rels in fk_count.items():
            unique_targets = set(r.to_table for r in rels)
            if len(unique_targets) >= 2:
                junctions.append(table)

        return junctions

    def generate_mermaid(self, relationships: list[Relationship]) -> str:
        lines = ["erDiagram"]

        seen_tables: set[str] = set()
        for rel in relationships:
            seen_tables.add(rel.from_table)
            seen_tables.add(rel.to_table)

        for table in sorted(seen_tables):
            lines.append(f"    {table} {{")

            rels_from = [r for r in relationships if r.from_table == table]
            for rel in rels_from:
                lines.append(f"        string {rel.from_column} FK")

            lines.append("    }")

        for rel in relationships:
            if rel.cardinality == "many-to-one":
                lines.append("    " + rel.to_table + " ||--o{ " + rel.from_table + " : has")
            elif rel.cardinality == "one-to-one":
                lines.append("    " + rel.to_table + " ||--|| " + rel.from_table + " : has")
            elif rel.cardinality == "many-to-many":
                lines.append("    " + rel.to_table + " }o--o{ " + rel.from_table + " : has")

        return "\n".join(lines)