felicity-lims/felicity/api/gql/notification/subscription.py

import asyncio
import logging
from typing import AsyncGenerator

import strawberry  # noqa
from api.gql.notification.types import ActivityStreamType
from apps.common.channel import BroadcastEvent, Subscriber, broadcast
from apps.notification.models import ActivityStream

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@strawberry.type
class StreamSubscription:
    @strawberry.subscription
    async def latest_activity(self) -> AsyncGenerator[ActivityStreamType, None]:  # noqa
        subscriber: Subscriber
        async with broadcast.subscribe(channel="activities") as subscriber:
            logger.info("Subscribed")
            event: BroadcastEvent
            try:
                async for event in subscriber:
                    logger.info(event)
                    yield event.message
            finally:
                logger.info("Unsubscribed")

    @strawberry.subscription
    async def stream_all(self) -> AsyncGenerator[ActivityStreamType, None]:  # noqa
        streams = await ActivityStream.all()
        for stream in streams:
            yield stream
            await asyncio.sleep(1)