felicity-lims/felicity/api/gql/notification/subscription.py
2023-04-10 14:23:31 +02:00

34 lines
1.1 KiB
Python

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)