56 lines
2.1 KiB
Python
56 lines
2.1 KiB
Python
"""Redis-backed notification registry for multi-pod deployments.
|
|
|
|
Thin adapter around :class:`RedisPubSub` — all Redis pub/sub logic
|
|
lives in the generic layer; this class maps the ``organization_id``
|
|
domain concept onto generic string keys.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
from redis.asyncio import Redis
|
|
|
|
from skyvern.forge.sdk.notification.base import BaseNotificationRegistry
|
|
from skyvern.forge.sdk.redis.pubsub import RedisPubSub
|
|
|
|
|
|
class RedisNotificationRegistry(BaseNotificationRegistry):
|
|
"""Fan-out pub/sub backed by Redis. One Redis PubSub channel per org."""
|
|
|
|
def __init__(self, redis_client: Redis) -> None:
|
|
self._pubsub = RedisPubSub(redis_client, channel_prefix="skyvern:notifications:")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Property accessors (used by existing tests)
|
|
# ------------------------------------------------------------------
|
|
|
|
@property
|
|
def _listener_tasks(self) -> dict[str, asyncio.Task[None]]:
|
|
return self._pubsub._listener_tasks
|
|
|
|
@property
|
|
def _subscribers(self) -> dict[str, list[asyncio.Queue[dict]]]:
|
|
return self._pubsub._subscribers
|
|
|
|
# ------------------------------------------------------------------
|
|
# Public interface
|
|
# ------------------------------------------------------------------
|
|
|
|
def subscribe(self, organization_id: str) -> asyncio.Queue[dict]:
|
|
return self._pubsub.subscribe(organization_id)
|
|
|
|
def unsubscribe(self, organization_id: str, queue: asyncio.Queue[dict]) -> None:
|
|
self._pubsub.unsubscribe(organization_id, queue)
|
|
|
|
def publish(self, organization_id: str, message: dict) -> None:
|
|
self._pubsub.publish(organization_id, message)
|
|
|
|
async def close(self) -> None:
|
|
await self._pubsub.close()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Internal helper (exposed for tests)
|
|
# ------------------------------------------------------------------
|
|
|
|
def _dispatch_local(self, organization_id: str, message: dict) -> None:
|
|
self._pubsub._dispatch_local(organization_id, message)
|