aetherdialect 0.1.0__py3-none-any.whl

text2sql/schema.py

@@ -0,0 +1,973 @@
+"""Schema reflection, graph building, and profiling orchestration.
+
+Reflects table and column metadata from a live database or SQL file, infers missing foreign keys by naming convention, and builds a shortest-path join graph between all table pairs. A local JSON cache avoids repeated reflection; the cache is invalidated by a content hash of the schema structure. Profiling data (statistics, value domains, LLM-assigned roles, and allowed operations) is layered on top of the reflected schema before caching. Adaptive limits derived from schema statistics are also stored.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import asdict
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from .config import EngineConfig, PolicyConfig
+from .contracts_base import (
+    ColumnMetadata,
+    FKEdge,
+    SchemaGraph,
+    SchemaLimits,
+    TableMetadata,
+)
+from .core_utils import debug, schema_hash_fp
+from .dialect import get_dialect
+from .schema_profiling import (
+    apply_column_roles_llm,
+    assign_column_ops,
+    extract_tables_from_catalog,
+    extract_tables_from_catalog_sql_connector,
+    parse_sql_file,
+    profile_schema,
+    profile_schema_spark,
+    profile_schema_sql_connector,
+)
+
+
+def compute_schema_stats(schema: SchemaGraph) -> dict[str, Any]:
+    """Compute schema-wide column availability statistics for adaptive limit calculation.
+
+    Args:
+
+        schema: Populated ``SchemaGraph`` with profiled column metadata.
+
+    Returns:
+
+        Dictionary with keys: ``total_filterable``, ``total_groupable``, ``total_aggregatable``, ``min/max_filterable_per_table``, ``min/max_groupable_per_table``, ``table_count``, and ``filterable_per_table`` (list of per-table counts).
+    """
+    debug("[schema.compute_schema_stats] computing schema statistics for adaptive limits")
+
+    stats = {
+        "total_filterable": 0,
+        "total_groupable": 0,
+        "total_aggregatable": 0,
+        "min_filterable_per_table": float("inf"),
+        "max_filterable_per_table": 0,
+        "min_groupable_per_table": float("inf"),
+        "max_groupable_per_table": 0,
+        "table_count": len(schema.tables),
+        "filterable_per_table": [],
+    }
+
+    table_details = []
+    for table_name, table in schema.tables.items():
+        filterable_count = sum(1 for col in table.columns.values() if col.is_filterable)
+        groupable_count = sum(1 for col in table.columns.values() if col.is_groupable)
+        aggregatable_count = sum(1 for col in table.columns.values() if col.is_aggregatable)
+
+        table_details.append(
+            {
+                "table": table_name,
+                "filterable": filterable_count,
+                "groupable": groupable_count,
+                "aggregatable": aggregatable_count,
+            }
+        )
+
+        stats["total_filterable"] += filterable_count
+        stats["total_groupable"] += groupable_count
+        stats["total_aggregatable"] += aggregatable_count
+
+        if filterable_count > 0:
+            stats["min_filterable_per_table"] = min(stats["min_filterable_per_table"], filterable_count)
+            stats["max_filterable_per_table"] = max(stats["max_filterable_per_table"], filterable_count)
+            stats["filterable_per_table"].append(filterable_count)
+
+        if groupable_count > 0:
+            stats["min_groupable_per_table"] = min(stats["min_groupable_per_table"], groupable_count)
+            stats["max_groupable_per_table"] = max(stats["max_groupable_per_table"], groupable_count)
+
+    if stats["min_filterable_per_table"] == float("inf"):
+        stats["min_filterable_per_table"] = 0
+    if stats["min_groupable_per_table"] == float("inf"):
+        stats["min_groupable_per_table"] = 0
+
+    debug("[schema.compute_schema_stats] per-table column counts:")
+    for td in table_details:
+        debug(
+            f"  {td['table']}: filterable={td['filterable']}, groupable={td['groupable']}, aggregatable={td['aggregatable']}"
+        )
+
+    debug("[schema.compute_schema_stats] schema-wide statistics:")
+    debug(f"  table_count: {stats['table_count']}")
+    debug(f"  total_filterable: {stats['total_filterable']}")
+    debug(f"  total_groupable: {stats['total_groupable']}")
+    debug(f"  total_aggregatable: {stats['total_aggregatable']}")
+    debug(f"  min_filterable_per_table: {stats['min_filterable_per_table']}")
+    debug(f"  max_filterable_per_table: {stats['max_filterable_per_table']}")
+    debug(f"  min_groupable_per_table: {stats['min_groupable_per_table']}")
+    debug(f"  max_groupable_per_table: {stats['max_groupable_per_table']}")
+    debug(f"  filterable_per_table distribution: {stats['filterable_per_table']}")
+
+    return stats
+
+
+def compute_schema_limits(schema_stats: dict[str, Any]) -> SchemaLimits:
+    """Compute adaptive pipeline limits from schema statistics.
+
+    Args:
+
+        schema_stats: Dictionary as returned by ``compute_schema_stats``.
+
+    Returns:
+
+        ``SchemaLimits`` with ``max_filters``, ``max_groupby``, and ``max_tables``.
+    """
+    table_count = schema_stats.get("table_count", 1)
+    total_filterable = schema_stats.get("total_filterable", 0)
+    total_groupable = schema_stats.get("total_groupable", 0)
+
+    max_filters = max(1, total_filterable // table_count) if table_count > 0 else 1
+    max_groupby = max(1, total_groupable // table_count) if table_count > 0 else 1
+
+    if table_count <= 3:
+        max_tables = table_count
+    elif table_count <= 10:
+        max_tables = 3
+    else:
+        max_tables = 4
+
+    return SchemaLimits(
+        max_filters=max_filters,
+        max_groupby=max_groupby,
+        max_tables=max_tables,
+    )
+
+
+def _edge_key(e: FKEdge) -> tuple[str, tuple[str, ...], str, tuple[str, ...]]:
+    """Generate a stable, sortable tuple key for an FK edge.
+
+    Args:
+
+        e: ``FKEdge`` to key.
+
+    Returns:
+
+        Tuple of ``(src_table, src_cols_tuple, dst_table, dst_cols_tuple)``.
+    """
+    return (e.src_table, tuple(e.src_cols), e.dst_table, tuple(e.dst_cols))
+
+
+def _reflect_enum_values(engine: Any, schema_name: str) -> dict[str, list[str]]:
+    """Reflect all enum types and their ordered values from the database.
+
+    Args:
+
+        engine: SQLAlchemy ``Engine`` connected to the target database.
+
+        schema_name: Database schema name to introspect.
+
+    Returns:
+
+        Dictionary mapping enum type name to a list of its allowed values, or an empty dict if reflection fails.
+    """
+    dialect = get_dialect()
+    try:
+        enum_values = dialect.reflect_enums(engine, schema_name)
+        debug(f"[schema.reflect_enum_values] found {len(enum_values)} enum types")
+        return enum_values
+    except Exception as e:
+        debug(f"[schema.reflect_enum_values] failed to reflect enums: {e}")
+        return {}
+
+
+def _table_to_dict(table: TableMetadata) -> dict[str, Any]:
+    """Serialize a TableMetadata instance to a plain dictionary.
+
+    Args:
+
+        table: ``TableMetadata`` to serialize.
+
+    Returns:
+
+        Dictionary with ``name``, ``columns``, ``primary_key``, ``foreign_keys``, ``role``, ``row_count``, and ``description`` fields.
+    """
+    return {
+        "name": table.name,
+        "columns": {k: asdict(v) for k, v in table.columns.items()},
+        "primary_key": table.primary_key,
+        "foreign_keys": [asdict(fk) for fk in table.foreign_keys],
+        "partition_columns": table.partition_columns,
+        "role": table.role,
+        "row_count": table.row_count,
+        "description": table.description,
+    }
+
+
+def _table_from_dict(d: dict[str, Any]) -> TableMetadata:
+    """Deserialize a TableMetadata instance from a plain dictionary.
+
+    Args:
+
+        d: Dictionary with keys matching ``TableMetadata`` fields.
+
+    Returns:
+
+        Populated ``TableMetadata`` with ``ColumnMetadata`` and ``FKEdge`` sub-objects.
+    """
+    columns = {k: ColumnMetadata.from_dict(v) for k, v in d["columns"].items()}
+    return TableMetadata(
+        name=d["name"],
+        columns=columns,
+        primary_key=d["primary_key"],
+        foreign_keys=[FKEdge(**fk) for fk in d["foreign_keys"]],
+        partition_columns=d.get("partition_columns", []),
+        role=d.get("role"),
+        row_count=d.get("row_count", 0),
+        description=d.get("description", ""),
+    )
+
+
+def _reverse_fk_path(path: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Reverse a FK path by flipping each edge's direction and reversing list order.
+
+    Args:
+
+        path: List of FK edge dicts with ``src_table``, ``src_cols``, ``dst_table``, and ``dst_cols``.
+
+    Returns:
+
+        New list of edge dicts with src/dst swapped and order reversed.
+    """
+    reversed_path = []
+    for e in reversed(path):
+        reversed_path.append(
+            {
+                "src_table": e["dst_table"],
+                "src_cols": e["dst_cols"],
+                "dst_table": e["src_table"],
+                "dst_cols": e["src_cols"],
+            }
+        )
+    return reversed_path
+
+
+def _analyze_fk_path_topology(path: list[dict[str, Any]]) -> tuple[str, str, list[str]]:
+    """Analyze an FK path to determine its topology type, anchor table, and leaf tables.
+
+    Args:
+
+        path: List of FK edge dicts representing a join path.
+
+    Returns:
+
+        Tuple of ``(topology_type, anchor_table, leaf_tables)`` where ``topology_type`` is one of ``'none'``, ``'linear'``, ``'star'``, or ``'tree'``.
+    """
+    if not path:
+        return ("none", "", [])
+    table_counts: dict[str, int] = {}
+    for e in path:
+        src = e["src_table"]
+        dst = e["dst_table"]
+        table_counts[src] = table_counts.get(src, 0) + 1
+        table_counts[dst] = table_counts.get(dst, 0) + 1
+    if not table_counts:
+        return ("none", "", [])
+    leaves = sorted([t for t, c in table_counts.items() if c == 1])
+    hubs = sorted(
+        [t for t, c in table_counts.items() if c > 1],
+        key=lambda t: (-table_counts[t], t),
+    )
+    if len(leaves) == 2 and len(hubs) == len(table_counts) - 2:
+        return ("linear", min(leaves), leaves)
+    if len(hubs) == 1:
+        return ("star", hubs[0], leaves)
+    if hubs:
+        return ("tree", hubs[0], leaves)
+    return ("linear", min(table_counts.keys()), list(table_counts.keys()))
+
+
+def _normalize_fk_path(path: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Normalize an FK join path to a canonical form based on its topology.
+
+    For linear paths, edges are ordered so traversal starts from the lexicographically smallest leaf. For star or tree paths, edges are reoriented to flow outward from the hub (anchor) table.
+
+    Args:
+
+        path: List of FK edge dicts to normalize.
+
+    Returns:
+
+        Reordered and/or flipped list of edge dicts in canonical form.
+    """
+    if not path:
+        return path
+    topology_type, anchor, leaves = _analyze_fk_path_topology(path)
+    if topology_type == "none":
+        return path
+    if topology_type == "linear":
+        start_table = path[0]["src_table"]
+        if start_table == anchor:
+            return path
+        return _reverse_fk_path(path)
+    edge_map: dict[str, list[dict[str, Any]]] = {}
+    for e in path:
+        src = e["src_table"]
+        dst = e["dst_table"]
+        if src == anchor:
+            edge_map.setdefault(dst, []).append(e)
+        elif dst == anchor:
+            flipped = {
+                "src_table": dst,
+                "src_cols": e["dst_cols"],
+                "dst_table": src,
+                "dst_cols": e["src_cols"],
+            }
+            edge_map.setdefault(src, []).append(flipped)
+        else:
+            for branch_key in sorted(edge_map.keys()):
+                branch_tables = set()
+                for be in edge_map[branch_key]:
+                    branch_tables.add(be["src_table"])
+                    branch_tables.add(be["dst_table"])
+                if src in branch_tables or dst in branch_tables:
+                    edge_map[branch_key].append(e)
+                    break
+    normalized = []
+    for branch_key in sorted(edge_map.keys()):
+        normalized.extend(edge_map[branch_key])
+    return normalized if normalized else path
+
+
+def _infer_missing_fks(tables: dict[str, TableMetadata]) -> list[FKEdge]:
+    """Infer missing foreign keys from column naming conventions.
+
+    Matches columns ending in ``'_id'``, ``'_key'``, ``'id'``, or ``'key'`` against table names, then verifies the candidate target table has a matching primary key column.
+
+    Args:
+
+        tables: Mapping of table name to ``TableMetadata`` as reflected from the database.
+
+    Returns:
+
+        List of ``FKEdge`` instances for inferred relationships not already declared in the schema.
+    """
+    inferred = []
+    suffixes = ["_id", "_key", "id", "key"]
+    for table_name, table in tables.items():
+        for col_name, col in table.columns.items():
+            if col.is_foreign_key:
+                continue
+            if col.is_primary_key:
+                continue
+            matched_suffix = None
+            for suffix in suffixes:
+                if col_name.endswith(suffix):
+                    matched_suffix = suffix
+                    break
+            if not matched_suffix:
+                continue
+            prefix_full = col_name[: -len(matched_suffix)]
+            if not prefix_full:
+                continue
+            candidate_prefixes = [prefix_full]
+            if "_" in prefix_full:
+                parts = prefix_full.split("_")
+                candidate_prefixes.append(parts[-1])
+                for i in range(len(parts) - 1, 0, -1):
+                    candidate_prefixes.append("_".join(parts[i:]))
+            for prefix in candidate_prefixes:
+                if prefix not in tables:
+                    continue
+                if table_name == prefix:
+                    continue
+                target_table = tables[prefix]
+                if not target_table.primary_key:
+                    continue
+                target_pk = target_table.primary_key[0]
+                target_matched = False
+                for target_suffix in suffixes:
+                    if target_pk.endswith(target_suffix):
+                        target_prefix = target_pk[: -len(target_suffix)]
+                        if target_prefix == prefix:
+                            target_matched = True
+                            break
+                if target_matched:
+                    debug(f"[schema.infer_missing_fks] candidate: {table_name}.{col_name} -> {prefix}.{target_pk}")
+                    inferred.append(
+                        FKEdge(
+                            src_table=table_name,
+                            src_cols=[col_name],
+                            dst_table=prefix,
+                            dst_cols=[target_pk],
+                        )
+                    )
+                    break
+    return inferred
+
+
+def _reflect_schema(engine: Any, schema_name: str = None) -> SchemaGraph:
+    """Reflect a database schema using SQLAlchemy and build a join-path graph.
+
+    Reflects all tables and columns, resolves explicit FK constraints, infers missing FKs by naming convention, reflects enum types, and runs BFS to compute shortest join paths between every table pair.
+
+    Args:
+
+        engine: SQLAlchemy ``Engine`` connected to the target database.
+
+        schema_name: Database schema to reflect. Defaults to ``EngineConfig.RUNTIME.SCHEMA`` or ``'public'`` if not set.
+
+    Returns:
+
+        Fully populated ``SchemaGraph`` with ``tables``, ``join_paths_multi``, ``schema_hash``, ``created_at`` timestamp, and ``enum_values``.
+    """
+    if schema_name is None:
+        schema_name = EngineConfig.RUNTIME.SCHEMA if hasattr(EngineConfig.RUNTIME, "SCHEMA") else "public"
+
+    from sqlalchemy import MetaData
+
+    debug(f"[schema.reflect_schema] reflecting from database schema '{schema_name}'")
+    md = MetaData(schema=schema_name)
+    md.reflect(bind=engine)
+
+    tables: dict[str, TableMetadata] = {}
+
+    for t in md.tables.values():
+        columns: dict[str, ColumnMetadata] = {}
+        for c in t.columns:
+            columns[c.name] = ColumnMetadata(
+                name=c.name,
+                data_type=str(c.type),
+                is_primary_key=c.name in [pk.name for pk in t.primary_key.columns],
+                is_foreign_key=False,
+                fk_target=None,
+            )
+
+        tables[t.name] = TableMetadata(
+            name=t.name,
+            columns=columns,
+            primary_key=[c.name for c in t.primary_key.columns],
+            foreign_keys=[],
+        )
+
+    debug(f"[schema.reflect_schema] found {len(tables)} tables")
+
+    for t in md.tables.values():
+        for fk in t.foreign_key_constraints:
+            e = FKEdge(
+                src_table=t.name,
+                src_cols=[el.parent.name for el in fk.elements],
+                dst_table=fk.elements[0].column.table.name,
+                dst_cols=[el.column.name for el in fk.elements],
+            )
+            tables[t.name].foreign_keys.append(e)
+
+            for i, src_col in enumerate(e.src_cols):
+                if src_col in tables[t.name].columns:
+                    tables[t.name].columns[src_col].is_foreign_key = True
+                    tables[t.name].columns[src_col].fk_target = (
+                        e.dst_table,
+                        e.dst_cols[i],
+                    )
+
+            debug(
+                f"[schema.reflect_schema] explicit FK: {e.src_table}.{e.src_cols[0]} -> {e.dst_table}.{e.dst_cols[0]}"
+            )
+
+    fk_count = sum(len(tbl.foreign_keys) for tbl in tables.values())
+    debug(f"[schema.reflect_schema] found {fk_count} foreign key edges")
+
+    inferred_fks = _infer_missing_fks(tables)
+    if inferred_fks:
+        debug(f"[schema.reflect_schema] inferred {len(inferred_fks)} missing FKs from naming conventions")
+        for e in inferred_fks:
+            tables[e.src_table].foreign_keys.append(e)
+            for i, src_col in enumerate(e.src_cols):
+                if src_col in tables[e.src_table].columns:
+                    tables[e.src_table].columns[src_col].is_foreign_key = True
+                    tables[e.src_table].columns[src_col].fk_target = (
+                        e.dst_table,
+                        e.dst_cols[i],
+                    )
+            debug(f"[schema.reflect_schema]   {e.src_table}.{e.src_cols[0]} -> {e.dst_table}.{e.dst_cols[0]}")
+
+    enum_values = _reflect_enum_values(engine, schema_name)
+
+    adj: dict[str, list[FKEdge]] = {t: [] for t in tables}
+    for tbl in tables.values():
+        for e in tbl.foreign_keys:
+            if e.src_table not in tables or e.dst_table not in tables:
+                continue
+            adj[e.src_table].append(e)
+            adj[e.dst_table].append(
+                FKEdge(
+                    src_table=e.dst_table,
+                    src_cols=e.dst_cols,
+                    dst_table=e.src_table,
+                    dst_cols=e.src_cols,
+                )
+            )
+    for t in adj:
+        adj[t] = sorted(adj[t], key=lambda x: _edge_key(x))
+
+    debug("[schema.reflect_schema] computing shortest join paths")
+    join_paths_multi: dict[str, dict[str, list[list[dict[str, Any]]]]] = {}
+
+    tlist = sorted(tables.keys())
+    for s in tlist:
+        join_paths_multi[s] = {}
+        queue = [s]
+        prev: dict[str, tuple[str, FKEdge]] = {}
+        seen = {s}
+        while queue:
+            cur = queue.pop(0)
+            for e in adj[cur]:
+                nxt = e.dst_table
+                if nxt in seen:
+                    continue
+                seen.add(nxt)
+                prev[nxt] = (cur, e)
+                queue.append(nxt)
+
+        for t in tlist:
+            if t == s:
+                join_paths_multi[s][t] = [[]]
+                continue
+            if t in prev:
+                path = []
+                cur = t
+                while cur != s:
+                    p, e = prev[cur]
+                    path.append(e)
+                    cur = p
+                path.reverse()
+                sp = [asdict(e) for e in path]
+                join_paths_multi[s][t] = [_normalize_fk_path(sp)]
+            else:
+                join_paths_multi[s][t] = []
+
+    schema_hash = schema_hash_fp({k: _table_to_dict(v) for k, v in tables.items()})
+
+    sg = SchemaGraph(
+        tables=tables,
+        join_paths_multi=join_paths_multi,
+        schema_hash=schema_hash,
+        created_at=datetime.now().isoformat(),
+        enum_values=enum_values,
+    )
+
+    return sg
+
+
+def load_or_create_schema_graph(engine: Any = None) -> SchemaGraph:
+    """Load the schema graph from the JSON cache or build it fresh from the database.
+
+    If a valid cached schema exists (hash matches) it is returned without any database access. Otherwise the schema is reflected, profiled, and saved.
+
+    Args:
+
+        engine: SQLAlchemy ``Engine`` required for PostgreSQL reflection. Not used for Databricks.
+
+    Returns:
+
+        Fully populated and profiled ``SchemaGraph``.
+
+    Raises:
+
+        ValueError: If ``engine`` is ``None`` for PostgreSQL or the engine type is unsupported.
+    """
+    schema_json_path = EngineConfig.SCHEMA_JSON_PATH
+
+    debug(f"[schema.load_or_create_schema_graph] engine_type={EngineConfig.TYPE}")
+
+    if not PolicyConfig.REGENERATE_SCHEMA_GRAPH and os.path.exists(schema_json_path):
+        debug(f"[schema.load_or_create_schema_graph] loading from cache '{schema_json_path}'")
+        with open(schema_json_path, encoding="utf-8") as f:
+            d = json.load(f)
+
+        tables = {k: _table_from_dict(v) for k, v in d["tables"].items()}
+
+        recomputed_hash = schema_hash_fp(d["tables"])
+        cached_hash = d.get("schema_hash", "")
+        if recomputed_hash != cached_hash:
+            print("Schema cache outdated, rebuilding...")
+            debug(
+                f"[schema.load_or_create_schema_graph] hash mismatch: cached={cached_hash[:16]} recomputed={recomputed_hash[:16]}"
+            )
+            os.remove(schema_json_path)
+        else:
+            join_paths_multi = d.get("join_paths_multi")
+            if not join_paths_multi:
+                join_paths_multi = {}
+                jp = d.get("join_paths", {})
+                for a in jp:
+                    join_paths_multi[a] = {}
+                    for b in jp[a]:
+                        join_paths_multi[a][b] = [jp[a][b]] if jp[a][b] is not None else []
+
+            debug(f"[schema.load_or_create_schema_graph] loaded {len(tables)} tables from cache")
+
+            sg = SchemaGraph(
+                tables=tables,
+                join_paths_multi=join_paths_multi,
+                schema_hash=d["schema_hash"],
+                enum_values=d.get("enum_values", {}),
+                schema_stats=d["schema_stats"],
+            )
+
+            return sg
+
+    debug("[schema.load_or_create_schema_graph] no cache found, building schema")
+
+    spark_session = None
+    databricks_connection = None
+    try:
+        if EngineConfig.TYPE == "databricks":
+            if getattr(EngineConfig.RUNTIME, "has_native_connection", lambda: False)():
+                try:
+                    from databricks import sql as dbsql
+
+                    databricks_connection = dbsql.connect(
+                        server_hostname=EngineConfig.RUNTIME.SERVER_HOSTNAME,
+                        http_path=EngineConfig.RUNTIME.HTTP_PATH,
+                        access_token=EngineConfig.RUNTIME.ACCESS_TOKEN,
+                    )
+                except Exception as e:
+                    raise RuntimeError(f"Failed to connect via databricks-sql-connector: {e}") from e
+            else:
+                from pyspark.sql import SparkSession
+
+                spark_session = SparkSession.builder.getOrCreate()
+            sg = _load_or_create_schema_databricks(
+                spark_session=spark_session,
+                connection=databricks_connection,
+            )
+        elif EngineConfig.TYPE == "postgresql":
+            if engine is None:
+                raise ValueError("Engine is required for PostgreSQL schema reflection")
+            sg = _load_or_create_schema_postgresql(engine)
+        else:
+            raise ValueError(f"Unsupported engine type: {EngineConfig.TYPE}")
+
+        _add_profiling_data(engine, sg, spark_session, databricks_connection)
+    finally:
+        if databricks_connection is not None:
+            try:
+                databricks_connection.close()
+            except Exception:
+                pass
+
+    debug("[schema.load_or_create_schema_graph] computing schema stats after profiling")
+    sg.schema_stats = compute_schema_stats(sg)
+
+    _save_schema_to_cache(sg, schema_json_path)
+    debug("[schema.load_or_create_schema_graph] cache saved with profiling data")
+
+    return sg
+
+
+def _tables_meta_to_schema_graph(tables_meta: dict[str, dict]) -> SchemaGraph:
+    """Convert a raw table metadata dictionary to a fully connected ``SchemaGraph``.
+
+    Builds ``ColumnMetadata`` and ``FKEdge`` objects, then runs BFS join-path computation between all table pairs.
+
+    Args:
+
+        tables_meta: Dictionary mapping table name to metadata dicts with keys ``column_names_original``, ``column_types``, ``primary_keys``, and ``foreign_keys``.
+
+    Returns:
+
+        ``SchemaGraph`` with ``tables``, ``join_paths_multi``, ``schema_hash``, and ``created_at``.
+    """
+    tables: dict[str, TableMetadata] = {}
+
+    for table_name, meta in tables_meta.items():
+        columns: dict[str, ColumnMetadata] = {}
+        col_names = meta.get("column_names_original", [])
+        col_types = meta.get("column_types", [])
+        pk_cols = meta.get("primary_keys", [])
+
+        for i, col_name in enumerate(col_names):
+            col_type = col_types[i] if i < len(col_types) else "UNKNOWN"
+            columns[col_name] = ColumnMetadata(
+                name=col_name,
+                data_type=col_type,
+                is_primary_key=col_name in pk_cols,
+                is_foreign_key=False,
+                fk_target=None,
+            )
+
+        fk_edges = []
+        for fk in meta.get("foreign_keys", []):
+            edge = FKEdge(
+                src_table=table_name,
+                src_cols=fk["src_cols"],
+                dst_table=fk["dst_table"],
+                dst_cols=fk["dst_cols"],
+            )
+            fk_edges.append(edge)
+
+            for i, src_col in enumerate(fk["src_cols"]):
+                if src_col in columns:
+                    columns[src_col].is_foreign_key = True
+                    columns[src_col].fk_target = (fk["dst_table"], fk["dst_cols"][i])
+
+        partition_cols = meta.get("partition_columns", [])
+        tables[table_name] = TableMetadata(
+            name=table_name,
+            columns=columns,
+            primary_key=pk_cols,
+            foreign_keys=fk_edges,
+            partition_columns=partition_cols,
+        )
+
+    fk_count = sum(len(tbl.foreign_keys) for tbl in tables.values())
+    debug(f"[schema.tables_meta_to_schema_graph] {len(tables)} tables, {fk_count} FK edges")
+
+    adj: dict[str, list[FKEdge]] = {t: [] for t in tables}
+    for tbl in tables.values():
+        for e in tbl.foreign_keys:
+            if e.src_table not in tables or e.dst_table not in tables:
+                continue
+            adj[e.src_table].append(e)
+            adj[e.dst_table].append(
+                FKEdge(
+                    src_table=e.dst_table,
+                    src_cols=e.dst_cols,
+                    dst_table=e.src_table,
+                    dst_cols=e.src_cols,
+                )
+            )
+    for t in adj:
+        adj[t] = sorted(adj[t], key=lambda x: _edge_key(x))
+
+    debug("[schema.tables_meta_to_schema_graph] computing shortest join paths")
+
+    join_paths_multi: dict[str, dict[str, list[list[dict[str, Any]]]]] = {}
+    tlist = sorted(tables.keys())
+
+    for s in tlist:
+        join_paths_multi[s] = {}
+        queue = [s]
+        prev: dict[str, tuple[str, FKEdge]] = {}
+        seen = {s}
+        while queue:
+            cur = queue.pop(0)
+            for e in adj[cur]:
+                nxt = e.dst_table
+                if nxt in seen:
+                    continue
+                seen.add(nxt)
+                prev[nxt] = (cur, e)
+                queue.append(nxt)
+
+        for t in tlist:
+            if t == s:
+                join_paths_multi[s][t] = [[]]
+                continue
+            if t in prev:
+                path = []
+                cur = t
+                while cur != s:
+                    p, e = prev[cur]
+                    path.append(e)
+                    cur = p
+                path.reverse()
+                sp = [asdict(e) for e in path]
+                join_paths_multi[s][t] = [_normalize_fk_path(sp)]
+            else:
+                join_paths_multi[s][t] = []
+
+    schema_hash = schema_hash_fp({k: _table_to_dict(v) for k, v in tables.items()})
+
+    sg = SchemaGraph(
+        tables=tables,
+        join_paths_multi=join_paths_multi,
+        schema_hash=schema_hash,
+        created_at=datetime.now().isoformat(),
+        enum_values={},
+    )
+
+    return sg
+
+
+def _load_or_create_schema_postgresql(engine: Any) -> SchemaGraph:
+    """Build a ``SchemaGraph`` for PostgreSQL from a live database or SQL file fallback.
+
+    Attempts live reflection first; falls back to parsing a SQL DDL file if reflection fails and ``EngineConfig.RUNTIME.SQL_FILE_PATH`` is configured.
+
+    Args:
+
+        engine: SQLAlchemy ``Engine`` connected to the PostgreSQL database.
+
+    Returns:
+
+        ``SchemaGraph`` built from the database or SQL file.
+
+    Raises:
+
+        ValueError: If reflection fails and no SQL file is available or parseable.
+    """
+    try:
+        debug("[schema.load_or_create_schema_postgresql] reflecting_database")
+        sg = _reflect_schema(engine)
+        debug(f"[schema.load_or_create_schema_postgresql] reflected: {len(sg.tables)} tables")
+        return sg
+    except Exception as e:
+        debug(f"[schema.load_or_create_schema_postgresql] reflection_failed: {e}")
+        sql_file_path = getattr(EngineConfig.RUNTIME, "SQL_FILE_PATH", None)
+
+        if sql_file_path and os.path.exists(sql_file_path):
+            debug(f"[schema.load_or_create_schema_postgresql] parsing_sql_file: {sql_file_path}")
+            tables_meta = parse_sql_file(Path(sql_file_path))
+
+            if not tables_meta or len(tables_meta) == 0:
+                raise ValueError("Both database reflection and SQL file parsing failed") from e
+
+            return _tables_meta_to_schema_graph(tables_meta)
+        else:
+            raise ValueError(f"Database reflection failed and no SQL file available: {e}") from e
+
+
+def _add_profiling_data(
+    engine: Any,
+    sg: SchemaGraph,
+    spark_session=None,
+    connection=None,
+) -> None:
+    """Add column profiling data to a SchemaGraph in-place.
+
+    Runs the three-step profiling pipeline: 1) ``profile_schema`` — collect statistics
+    (distinct count, null ratio, min/max, top-K). 2) ``apply_column_roles_llm`` — LLM
+    assigns table and column roles with overrides. 3) ``assign_column_ops`` —
+    deterministic operation flags based on final roles.
+
+    Falls back to heuristics after ``MAX_ROLE_CLASSIFICATION_RETRIES`` LLM failures.
+
+    Args:
+
+        engine: SQLAlchemy ``Engine`` for PostgreSQL profiling; unused for Databricks.
+
+        sg: ``SchemaGraph`` to enrich in-place.
+
+        spark_session: Active ``SparkSession`` for Databricks Spark-based profiling.
+
+        connection: Active ``databricks.sql`` connection for connector-based profiling.
+    """
+    debug("[schema.add_profiling_data] Step 1: profiling columns (statistics)")
+
+    if EngineConfig.TYPE == "databricks":
+        catalog = EngineConfig.RUNTIME.CATALOG
+        schema_name = EngineConfig.RUNTIME.SCHEMA
+        if connection is not None:
+            profile_schema_sql_connector(connection, catalog, schema_name, sg)
+        else:
+            if spark_session is None:
+                from pyspark.sql import SparkSession
+
+                spark_session = SparkSession.builder.getOrCreate()
+            profile_schema_spark(spark_session, catalog, schema_name, sg)
+    else:
+        profile_schema(engine, sg)
+
+    debug("[schema.add_profiling_data] Step 2: inferring column and table roles via LLM")
+    apply_column_roles_llm(sg)
+
+    debug("[schema.add_profiling_data] Step 3: assigning column operations (deterministic)")
+    assign_column_ops(sg)
+
+
+def _save_schema_to_cache(sg: SchemaGraph, schema_json_path: str) -> None:
+    """Save a SchemaGraph to a JSON cache file.
+
+    Args:
+
+        sg: ``SchemaGraph`` to persist.
+
+        schema_json_path: Absolute path to the output JSON file.
+    """
+
+    def _dump_mixed(fh, data: dict[str, Any]):
+        fh.write("{")
+        first = True
+        for k in sorted(data.keys()):
+            if not first:
+                fh.write(",")
+            first = False
+            fh.write(json.dumps(k, ensure_ascii=False))
+            fh.write(":")
+            v = data[k]
+            if isinstance(v, dict):
+                json.dump(v, fh, indent=2, sort_keys=True, ensure_ascii=False)
+            else:
+                json.dump(v, fh, separators=(",", ":"), ensure_ascii=False)
+        fh.write("}")
+
+    cache_data = {
+        "tables": {k: _table_to_dict(v) for k, v in sg.tables.items()},
+        "join_paths_multi": sg.join_paths_multi,
+        "schema_hash": sg.schema_hash,
+        "created_at": sg.created_at,
+        "enum_values": sg.enum_values or {},
+        "schema_stats": sg.schema_stats or {},
+    }
+
+    debug(f"[schema.save_schema_to_cache] saving to '{schema_json_path}'")
+    with open(schema_json_path, "w", encoding="utf-8") as f:
+        _dump_mixed(f, cache_data)
+
+
+def _load_or_create_schema_databricks(
+    spark_session=None,
+    connection=None,
+) -> SchemaGraph:
+    """Build a ``SchemaGraph`` for Databricks from a SQL DDL file or catalog introspection.
+
+    Parses ``EngineConfig.RUNTIME.SQL_FILE_PATH`` if it exists and is non-empty; otherwise
+    extracts table metadata from the Databricks catalog via databricks-sql-connector
+    (when ``connection`` is provided) or Spark.
+
+    Args:
+
+        spark_session: Active ``SparkSession`` for Spark-based catalog extraction.
+
+        connection: Active ``databricks.sql`` connection for connector-based extraction.
+
+    Returns:
+
+        ``SchemaGraph`` built from the SQL file or catalog.
+
+    Raises:
+
+        RuntimeError: If no SQL file is available and catalog extraction fails.
+    """
+    sql_file_path = getattr(EngineConfig.RUNTIME, "SQL_FILE_PATH", None)
+
+    if sql_file_path and os.path.exists(sql_file_path):
+        debug(f"[schema.load_or_create_schema_databricks] parsing SQL file '{sql_file_path}'")
+        tables_meta = parse_sql_file(Path(sql_file_path))
+
+        if not tables_meta or len(tables_meta) == 0:
+            debug(
+                "[schema.load_or_create_schema_databricks] SQL file parsing returned 0 tables, falling back to catalog extraction"
+            )
+            sql_file_path = None
+
+    if not sql_file_path or not os.path.exists(sql_file_path):
+        debug("[schema.load_or_create_schema_databricks] SQL file not found or empty, attempting catalog extraction")
+        catalog = EngineConfig.RUNTIME.CATALOG
+        schema = EngineConfig.RUNTIME.SCHEMA
+        try:
+            if connection is not None:
+                tables_meta = extract_tables_from_catalog_sql_connector(
+                    connection, catalog, schema
+                )
+            else:
+                from pyspark.sql import SparkSession
+
+                spark = spark_session if spark_session else SparkSession.builder.getOrCreate()
+                tables_meta = extract_tables_from_catalog(spark, catalog, schema)
+        except Exception as e:
+            raise RuntimeError(f"Cannot build schema: no SQL file and catalog extraction failed: {e}") from e
+
+    return _tables_meta_to_schema_graph(tables_meta)