knut_bifrost/

codegen.rs

//! **rel2graph CODEGEN** — one [`MappingSpec`] projected onto executable
//! backends.
//!
//! bifröst **is** the rel2graph engine: a relational schema → a *mapping spec*
//! (nodes, merge keys, edges, directions, props) → materialised graph. Codegen
//! is *one tool* of that engine (its siblings: the live CDC listener, the
//! pipeline-config authoring, the schema→mapping inference in
//! [`crate::mapping::MappingSpec::from_arrow_schema`]). This module is the
//! **describe-first, execute-anywhere** half: it lowers a reviewed
//! [`MappingSpec`] into a single canonical [`GraphPlan`] and then *projects*
//! that plan onto a backend.
//!
//! ## The two-projections invariant
//!
//! A [`GraphPlan`] is projected two ways, and **both must agree exactly**:
//!
//! * [`GraphPlan::to_pyspark`] — a PySpark job that loads the schema into the
//!   graph sink (the **Spark** backend projection). Lives here, pure Rust.
//! * a **thund [`Pipeline`]** the all-Rust engine runs on *any*
//!   [`ExecBackend`] (Spark-Connect-SDP **or** native DataFusion) — lives in
//!   `knut-thund` (`knut_thund::rel2graph`, feature `rel2graph`), because that
//!   crate owns the execution IR. It projects from *this same* [`GraphPlan`].
//!
//! bifröst maps rel→graph **once**; thund picks the tech. Any divergence
//! between the two projections (a label, a merge key, an edge direction, the
//! batch size) is a bug — pinned by the `two_projections_agree` test in
//! `knut-thund` (which can see both projections without a dependency cycle:
//! `knut-pipelines` already depends on `knut-bifrost`, so bifröst must not
//! depend on thund; thund depending on bifröst is the acyclic direction).
//!
//! [`MappingSpec`]: crate::mapping::MappingSpec
//! [`Pipeline`]: https://docs.rs/knut-thund
//! [`ExecBackend`]: https://docs.rs/knut-thund

use std::fmt::Write as _;

use crate::error::{BifrostError, Result};
use crate::mapping::{MappingSpec, PipelineStep, SyncMode};

// ═══════════════════════════════════════════════════════════════════════════
// Options + backend/mode enums
// ═══════════════════════════════════════════════════════════════════════════

/// The relational **source plane** a schema is read from. In the intended
/// deployment Iceberg is the *data plane* and Postgres the *schema plane*
/// (the T1 catalog), so both are first-class.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SourceKind {
    /// A Postgres (or Postgres-catalog-fronted) source.
    Postgres,
    /// An Iceberg table (the streaming-capable data plane).
    Iceberg,
}

impl SourceKind {
    /// The lower-case token used in generated artifacts / RON.
    pub fn token(self) -> &'static str {
        match self {
            SourceKind::Postgres => "postgres",
            SourceKind::Iceberg => "iceberg",
        }
    }
}

/// The graph **sink** a plan targets.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SinkKind {
    /// FalkorDB (Cypher `UNWIND … MERGE`, the runtime sink bifröst already
    /// drives in [`crate::sink`]).
    FalkorDb,
    /// GraphAr columnar graph files.
    GraphAr,
}

impl SinkKind {
    /// The lower-case token used in generated artifacts / RON.
    pub fn token(self) -> &'static str {
        match self {
            SinkKind::FalkorDb => "falkordb",
            SinkKind::GraphAr => "graphar",
        }
    }
}

/// Batch vs streaming. Streaming requires an Iceberg source (the changelog
/// scan) — enforced by [`GraphPlan::from_mapping`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PlanMode {
    /// A bounded, idempotent full/incremental load.
    Batch,
    /// A continuous changelog-driven load (needs Iceberg).
    Stream,
}

impl PlanMode {
    /// The lower-case token used in generated artifacts / RON.
    pub fn token(self) -> &'static str {
        match self {
            PlanMode::Batch => "batch",
            PlanMode::Stream => "stream",
        }
    }
}

/// **Resolution tier** — where a node/edge mapping's authority *came from*, per
/// the rel2graph spec. Recorded on every plan element and surfaced in the
/// WARNINGS provenance block so a reviewer can see how confident each mapping
/// is (T1 catalog fact ▸ T4 inference proposal).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Tier {
    /// **T1** — the upstream Postgres catalog (PK/FK constraints). The answer
    /// for most schemas: Iceberg is the data plane, Postgres the schema plane.
    T1PostgresCatalog,
    /// **T2** — Iceberg identifier-field-ids (the table's declared identity).
    T2IcebergFieldIds,
    /// **T3** — `knut.*` TBLPROPERTIES on the table
    /// (`knut.node.label`/`knut.node.key`/`knut.edge.*`). What a hand-authored
    /// or reviewed [`MappingSpec`] encodes.
    T3TblProperties,
    /// **T4** — inclusion-dependency *inference* (the LEFT-ANTI-JOIN orphan
    /// test). **Proposals only** — emitted COMMENTED-OUT with evidence, never
    /// applied. A T4 element must never appear in the active plan.
    T4InclusionDependency,
}

