src/catalog/fk_inference.py

"""Heuristic foreign-key inference for catalogs that ship no declared FKs.

The dedorch catalog (written by Go's introspection) currently carries **no**
`foreign_keys`, so the FK-backed-joins-only IR validator rejects every join the
planner proposes — cross-table questions ("revenue by product") can't run even
though the planner picks the right columns. Until Go captures real FK
constraints, we infer the obvious relational edges from naming conventions so the
planner and the validator agree on the same catalog.

Conservative by design (a wrong edge would silently corrupt joined results):
  - `schema` (database) sources only — joins are DB-only anyway
  - a foreign key is only inferred from a column named ``<base>_id``
  - the target must be the SINGLE other table whose name matches ``<base>``
    (singular/plural) and exposes an ``id`` column of the SAME data_type
  - ambiguous matches (0 or >1 candidate tables) are skipped, never guessed
  - sources that already declare ANY foreign key are left untouched (trust Go)
"""

from __future__ import annotations

import re

from src.catalog.models import ForeignKey, Source
from src.middlewares.logging import get_logger

from .models import Catalog

logger = get_logger("fk_inference")

# `<base>_id` — the conventional foreign-key column name (base must be non-empty).
_ID_COL = re.compile(r"^(?P<base>.+)_id$", re.IGNORECASE)


def _table_matches_base(table_name: str, base: str) -> bool:
    """Whether `table_name` is the table `<base>` refers to (singular/plural)."""
    n = table_name.lower()
    b = base.lower()
    # `orders`↔`order`, `products`↔`product`, `sales_agents`↔`agent` (suffix),
    # plus the singular form and the `-es` plural.
    return n == b or n == b + "es" or n.endswith(b + "s")


def _infer_source(source: Source) -> int:
    """Add inferred FK edges to one source's tables in place; return the count."""
    added = 0
    for table in source.tables:
        for col in table.columns:
            m = _ID_COL.match(col.name)
            if not m:
                continue
            base = m.group("base")
            candidates: list[tuple[str, str]] = []  # (target_table_id, target_column_id)
            for tgt in source.tables:
                if tgt.table_id == table.table_id:
                    continue
                if not _table_matches_base(tgt.name, base):
                    continue
                id_col = next(
                    (
                        c
                        for c in tgt.columns
                        if c.name.lower() == "id" and c.data_type == col.data_type
                    ),
                    None,
                )
                if id_col is not None:
                    candidates.append((tgt.table_id, id_col.column_id))
            # Only act on an unambiguous single match — never guess between many.
            if len(candidates) != 1:
                continue
            target_table_id, target_column_id = candidates[0]
            table.foreign_keys.append(
                ForeignKey(
                    column_id=col.column_id,
                    target_table_id=target_table_id,
                    target_column_id=target_column_id,
                )
            )
            added += 1
    return added


def infer_foreign_keys(catalog: Catalog) -> Catalog:
    """Infer FK edges in place for schema sources that declare none. Returns `catalog`.

    Sources that already carry any declared FK are left as-is (Go's real FKs win).
    """
    total = 0
    for source in catalog.sources:
        if source.source_type != "schema":
            continue
        if any(t.foreign_keys for t in source.tables):
            continue  # real FKs present — trust them, infer nothing
        total += _infer_source(source)
    if total:
        logger.info("inferred foreign keys", user_id=catalog.user_id, count=total)
    return catalog