"""Self-heal compliance id sequences on every deploy

Revision ID: 018
Revises: 017
Create Date: 2026-05-19 16:00:00.000000

Background: when Base.metadata.create_all() runs at app startup before
alembic upgrade (or after a failed partial migration), Postgres ends
up with tables whose id column has no DEFAULT nextval(...). Inserts
then fail with NotNullViolation. Tester has hit this on every deploy
involving the 4 new Plan-D/E tables.

This migration is idempotent and runs every time `alembic upgrade
head` is executed. For each of the 4 affected tables:
  1. Create the sequence if missing.
  2. Set the column DEFAULT to nextval(sequence).
  3. ALTER SEQUENCE OWNED BY column.id so DROP TABLE cleans up.
  4. setval to MAX(id)+1 so future inserts pick up where rows ended.

Safe to re-run forever — every step is conditional / idempotent.
"""
from alembic import op


revision = '018'
down_revision = '017'
branch_labels = None
depends_on = None


_TABLES = (
    "compliance_results",
    "compliance_checks",
    "compliance_impacts",
    "asset_risk_snapshots",
)


def upgrade() -> None:
    for table in _TABLES:
        seq = f"{table}_id_seq"
        op.execute(f"""
        DO $$
        BEGIN
            -- Skip when the table itself does not exist (older DB).
            IF NOT EXISTS (
                SELECT 1 FROM information_schema.tables
                WHERE table_name = '{table}'
            ) THEN
                RAISE NOTICE '{table} does not exist — skipping seq heal';
                RETURN;
            END IF;

            -- 1. Create sequence if missing.
            IF NOT EXISTS (
                SELECT 1 FROM pg_class WHERE relkind = 'S' AND relname = '{seq}'
            ) THEN
                EXECUTE 'CREATE SEQUENCE {seq}';
            END IF;

            -- 2. Make sure the id column uses it as DEFAULT.
            EXECUTE 'ALTER TABLE {table} ALTER COLUMN id SET DEFAULT nextval(''{seq}''::regclass)';

            -- 3. Ownership so DROP TABLE drops the sequence too.
            EXECUTE 'ALTER SEQUENCE {seq} OWNED BY {table}.id';

            -- 4. setval to MAX(id)+1, defensive coalesce for empty tables.
            PERFORM setval('{seq}', COALESCE((SELECT MAX(id) FROM {table}), 0) + 1, false);
        END
        $$;
        """)


def downgrade() -> None:
    # No-op: dropping sequences would break existing data on rollback.
    pass