impl Tier {
    /// A stable short token (`T1`…`T4`).
    pub fn token(self) -> &'static str {
        match self {
            Tier::T1PostgresCatalog => "T1",
            Tier::T2IcebergFieldIds => "T2",
            Tier::T3TblProperties => "T3",
            Tier::T4InclusionDependency => "T4",
        }
    }

    /// A one-line human description of the tier.
    pub fn describe(self) -> &'static str {
        match self {
            Tier::T1PostgresCatalog => "upstream Postgres catalog (PK/FK constraints)",
            Tier::T2IcebergFieldIds => "Iceberg identifier-field-ids",
            Tier::T3TblProperties => "knut.* TBLPROPERTIES on the table",
            Tier::T4InclusionDependency => {
                "inclusion-dependency inference (LEFT ANTI JOIN orphan test) — PROPOSAL only"
            }
        }
    }
}

/// Knobs for [`GraphPlan::from_mapping`]. Defaults match the rel2graph spec
/// (1000-row UNWIND batches, single coalesced writer, Iceberg→FalkorDB batch).
#[derive(Debug, Clone)]
pub struct CodegenOptions {
    /// The source plane the schema was read from.
    pub source: SourceKind,
    /// The graph sink to load.
    pub sink: SinkKind,
    /// Batch or streaming (streaming needs an Iceberg source).
    pub mode: PlanMode,
    /// Rows per `UNWIND … MERGE` batch (the idempotent write unit).
    pub batch_rows: usize,
    /// Writer concurrency — the `coalesce(N)` partition count the sink client
    /// runs inside.
    pub concurrency: usize,
    /// The single checkpoint / storage location (one per query).
    pub checkpoint: String,
    /// The tier the *declared* mapping elements are attributed to. A reviewed
    /// [`MappingSpec`] is T3 (its `knut.*` properties); a spec derived straight
    /// from an Iceberg schema is T2. The live-catalog T1 resolution is a
    /// worklist item (see the design doc).
    pub declared_tier: Tier,
}

impl Default for CodegenOptions {
    fn default() -> Self {
        CodegenOptions {
            source: SourceKind::Iceberg,
            sink: SinkKind::FalkorDb,
            mode: PlanMode::Batch,
            batch_rows: 1000,
            concurrency: 1,
            checkpoint: "/tmp/rel2graph/checkpoint".to_string(),
            declared_tier: Tier::T3TblProperties,
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// The canonical plan — the ONE object both projections derive from
// ═══════════════════════════════════════════════════════════════════════════

/// A node projection: one merge-key table → one graph node.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NodePlan {
    /// The fully-qualified source table.
    pub table: String,
    /// The graph node label (PascalCase singular).
    pub label: String,
    /// The merge key column — node identity.
    pub key: String,
    /// The non-key, non-FK scalar property columns, in **ordinal** order.
    pub props: Vec<String>,
    /// Where this node's authority came from.
    pub tier: Tier,
}

/// An edge projection: a directed relationship `from → to`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EdgePlan {
    /// The source table the edge rows are read from.
    pub table: String,
    /// The tail (child / "many") node label.
    pub from_label: String,
    /// The relationship type (SCREAMING_SNAKE).
    pub rel: String,
    /// The head (parent / "one") node label.
    pub to_label: String,
    /// The tail node's **merge-key property** — the property name the `from`
    /// node is identified/MERGEd by in the graph (its `id_column`). This is the
    /// property NAME on the node, which is NOT necessarily the column the value
    /// is read from — see [`Self::from_col`].
    pub from_key: String,
    /// The head node's **merge-key property** — the property name the `to` node
    /// is identified/MERGEd by in the graph (its `id_column`). For a self-ref /
    /// renamed FK this differs from the column the value is read from
    /// ([`Self::to_col`]): e.g. `Employee-REPORTS_TO->Employee` MERGEs on
    /// `employee_id` but reads the endpoint value from the `manager_id` column.
    pub to_key: String,
    /// The column in `table` the **tail** endpoint VALUE is read from
    /// (`src.column`). Usually equal to [`Self::from_key`]; differs only when the
    /// tail FK column is named differently from the node's key.
    pub from_col: String,
    /// The column in `table` the **head** endpoint VALUE is read from
    /// (`dst.column`, i.e. the FK column). Differs from [`Self::to_key`] for a
    /// self-ref / renamed FK (`manager_id` read, MERGEd on `employee_id`).
    pub to_col: String,
    /// Edge property columns, ordinal order.
    pub props: Vec<String>,
    /// Where this edge's authority came from.
    pub tier: Tier,
}

/// A **T4 proposal** — an inferred relationship the codegen must NOT apply,
/// only surface commented-out with its evidence. Carried through the plan so
/// both projections emit it identically (as a comment), never as live code.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EdgeProposal {
    /// The candidate edge (its `tier` is always [`Tier::T4InclusionDependency`]).
    pub edge: EdgePlan,
    /// Human-readable evidence (e.g. the orphan-row count of the LEFT ANTI
    /// JOIN test) a reviewer confirms before promoting it.
    pub evidence: String,
}

