state-hub/migrations/versions/c4d5e6f7a8b9_inbox_feed_revision.py

"""inbox change-feed revisions and tombstones (STATE-WP-0098-T05)

Implements §2 of docs/inbox-change-feed-contract.md (accepted 2026-10-13, D4):

- sequence ``inbox_revision_seq``;
- ``revision bigint NOT NULL`` on ``agent_messages`` and ``message_receipts``,
  backfilled with ONE sequence value (one snapshot = one revision), indexed;
- ``message_tombstones``;
- a BEFORE INSERT OR UPDATE trigger on both tables that skips no-op updates,
  takes a transaction-scoped advisory lock (commit order = revision order) and
  assigns one revision per transaction (cached per txid in a txn-local GUC);
- a BEFORE DELETE trigger on ``agent_messages`` that writes a tombstone.

The SQL is a frozen copy of ``api.models.agent_message.INBOX_FEED_FUNCTIONS_SQL``;
tests/test_message_feed.py checks the two agree.

Revision ID: c4d5e6f7a8b9
Revises: b3c4d5e6f7a8
"""
import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects.postgresql import UUID

revision = "c4d5e6f7a8b9"
down_revision = "b3c4d5e6f7a8"
branch_labels = None
depends_on = None

INBOX_ADVISORY_LOCK_KEY = 5_348_490_098

FUNCTIONS_SQL = f"""
CREATE OR REPLACE FUNCTION inbox_feed_revision() RETURNS bigint
LANGUAGE plpgsql AS $fn$
DECLARE
    tx text := txid_current()::text;
    cached text := coalesce(current_setting('statehub.inbox_rev', true), '');
    rev bigint;
BEGIN
    PERFORM pg_advisory_xact_lock({INBOX_ADVISORY_LOCK_KEY});
    IF split_part(cached, ':', 1) = tx THEN
        RETURN split_part(cached, ':', 2)::bigint;
    END IF;
    rev := nextval('inbox_revision_seq');
    PERFORM set_config('statehub.inbox_rev', tx || ':' || rev::text, true);
    RETURN rev;
END
$fn$;

CREATE OR REPLACE FUNCTION inbox_feed_stamp_revision() RETURNS trigger
LANGUAGE plpgsql AS $fn$
BEGIN
    IF TG_OP = 'UPDATE' THEN
        IF NEW IS NOT DISTINCT FROM OLD THEN
            RETURN NEW;
        END IF;
    END IF;
    NEW.revision := inbox_feed_revision();
    RETURN NEW;
END
$fn$;

CREATE OR REPLACE FUNCTION inbox_feed_tombstone() RETURNS trigger
LANGUAGE plpgsql AS $fn$
BEGIN
    INSERT INTO message_tombstones (message_id, revision, deleted_at, reason)
    VALUES (
        OLD.id,
        inbox_feed_revision(),
        now(),
        nullif(current_setting('statehub.delete_reason', true), '')
    )
    ON CONFLICT (message_id) DO UPDATE
        SET revision = EXCLUDED.revision,
            deleted_at = EXCLUDED.deleted_at,
            reason = EXCLUDED.reason;
    RETURN OLD;
END
$fn$;
"""

TRIGGERS_SQL = """
CREATE TRIGGER agent_messages_inbox_revision
    BEFORE INSERT OR UPDATE ON agent_messages
    FOR EACH ROW EXECUTE FUNCTION inbox_feed_stamp_revision();
CREATE TRIGGER agent_messages_inbox_tombstone
    BEFORE DELETE ON agent_messages
    FOR EACH ROW EXECUTE FUNCTION inbox_feed_tombstone();
CREATE TRIGGER message_receipts_inbox_revision
    BEFORE INSERT OR UPDATE ON message_receipts
    FOR EACH ROW EXECUTE FUNCTION inbox_feed_stamp_revision();
"""

_TABLES = ("agent_messages", "message_receipts")


def upgrade() -> None:
    op.execute(sa.text("CREATE SEQUENCE inbox_revision_seq AS bigint"))
    for table in _TABLES:
        op.add_column(table, sa.Column("revision", sa.BigInteger(), nullable=True))
    # One sequence value for all existing rows: the pre-feed state is one revision.
    op.execute(
        sa.text(
            """
            DO $$
            DECLARE base bigint := nextval('inbox_revision_seq');
            BEGIN
                UPDATE agent_messages SET revision = base;
                UPDATE message_receipts SET revision = base;
            END
            $$
            """
        )
    )
    for table in _TABLES:
        op.alter_column(table, "revision", nullable=False)
        op.create_index(f"ix_{table}_revision", table, ["revision"])
    op.create_table(
        "message_tombstones",
        sa.Column("message_id", UUID(as_uuid=True), primary_key=True),
        sa.Column("revision", sa.BigInteger(), nullable=False),
        sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=False),
        sa.Column("reason", sa.Text(), nullable=True),
    )
    op.create_index("ix_message_tombstones_revision", "message_tombstones", ["revision"])
    op.execute(sa.text(FUNCTIONS_SQL))
    op.execute(sa.text(TRIGGERS_SQL))


def downgrade() -> None:
    op.execute(
        sa.text(
            """
            DROP TRIGGER IF EXISTS message_receipts_inbox_revision ON message_receipts;
            DROP TRIGGER IF EXISTS agent_messages_inbox_tombstone ON agent_messages;
            DROP TRIGGER IF EXISTS agent_messages_inbox_revision ON agent_messages;
            DROP FUNCTION IF EXISTS inbox_feed_tombstone();
            DROP FUNCTION IF EXISTS inbox_feed_stamp_revision();
            DROP FUNCTION IF EXISTS inbox_feed_revision();
            """
        )
    )
    op.drop_index("ix_message_tombstones_revision", table_name="message_tombstones")
    op.drop_table("message_tombstones")
    for table in reversed(_TABLES):
        op.drop_index(f"ix_{table}_revision", table_name=table)
        op.drop_column(table, "revision")
    op.execute(sa.text("DROP SEQUENCE IF EXISTS inbox_revision_seq"))