Model/scripts/backfill_boiler_efficiency_band.py

"""Backfill the ``boiler_efficiency_band`` override for already-ingested boilers.

The SEDBUK band (ADR-0068) is losslessly present in every ``main_heating_system``
override's ``original_spreadsheet_description`` (``Boiler: C rated Combi``) but was
dropped at classification. This one-off populates the new ``boiler_efficiency_band``
override rows for the existing stock by re-parsing that text with the SAME guard the
live classifier uses (``boiler_efficiency_band_guard``), so the backfill and the
forward path cannot drift.

One ``boiler_efficiency_band`` row per ``main_heating_system`` row whose description
carries a band (``A``–``G``), keyed to the same ``(property_id, building_part)``.
Descriptions with no band (a plain boiler, a non-boiler heating system) get no row.
The modelling gate (gas/oil boilers only, ADR-0068) is applied at overlay time, not
here — so a band parsed off an electric ``Boiler: A rated NA`` is stored for
fidelity but ignored by the calculator, exactly as the live classifier stores it.

GATED: the ``boiler_efficiency_band`` value of the FE-owned ``override_component``
pgEnum must exist first (the Drizzle migration in the assessment-model repo). Until
then this writes nothing in prod — the Class-A/B deferred pattern.

DRY-RUN BY DEFAULT: prints the counts it would write and writes nothing. Pass
``--apply`` to execute inside a transaction. Idempotent — a row already carrying the
target band is left untouched, so re-running is a no-op.
"""

from __future__ import annotations

import argparse
from collections.abc import Iterable
from dataclasses import dataclass

from sqlalchemy import Connection, text

from domain.epc.property_overrides.boiler_efficiency_band import BoilerEfficiencyBand
from domain.epc.property_overrides.boiler_efficiency_band_guard import (
    boiler_efficiency_band_guard,
)
from scripts.e2e_common import build_engine, load_env

_OVERRIDE_COMPONENT = "boiler_efficiency_band"


@dataclass(frozen=True)
class MainHeatingRow:
    """The fields of a ``main_heating_system`` ``property_overrides`` row the
    backfill reads."""

    property_id: int
    portfolio_id: int
    building_part: int
    original_spreadsheet_description: str


@dataclass(frozen=True)
class BandBackfillRow:
    """A ``boiler_efficiency_band`` override row to upsert for a boiler."""

    property_id: int
    portfolio_id: int
    building_part: int
    override_value: str
    original_spreadsheet_description: str


def band_backfill_rows(rows: Iterable[MainHeatingRow]) -> list[BandBackfillRow]:
    """The ``boiler_efficiency_band`` rows to write for a set of
    ``main_heating_system`` rows — one per row whose description carries a
    parseable SEDBUK band, keyed to the same property + building part. Rows with
    no band (or the ``UNKNOWN`` sentinel) produce nothing.

    Pure and DB-free, so the parse/decide logic is unit-tested without a
    database; the connection wrapper below just feeds it rows and upserts."""
    backfill: list[BandBackfillRow] = []
    for row in rows:
        band = boiler_efficiency_band_guard(row.original_spreadsheet_description)
        if band is None or band is BoilerEfficiencyBand.UNKNOWN:
            continue
        backfill.append(
            BandBackfillRow(
                property_id=row.property_id,
                portfolio_id=row.portfolio_id,
                building_part=row.building_part,
                override_value=band.value,
                original_spreadsheet_description=row.original_spreadsheet_description,
            )
        )
    return backfill


_SELECT_MAIN_HEATING = text(
    """
    SELECT property_id, portfolio_id, building_part, original_spreadsheet_description
    FROM property_overrides
    WHERE override_component = 'main_heating_system'
    """
)
# Idempotent upsert on the (property, component, part) unique constraint — a row
# already carrying the target band is left untouched (updated_at unchanged).
_UPSERT_BAND = text(
    """
    INSERT INTO property_overrides
        (id, property_id, portfolio_id, building_part, override_component,
         override_value, original_spreadsheet_description, created_at, updated_at)
    VALUES
        (gen_random_uuid(), :property_id, :portfolio_id, :building_part,
         'boiler_efficiency_band', :override_value, :original_spreadsheet_description,
         now(), now())
    ON CONFLICT (property_id, override_component, building_part)
    DO UPDATE SET override_value = EXCLUDED.override_value,
                  original_spreadsheet_description =
                      EXCLUDED.original_spreadsheet_description,
                  updated_at = now()
    WHERE property_overrides.override_value <> EXCLUDED.override_value
    """
)


def backfill(conn: Connection, *, apply: bool) -> int:
    """Backfill the ``boiler_efficiency_band`` rows off the existing
    ``main_heating_system`` descriptions. Returns the number of band rows the run
    writes (or would write, in dry-run)."""
    rows = [
        MainHeatingRow(
            property_id=r.property_id,
            portfolio_id=r.portfolio_id,
            building_part=r.building_part,
            original_spreadsheet_description=r.original_spreadsheet_description or "",
        )
        for r in conn.execute(_SELECT_MAIN_HEATING)
    ]
    to_write = band_backfill_rows(rows)
    if apply:
        for band_row in to_write:
            conn.execute(
                _UPSERT_BAND,
                {
                    "property_id": band_row.property_id,
                    "portfolio_id": band_row.portfolio_id,
                    "building_part": band_row.building_part,
                    "override_value": band_row.override_value,
                    "original_spreadsheet_description": (
                        band_row.original_spreadsheet_description
                    ),
                },
            )
    return len(to_write)


def main() -> None:
    load_env()
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--apply",
        action="store_true",
        help="execute the writes (default: dry-run, writes nothing)",
    )
    args = parser.parse_args()

    engine = build_engine()
    with engine.begin() as conn:
        conn.execute(text("SET statement_timeout = 120000"))
        total = backfill(conn, apply=args.apply)

    verb = "backfilled" if args.apply else "would backfill"
    print(
        f"{verb} {total} boiler_efficiency_band override row(s) from the existing "
        "main_heating_system descriptions (Boiler: <A-G> rated ...). The modelling "
        "gate (gas/oil boilers only) is applied at overlay time (ADR-0068)."
    )
    if not args.apply:
        print("\nDRY-RUN — nothing written. Re-run with --apply to execute.")


if __name__ == "__main__":
    main()