/// The canonical rel2graph plan: **the single source both projections read**.
///
/// Deterministic by construction (see [`GraphPlan::from_mapping`]):
/// tables schema-then-name, node props ordinal, nodes-before-edges, edges
/// sorted `(from_label, rel, to_label)`. No timestamps, no randomness, no
/// hashmap iteration order — the same input is the same bytes, always.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GraphPlan {
    /// Source plane.
    pub source: SourceKind,
    /// Graph sink.
    pub sink: SinkKind,
    /// Batch/streaming.
    pub mode: PlanMode,
    /// Rows per `UNWIND … MERGE` batch.
    pub batch_rows: usize,
    /// Writer concurrency (`coalesce(N)`).
    pub concurrency: usize,
    /// The single checkpoint / storage location.
    pub checkpoint: String,
    /// Nodes, in canonical order.
    pub nodes: Vec<NodePlan>,
    /// Edges, in canonical order.
    pub edges: Vec<EdgePlan>,
    /// T4 inference proposals — emitted commented-out, never applied.
    pub proposals: Vec<EdgeProposal>,
    /// Provenance + honesty notes, closing every artifact.
    pub warnings: Vec<String>,
}

impl GraphPlan {
    /// **Lower a reviewed [`MappingSpec`] into the canonical plan.**
    ///
    /// Walks tables in schema-then-name order; each `vertex` step becomes a
    /// [`NodePlan`], each `edge`/`fk_edges` step a [`EdgePlan`] (direction =
    /// `src → dst`, the child→parent FK direction). `custom` steps are opaque
    /// and cannot be projected to a declarative graph load, so they are
    /// SURFACED in [`warnings`](Self::warnings), never silently dropped.
    ///
    /// Errors on the blocking gaps the spec names: a streaming plan on a
    /// non-Iceberg source, or a node with no merge key.
    pub fn from_mapping(spec: &MappingSpec, opts: &CodegenOptions) -> Result<GraphPlan> {
        if opts.mode == PlanMode::Stream && opts.source != SourceKind::Iceberg {
            return Err(BifrostError::Other(format!(
                "MODE/SOURCE conflict: streaming needs an Iceberg source, got {:?}",
                opts.source
            )));
        }

        let mut nodes: Vec<NodePlan> = Vec::new();
        let mut edges: Vec<EdgePlan> = Vec::new();
        let mut warnings: Vec<String> = Vec::new();

        // Tables schema-then-name (namespace before leaf), deterministically.
        let mut tables: Vec<&crate::mapping::TableMapping> = spec.tables.iter().collect();
        tables.sort_by(|a, b| table_sort_key(&a.table).cmp(&table_sort_key(&b.table)));

        for tm in tables {
            if tm.mode == SyncMode::Cdc && opts.mode == PlanMode::Batch {
                warnings.push(format!(
                    "table `{}` declares CDC sync but the plan mode is batch — the changelog \
                     is read as a full insert load (see MODE)",
                    tm.table
                ));
            }
            for step in &tm.pipeline {
                match step {
                    PipelineStep::Vertex { vertex } => {
                        if vertex.id_column.is_empty() {
                            return Err(BifrostError::Other(format!(
                                "node `{}` (table `{}`) has no merge key",
                                vertex.label, tm.table
                            )));
                        }
                        nodes.push(NodePlan {
                            table: tm.table.clone(),
                            label: vertex.label.clone(),
                            key: vertex.id_column.clone(),
                            props: vertex.properties.clone(),
                            tier: opts.declared_tier,
                        });
                    }
                    PipelineStep::Edge { edge: t } | PipelineStep::FkEdges { fk_edges: t } => {
                        edges.push(EdgePlan {
                            table: tm.table.clone(),
                            from_label: t.src.label.clone(),
                            rel: t.edge_type.clone(),
                            to_label: t.dst.label.clone(),
                            from_key: t
                                .src
                                .id_column
                                .clone()
                                .unwrap_or_else(|| t.src.column.clone()),
                            to_key: t
                                .dst
                                .id_column
                                .clone()
                                .unwrap_or_else(|| t.dst.column.clone()),
                            // The VALUE is always read from the endpoint's own
                            // `column` (the transform fabric reads `src.column`
                            // / `dst.column`); only the MERGE *property name*
                            // comes from `id_column`. Keeping these separate is
                            // what makes self-ref / renamed FKs correct.
                            from_col: t.src.column.clone(),
                            to_col: t.dst.column.clone(),
                            props: t.properties.clone(),
                            tier: opts.declared_tier,
                        });
                    }
                    PipelineStep::Sql { .. } => {
                        warnings.push(format!(
                            "table `{}` uses a `sql:` step — its output is decided by the query \
                             at run time and cannot be projected to a declarative graph load; wire \
                             it by hand on the target backend (NOT dropped)",
                            tm.table
                        ));
                    }
                    PipelineStep::Custom { custom } => {
                        warnings.push(format!(
                            "table `{}` uses custom transform `{custom}` — an opaque Rust step \
                             that cannot be projected to a declarative graph load; wire it by \
                             hand on the target backend (NOT dropped)",
                            tm.table
                        ));
                    }
                }
            }
        }

        // DETERMINISM: edges sorted (from_label, rel, to_label). Nodes keep the
        // schema-then-name table order (one merge-key table → one node); props
        // stay ordinal.
        edges.sort_by(|a, b| {
            (&a.from_label, &a.rel, &a.to_label).cmp(&(&b.from_label, &b.rel, &b.to_label))
        });

        // Provenance summary — the WARNINGS block leads with it.
        warnings.insert(
            0,
            format!(
                "provenance: {} node(s) + {} edge(s) resolved at tier {} ({}); {} T4 proposal(s) \
                 held back (commented-out, never applied)",
                nodes.len(),
                edges.len(),
                opts.declared_tier.token(),
                opts.declared_tier.describe(),
                0
            ),
        );

        Ok(GraphPlan {
            source: opts.source,
            sink: opts.sink,
            mode: opts.mode,
            batch_rows: opts.batch_rows,
            concurrency: opts.concurrency,
            checkpoint: opts.checkpoint.clone(),
            nodes,
            edges,
            proposals: Vec::new(),
            warnings,
        })
    }

