108 lines
4.1 KiB
Python
108 lines
4.1 KiB
Python
"""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")
|