home-automation/alembic_app/versions/20260822_18_contract_scopes.py

"""add a billing scope to energy contracts

Revision ID: 20260822_18_contract_scopes
Revises: 20260822_17_warmtelink_readings
Create Date: 2026-08-22 00:00:00.000000

The upgrade preserves every existing contract, version and cost row. Existing
contracts predate scopes and therefore deterministically belong to electricity.
"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op


revision: str = "20260822_18_contract_scopes"
down_revision: Union[str, None] = "20260822_17_warmtelink_readings"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def _count(connection: sa.Connection, table: str) -> int:
    return int(connection.execute(sa.text(f"SELECT COUNT(*) FROM {table}")).scalar_one())


def _orphan_count(connection: sa.Connection) -> int:
    version_orphans = connection.execute(
        sa.text(
            "SELECT COUNT(*) FROM energy_contract_version v "
            "LEFT JOIN energy_contract c ON c.id = v.contract_id WHERE c.id IS NULL"
        )
    ).scalar_one()
    cost_orphans = connection.execute(
        sa.text(
            "SELECT COUNT(*) FROM energy_cost_period p "
            "LEFT JOIN energy_contract_version v ON v.id = p.contract_version_id "
            "WHERE p.contract_version_id IS NOT NULL AND v.id IS NULL"
        )
    ).scalar_one()
    return int(version_orphans) + int(cost_orphans)


def _audit_scope_upgrade(connection: sa.Connection, before: dict[str, int], orphan_before: int) -> None:
    after = {table: _count(connection, table) for table in before}
    if after != before:
        raise RuntimeError("contract scope migration row-count audit failed")
    if _orphan_count(connection) != orphan_before:
        raise RuntimeError("contract scope migration FK audit failed")
    invalid_scope_count = connection.execute(
        sa.text("SELECT COUNT(*) FROM energy_contract WHERE scope IS NULL OR scope != 'electricity'")
    ).scalar_one()
    if invalid_scope_count:
        raise RuntimeError("contract scope migration backfill audit failed")

    # Kept on Alembic's Config attributes rather than an environment switch so
    # isolated migration tests can deterministically exercise the rollback
    # boundary without changing production behavior.
    failure_injector = op.get_context().config.attributes.get("m8_t12_post_ddl_audit_failure")
    if callable(failure_injector):
        failure_injector()


def _apply_scope_schema() -> None:
    # SQLite batch mode reconstructs the table. The server default gives every
    # historical row its deterministic value during reconstruction.
    with op.batch_alter_table("energy_contract", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column("scope", sa.String(length=32), nullable=False, server_default="electricity")
        )
        batch_op.create_index("ix_energy_contract_scope", ["scope"])


def upgrade() -> None:
    connection = op.get_bind()
    before = {
        table: _count(connection, table)
        for table in ("energy_contract", "energy_contract_version", "energy_cost_period")
    }
    orphan_before = _orphan_count(connection)

    if connection.dialect.name != "sqlite":
        _apply_scope_schema()
        _audit_scope_upgrade(connection, before, orphan_before)
        return

    # Alembic marks SQLite batch DDL as non-transactional.  SQLite itself can
    # nevertheless atomically roll back CREATE/COPY/DROP/RENAME when an
    # explicit transaction owns the complete batch operation.  Keep the audit
    # inside that boundary so a failed audit cannot strand a revision-17 DB
    # with a revision-18 table shape.
    connection.exec_driver_sql("BEGIN IMMEDIATE")
    try:
        _apply_scope_schema()
        _audit_scope_upgrade(connection, before, orphan_before)
    except BaseException:
        connection.exec_driver_sql("ROLLBACK")
        raise
    else:
        connection.exec_driver_sql("COMMIT")


def downgrade() -> None:
    # Schema reversibility is only exercised against isolated temporary test DBs.
    with op.batch_alter_table("energy_contract", schema=None) as batch_op:
        batch_op.drop_index("ix_energy_contract_scope")
        batch_op.drop_column("scope")