    /// Attach T4 inference proposals (from an inclusion-dependency pass run
    /// elsewhere). They are stamped [`Tier::T4InclusionDependency`] and will
    /// only ever be emitted COMMENTED-OUT. Returns `self` for chaining.
    pub fn with_proposals(mut self, proposals: impl IntoIterator<Item = EdgeProposal>) -> Self {
        for mut p in proposals {
            p.edge.tier = Tier::T4InclusionDependency;
            self.proposals.push(p);
        }
        // Re-summarise the proposal count in the provenance line.
        if let Some(first) = self.warnings.first_mut() {
            if first.starts_with("provenance:") {
                *first = format!(
                    "provenance: {} node(s) + {} edge(s) resolved; {} T4 proposal(s) held back \
                     (commented-out, never applied)",
                    self.nodes.len(),
                    self.edges.len(),
                    self.proposals.len()
                );
            }
        }
        self
    }

    /// A canonical, human-readable digest of the plan's **labels, merge keys,
    /// edge directions and batch size** — the *contract* both projections must
    /// honour. The two-projections-agree test compares what it can parse back
    /// out of each artifact against this. Deterministic and stable.
    pub fn contract(&self) -> String {
        let mut s = String::new();
        let _ = writeln!(s, "batch_rows={}", self.batch_rows);
        for n in &self.nodes {
            let _ = writeln!(s, "node {}={}", n.label, n.key);
        }
        for e in &self.edges {
            let _ = writeln!(s, "edge {}-{}->{}", e.from_label, e.rel, e.to_label);
        }
        s
    }

    // ───────────────────────────────────────────────────────────────────────
    // Projection (a): the PySpark job — the Spark backend
    // ───────────────────────────────────────────────────────────────────────

