Source code for

"""Websockets outbound transport."""

import logging
from typing import Union

from aiohttp import ClientSession, DummyCookieJar

from ...config.injection_context import InjectionContext

from .base import BaseOutboundTransport

[docs]class WsTransport(BaseOutboundTransport): """Websockets outbound transport class.""" schemes = ("ws", "wss") def __init__(self) -> None: """Initialize an `WsTransport` instance.""" super(WsTransport, self).__init__() self.logger = logging.getLogger(__name__)
[docs] async def start(self): """Start the outbound transport.""" self.client_session = ClientSession(cookie_jar=DummyCookieJar()) return self
[docs] async def stop(self): """Stop the outbound transport.""" await self.client_session.close() self.client_session = None
[docs] async def handle_message( self, context: InjectionContext, payload: Union[str, bytes], endpoint: str ): """ Handle message from queue. Args: context: the context that produced the message payload: message payload in string or byte format endpoint: URI endpoint for delivery """ # aiohttp should automatically handle websocket sessions async with self.client_session.ws_connect(endpoint) as ws: if isinstance(payload, bytes): await ws.send_bytes(payload) else: await ws.send_str(payload)