| """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") |
|
|
| |
| _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() |
| |
| |
| 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]] = [] |
| 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)) |
| |
| 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 |
| total += _infer_source(source) |
| if total: |
| logger.info("inferred foreign keys", user_id=catalog.user_id, count=total) |
| return catalog |
|
|