    /// **Projection (a): a PySpark job** loading the schema into the graph
    /// sink. Obeys the rel2graph Python hard rules: a FAIR scheduler, indexes
    /// created at startup, **nodes-before-edges**, the sink client constructed
    /// *inside* the partition function, a `coalesce()` before
    /// `foreachPartition`, and idempotent `UNWIND $rows AS r MERGE` in
    /// `batch_rows`-sized chunks. Streaming wraps the same per-partition writer
    /// in a single `foreachBatch` with one `checkpointLocation`.
    ///
    /// Deterministic: same plan ⇒ byte-identical Python.
    pub fn to_pyspark(&self) -> String {
        let mut s = String::new();
        let p = &mut s;
        let _ = writeln!(
            p,
            "# ── rel2graph codegen · Spark backend projection ──────────────────"
        );
        let _ = writeln!(
            p,
            "# source={} sink={} mode={} batch_rows={} concurrency={}",
            self.source.token(),
            self.sink.token(),
            self.mode.token(),
            self.batch_rows,
            self.concurrency
        );
        let _ = writeln!(
            p,
            "# GENERATED — one MappingSpec, two projections; edits belong upstream."
        );
        let _ = writeln!(p, "from pyspark.sql import SparkSession");
        let _ = writeln!(p);
        let _ = writeln!(p, "BATCH_ROWS = {}", self.batch_rows);
        let _ = writeln!(p, "CHECKPOINT = {:?}", self.checkpoint);
        let _ = writeln!(p);

        // Sink client + per-batch MERGE helpers, built INSIDE the partition fn.
        let _ = writeln!(p, "def _sink():");
        let _ = writeln!(
            p,
            "    # sink client constructed INSIDE the partition fn (never shipped from the driver)"
        );
        match self.sink {
            SinkKind::FalkorDb => {
                let _ = writeln!(p, "    from falkordb import FalkorDB");
                let _ = writeln!(
                    p,
                    "    return FalkorDB(host='localhost', port=6379).select_graph('rel2graph')"
                );
            }
            SinkKind::GraphAr => {
                let _ = writeln!(p, "    import pygraphar  # columnar graph writer");
                let _ = writeln!(p, "    return pygraphar.Writer('rel2graph')");
            }
        }
        let _ = writeln!(p);

        // Node writer.
        let _ = writeln!(p, "def _write_nodes(label, key, rows):");
        let _ = writeln!(p, "    g = _sink(); buf = []");
        let _ = writeln!(p, "    def flush():");
        let _ = writeln!(p, "        if not buf: return");
        p.push_str(
            "        g.query('UNWIND $rows AS r MERGE (n:`' + label + '` {`' + key + '`: r[\"' + key + '\"]}) SET n += r', {'rows': buf})\n",
        );
        let _ = writeln!(p, "        buf.clear()");
        let _ = writeln!(p, "    for row in rows:");
        let _ = writeln!(p, "        buf.append(row.asDict())");
        let _ = writeln!(p, "        if len(buf) >= BATCH_ROWS: flush()");
        let _ = writeln!(p, "    flush()");
        let _ = writeln!(p);

        // Edge writer.
        let _ = writeln!(p, "def _write_edges(fl, rel, tl, fk, tk, rows):");
        let _ = writeln!(p, "    g = _sink(); buf = []");
        let _ = writeln!(p, "    def flush():");
        let _ = writeln!(p, "        if not buf: return");
        p.push_str(
            "        g.query('UNWIND $rows AS r MERGE (a:`' + fl + '` {`' + fk + '`: r.src}) MERGE (b:`' + tl + '` {`' + tk + '`: r.dst}) MERGE (a)-[e:`' + rel + '`]->(b) SET e += r.props', {'rows': buf})\n",
        );
        let _ = writeln!(p, "        buf.clear()");
        let _ = writeln!(p, "    for row in rows:");
        let _ = writeln!(
            p,
            "        d = row.asDict(); props = {{k: v for k, v in d.items() if k not in ('src', 'dst')}}"
        );
        let _ = writeln!(
            p,
            "        buf.append({{'src': d['src'], 'dst': d['dst'], 'props': props}})"
        );
        let _ = writeln!(p, "        if len(buf) >= BATCH_ROWS: flush()");
        let _ = writeln!(p, "    flush()");
        let _ = writeln!(p);

        // Startup indexes — one per node label+key.
        let _ = writeln!(p, "def ensure_indexes(g):");
        if self.nodes.is_empty() {
            let _ = writeln!(p, "    pass  # no nodes");
        }
        for n in &self.nodes {
            let _ = writeln!(
                p,
                "    g.query('CREATE INDEX FOR (n:`{}`) ON (n.`{}`)')  # {}",
                n.label,
                n.key,
                n.tier.token()
            );
        }
        let _ = writeln!(p);

        // main(): FAIR scheduler, nodes before edges.
        let _ = writeln!(p, "def main():");
        let _ = writeln!(
            p,
            "    spark = SparkSession.builder.config('spark.scheduler.mode', 'FAIR').getOrCreate()"
        );
        let _ = writeln!(p, "    ensure_indexes(_sink())");
        let stream = self.mode == PlanMode::Stream;
        let reader = if stream { "readStream.table" } else { "table" };

        let _ = writeln!(p, "    # ── nodes (before edges) ──");
        for n in &self.nodes {
            let cols = node_select_cols(n);
            let df = format!("df_node_{}", sanitize(&n.label));
            let _ = writeln!(
                p,
                "    {df} = spark.{reader}({:?}).selectExpr({cols})",
                n.table
            );
            if stream {
                let _ = writeln!(
                    p,
                    "    {df}.writeStream.option('checkpointLocation', CHECKPOINT + '/node_{}').foreachBatch(",
                    sanitize(&n.label)
                );
                let _ = writeln!(
                    p,
                    "        lambda b, _: b.coalesce({}).foreachPartition(lambda rows: _write_nodes({:?}, {:?}, rows))).start()",
                    self.concurrency, n.label, n.key
                );
            } else {
                let _ = writeln!(
                    p,
                    "    {df}.coalesce({}).foreachPartition(lambda rows: _write_nodes({:?}, {:?}, rows))",
                    self.concurrency, n.label, n.key
                );
            }
        }

        let _ = writeln!(p, "    # ── edges ──");
        for e in &self.edges {
            let cols = edge_select_cols(e);
            let df = format!(
                "df_edge_{}_{}_{}",
                sanitize(&e.from_label),
                sanitize(&e.rel),
                sanitize(&e.to_label)
            );
            let _ = writeln!(
                p,
                "    {df} = spark.{reader}({:?}).selectExpr({cols})",
                e.table
            );
            if stream {
                let _ = writeln!(
                    p,
                    "    {df}.writeStream.option('checkpointLocation', CHECKPOINT + '/edge_{}_{}_{}').foreachBatch(",
                    sanitize(&e.from_label),
                    sanitize(&e.rel),
                    sanitize(&e.to_label)
                );
                let _ = writeln!(
                    p,
                    "        lambda b, _: b.coalesce({}).foreachPartition(lambda rows: _write_edges({:?}, {:?}, {:?}, {:?}, {:?}, rows))).start()",
                    self.concurrency, e.from_label, e.rel, e.to_label, e.from_key, e.to_key
                );
            } else {
                let _ = writeln!(
                    p,
                    "    {df}.coalesce({}).foreachPartition(lambda rows: _write_edges({:?}, {:?}, {:?}, {:?}, {:?}, rows))",
                    self.concurrency, e.from_label, e.rel, e.to_label, e.from_key, e.to_key
                );
            }
        }
        if stream {
            let _ = writeln!(p, "    spark.streams.awaitAnyTermination()");
        }
        let _ = writeln!(p);
        let _ = writeln!(p, "if __name__ == '__main__':");
        let _ = writeln!(p, "    main()");

        // T4 proposals — commented out with evidence, never applied.
        self.emit_proposals_py(p);
        // WARNINGS block closes the artifact.
        self.emit_warnings(p, "#");
        s
    }

