"""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 ``_id`` - the target must be the SINGLE other table whose name matches ```` (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") # `_id` — the conventional foreign-key column name (base must be non-empty). _ID_COL = re.compile(r"^(?P.+)_id$", re.IGNORECASE) def _table_matches_base(table_name: str, base: str) -> bool: """Whether `table_name` is the table `` 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