SQLAlchemy + Postgres Cheatsheet 08 — LISTEN / NOTIFY

async with session.begin():
    session.add(user)
    await session.execute(text("NOTIFY user_events, :p"), {"p": json.dumps({"id": user.id})})

import asyncpg

async def listen():
    conn = await asyncpg.connect(DATABASE_URL)
    await conn.add_listener("user_events", on_notify)
    while True:
        await asyncio.sleep(60)        # keep connection alive

async def on_notify(connection, pid, channel, payload):
    data = json.loads(payload)
    log.info("got_event", channel=channel, data=data)

# Per-tenant
await conn.execute(f"LISTEN tenant_{tenant_id}")

# By type
await conn.execute("LISTEN orders")
await conn.execute("LISTEN invoices")

CREATE OR REPLACE FUNCTION notify_outbox() RETURNS TRIGGER AS $$
BEGIN
    PERFORM pg_notify('outbox', '');
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER outbox_notify
AFTER INSERT ON outbox_events
FOR EACH ROW EXECUTE FUNCTION notify_outbox();

async def outbox_worker():
    conn = await asyncpg.connect(DATABASE_URL)
    await conn.add_listener("outbox", lambda *_: asyncio.create_task(drain()))
    while True:
        await drain()                  # periodic safety drain
        await asyncio.sleep(60)

async def drain():
    async with AsyncSessionLocal() as s:
        async with s.begin():
            events = (await s.execute(
                select(OutboxEvent)
                .where(OutboxEvent.published_at == None)
                .limit(100)
                .with_for_update(skip_locked=True)
            )).scalars().all()
            for e in events:
                await bus.publish(e.event_type, e.payload)
                e.published_at = datetime.utcnow()

await session.execute(text("NOTIFY user_events, :id"), {"id": str(user.id)})

NOTIFY my_channel;

await conn.execute("NOTIFY my_channel")

CREATE OR REPLACE FUNCTION notify_user_change() RETURNS TRIGGER AS $$
BEGIN
    PERFORM pg_notify('user_changes', json_build_object('op', TG_OP, 'id', NEW.id)::text);
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER user_change_trigger
AFTER INSERT OR UPDATE OR DELETE ON users
FOR EACH ROW EXECUTE FUNCTION notify_user_change();

@asynccontextmanager
async def lifespan(app):
    listen_task = asyncio.create_task(listen())
    yield
    listen_task.cancel()
    try: await listen_task
    except asyncio.CancelledError: pass

async def keepalive():
    while True:
        await conn.execute("SELECT 1")
        await asyncio.sleep(60)

conn = await asyncpg.connect(
    DATABASE_URL,
    server_settings={"tcp_keepalives_idle": "60"},
)