    fn emit_proposals_py(&self, p: &mut String) {
        if self.proposals.is_empty() {
            return;
        }
        let _ = writeln!(p);
        let _ = writeln!(
            p,
            "# ── T4 PROPOSALS (inclusion-dependency inference) — NOT applied ──"
        );
        let _ = writeln!(
            p,
            "# Confirm the evidence, then add upstream in the MappingSpec. Never auto-applied."
        );
        for pr in &self.proposals {
            let e = &pr.edge;
            let _ = writeln!(
                p,
                "# ({}:`{}`)-[:`{}`]->({}:`{}`)  evidence: {}",
                sanitize(&e.from_label),
                e.from_label,
                e.rel,
                sanitize(&e.to_label),
                e.to_label,
                pr.evidence
            );
        }
    }

    /// The WARNINGS block, prefixed with `line_prefix` (e.g. `#` for Python,
    /// `//` for RON). Closes every artifact — provenance first, then notes.
    pub fn emit_warnings(&self, p: &mut String, line_prefix: &str) {
        let _ = writeln!(p);
        let _ = writeln!(p, "{line_prefix} ── WARNINGS ──");
        if self.warnings.is_empty() {
            let _ = writeln!(p, "{line_prefix} (none)");
        }
        for w in &self.warnings {
            let _ = writeln!(p, "{line_prefix} {w}");
        }
    }
}

// ── helpers ─────────────────────────────────────────────────────────────────

/// Schema-then-name sort key: `(namespace, leaf)` so `a.z` precedes `b.a` on
/// namespace but `s.accounts` precedes `s.users` on leaf. Deterministic.
fn table_sort_key(table: &str) -> (String, String) {
    match table.rsplit_once('.') {
        Some((ns, leaf)) => (ns.to_string(), leaf.to_string()),
        None => (String::new(), table.to_string()),
    }
}

/// Python identifier-safe form of a label/rel for a variable name.
fn sanitize(s: &str) -> String {
    s.chars()
        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
        .collect()
}

/// `selectExpr(...)` column list for a node: key first, then ordinal props,
/// each backtick-quoted.
fn node_select_cols(n: &NodePlan) -> String {
    let mut cols = vec![format!("\"`{}`\"", n.key)];
    cols.extend(n.props.iter().map(|c| format!("\"`{c}`\"")));
    cols.join(", ")
}

/// `selectExpr(...)` column list for an edge: `from_col AS src`, `to_col AS
/// dst`, then ordinal props (as themselves), each backtick-quoted.
///
/// The endpoint VALUES are read from the FK columns (`from_col`/`to_col`), NOT
/// the merge-key property names (`from_key`/`to_key`) — they differ for a
/// self-ref / renamed FK (`manager_id AS dst`, MERGEd on `employee_id`).
fn edge_select_cols(e: &EdgePlan) -> String {
    let mut parts = vec![
        format!("\"`{}` AS src\"", e.from_col),
        format!("\"`{}` AS dst\"", e.to_col),
    ];
    for c in &e.props {
        parts.push(format!("\"`{c}`\""));
    }
    parts.join(", ")
}

// ═══════════════════════════════════════════════════════════════════════════
// Tests — canonical plan + the PySpark (Spark backend) projection. The
// *two-projections-agree* invariant is pinned in `knut-thund` (which can see
// the thund `Pipeline` projection too); here we pin determinism, ordering,
// provenance, and the Python hard rules.
// ═══════════════════════════════════════════════════════════════════════════

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mapping::MappingSpec;

    /// A two-table social schema: users (with a company FK) + companies, whose
    /// declaration order is deliberately "wrong" (users before companies, edge
    /// listed before the plain node) so the canonical ordering has work to do.
    fn sample_spec() -> MappingSpec {
        let yaml = r#"
version: 1
tables:
  - table: lake.social.users
    pipeline:
      - vertex: { label: User, id_column: user_id, properties: [name, age] }
      - fk_edges:
          type: WORKS_AT
          src: { label: User,    column: user_id }
          dst: { label: Company, column: company_id }
          properties: [since]
  - table: lake.social.companies
    pipeline:
      - vertex: { label: Company, id_column: company_id, properties: [name] }
"#;
        MappingSpec::from_yaml(yaml).expect("sample yaml parses")
    }

    #[test]
    fn plan_is_canonical_and_deterministic() {
        let spec = sample_spec();
        let opts = CodegenOptions::default();
        let a = GraphPlan::from_mapping(&spec, &opts).unwrap();
        let b = GraphPlan::from_mapping(&spec, &opts).unwrap();
        // Byte-identical PySpark on repeat — no timestamps/random/hashmap order.
        assert_eq!(a.to_pyspark(), b.to_pyspark(), "codegen is deterministic");

        // Tables schema-then-name: `companies` (leaf) sorts before `users`, so
        // its node comes first even though it was declared last.
        assert_eq!(a.nodes.len(), 2);
        assert_eq!(a.nodes[0].label, "Company");
        assert_eq!(a.nodes[1].label, "User");
        // Props stay ordinal.
        assert_eq!(a.nodes[1].props, vec!["name", "age"]);

        // Edge: direction User -[WORKS_AT]-> Company, keys carried.
        assert_eq!(a.edges.len(), 1);
        let e = &a.edges[0];
        assert_eq!(
            (e.from_label.as_str(), e.rel.as_str(), e.to_label.as_str()),
            ("User", "WORKS_AT", "Company")
        );
        assert_eq!(e.from_key, "user_id");
        assert_eq!(e.to_key, "company_id");
    }

    #[test]
    fn pyspark_obeys_the_python_hard_rules() {
        let plan = GraphPlan::from_mapping(&sample_spec(), &CodegenOptions::default()).unwrap();
        let py = plan.to_pyspark();

        assert!(py.contains("spark.scheduler.mode"), "FAIR scheduler set");
        assert!(py.contains("'FAIR'"));
        assert!(py.contains("ensure_indexes"), "indexes created at startup");
        assert!(py.contains("CREATE INDEX FOR (n:`User`) ON (n.`user_id`)"));
        assert!(
            py.contains("def _sink():"),
            "sink client built inside partition fn"
        );
        assert!(
            py.contains(".coalesce(1).foreachPartition("),
            "coalesce before foreachPartition"
        );
        assert!(
            py.contains("UNWIND $rows AS r MERGE"),
            "idempotent UNWIND MERGE"
        );
        assert!(py.contains("if len(buf) >= BATCH_ROWS:"), "batched writes");
        assert!(py.contains("BATCH_ROWS = 1000"));

        // REGRESSION GUARD: the Cypher `MERGE (n:… {`key`: …})` carries REAL
        // SINGLE braces `{`/`}` in the emitted text. These are literal `{` in
        // the source; emitting them via `format!`/`writeln!` (which read `{` as
        // a placeholder) is a compile error that has bitten the kube-feature
        // build — so they are built with a raw `push_str`, and this pins that
        // the single-brace Cypher survives to the output verbatim.
        assert!(
            py.contains("MERGE (n:`' + label + '` {`' + key + '`: r[\"' + key + '\"]})"),
            "node MERGE keeps its single-brace Cypher: {py}"
        );
        assert!(
            py.contains("MERGE (a:`' + fl + '` {`' + fk + '`: r.src})"),
            "edge MERGE keeps its single-brace Cypher"
        );
        assert!(
            py.contains("{'rows': buf}"),
            "the params dict uses single braces"
        );
        // And never the doubled braces a `format!` escape would leave behind.
        assert!(!py.contains("{{"), "no doubled `{{` leaked into the output");
        assert!(!py.contains("}}"), "no doubled `}}` leaked into the output");

        // Nodes-before-edges: both node writers precede the edge writer in main().
        let main = &py[py.find("def main()").unwrap()..];
        let node_at = main.find("_write_nodes").unwrap();
        let edge_at = main.find("_write_edges").unwrap();
        assert!(node_at < edge_at, "nodes are loaded before edges");

        // WARNINGS block closes the artifact, provenance first.
        assert!(py.trim_end().contains("── WARNINGS ──"));
        assert!(py.contains("provenance:"));
    }

    #[test]
    fn streaming_needs_iceberg_and_wraps_one_foreachbatch() {
        let spec = sample_spec();
        // Stream on Postgres is refused up front.
        let bad = CodegenOptions {
            source: SourceKind::Postgres,
            mode: PlanMode::Stream,
            ..Default::default()
        };
        assert!(
            GraphPlan::from_mapping(&spec, &bad).is_err(),
            "stream needs iceberg"
        );

        let ok = CodegenOptions {
            source: SourceKind::Iceberg,
            mode: PlanMode::Stream,
            ..Default::default()
        };
        let py = GraphPlan::from_mapping(&spec, &ok).unwrap().to_pyspark();
        assert!(py.contains("readStream.table"), "streaming reader");
        assert!(
            py.contains("foreachBatch("),
            "foreachBatch wraps the per-partition writer"
        );
        assert!(
            py.contains("checkpointLocation"),
            "one checkpointLocation per query"
        );
        assert!(py.contains("awaitAnyTermination"));
    }

    #[test]
    fn custom_step_is_surfaced_not_dropped() {
        let yaml = r#"
version: 1
tables:
  - table: lake.x.t
    pipeline:
      - vertex: { label: T, id_column: id }
      - custom: my_special_transform
"#;
        let spec = MappingSpec::from_yaml(yaml).unwrap();
        let plan = GraphPlan::from_mapping(&spec, &CodegenOptions::default()).unwrap();
        assert_eq!(plan.nodes.len(), 1);
        assert!(
            plan.warnings
                .iter()
                .any(|w| w.contains("my_special_transform") && w.contains("NOT dropped")),
            "the opaque custom step is surfaced in WARNINGS: {:?}",
            plan.warnings
        );
    }

    #[test]
    fn t4_proposals_are_emitted_commented_out_never_applied() {
        let plan = GraphPlan::from_mapping(&sample_spec(), &CodegenOptions::default())
            .unwrap()
            .with_proposals([EdgeProposal {
                edge: EdgePlan {
                    table: "lake.social.users".into(),
                    from_label: "User".into(),
                    rel: "REFERRED_BY".into(),
                    to_label: "User".into(),
                    from_key: "user_id".into(),
                    to_key: "user_id".into(),
                    from_col: "user_id".into(),
                    to_col: "referrer_id".into(),
                    props: vec![],
                    tier: Tier::T1PostgresCatalog, // will be overwritten to T4
                },
                evidence: "LEFT ANTI JOIN found 0 orphan referrer_id rows (n=10000)".into(),
            }]);
        // Stamped T4 regardless of what the caller passed.
        assert_eq!(plan.proposals[0].edge.tier, Tier::T4InclusionDependency);
        let py = plan.to_pyspark();
        assert!(py.contains("T4 PROPOSALS"), "proposals section present");
        assert!(py.contains("REFERRED_BY"), "the proposal is shown");
        assert!(py.contains("evidence:"), "with its evidence");
        // But NEVER as an active edge writer.
        assert!(
            !py.contains("_write_edges(\"User\", \"REFERRED_BY\""),
            "a T4 proposal is never emitted as a live edge load"
        );
    }
}