fastapi-ws-rpc 0.1.0__py3-none-any.whl

fastapi_ws_rpc/__init__.py

@@ -0,0 +1,11 @@
+from fastapi_ws_rpc.rpc_methods import RpcMethodsBase
+from fastapi_ws_rpc.schemas import WebSocketFrameType
+from fastapi_ws_rpc.websocket_rpc_client import WebSocketRpcClient
+from fastapi_ws_rpc.websocket_rpc_endpoint import WebsocketRPCEndpoint
+
+__all__ = [
+    "RpcMethodsBase",
+    "WebSocketRpcClient",
+    "WebsocketRPCEndpoint",
+    "WebSocketFrameType",
+]

fastapi_ws_rpc/connection_manager.py

@@ -0,0 +1,15 @@
+from typing import List
+
+from fastapi import WebSocket
+
+
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections: List[WebSocket] = []
+
+    async def connect(self, websocket: WebSocket):
+        await websocket.accept()
+        self.active_connections.append(websocket)
+
+    def disconnect(self, websocket: WebSocket):
+        self.active_connections.remove(websocket)

fastapi_ws_rpc/logger.py

@@ -0,0 +1,113 @@
+import logging
+import os
+from enum import Enum
+from logging.config import dictConfig
+from typing import NewType
+
+ENV_VAR = "WS_RPC_LOGGING"
+
+
+class LoggingModes(Enum):
+    # don't produce logs
+    NO_LOGS = 0
+    # Log alongside uvicorn
+    UVICORN = 1
+    # Simple log calls (no config)
+    SIMPLE = 2
+    # log via the loguru module
+    LOGURU = 3
+
+
+LoggingMode = NewType("LoggingMode", LoggingModes)
+
+
+class LoggingConfig:
+    def __init__(self) -> None:
+        self._mode = None
+
+    config_template = {
+        "version": 1,
+        "disable_existing_loggers": True,
+        "formatters": {
+            "default": {
+                "()": "uvicorn.logging.DefaultFormatter",
+                "fmt": "%(levelprefix)s %(asctime)s %(message)s",
+                "datefmt": "%Y-%m-%d %H:%M:%S",
+            },
+        },
+        "handlers": {
+            "default": {
+                "formatter": "default",
+                "class": "logging.StreamHandler",
+            },
+        },
+        "loggers": {},
+    }
+
+    UVICORN_LOGGERS = {
+        "uvicorn.error": {
+            "propagate": False,
+            "handlers": ["default"],
+        },
+        "fastapi_ws_rpc": {
+            "handlers": ["default"],
+            "propagate": False,
+            "level": logging.INFO,
+        },
+    }
+
+    def get_mode(self):
+        # if no one set the mode - set default from ENV or hardcoded default
+        if self._mode is None:
+            mode = LoggingModes.__members__.get(
+                os.environ.get(ENV_VAR, "").upper(), LoggingModes.SIMPLE
+            )
+            self.set_mode(mode)
+        return self._mode
+
+    def set_mode(self, mode: LoggingMode = LoggingModes.UVICORN, level=logging.INFO):
+        """
+        Configure logging. this method calls 'logging.config.dictConfig()' to enable
+        quick setup of logging. Call this method before starting the app.
+        For more advanced cases use 'logging.config' directly (loggers used by this
+        library are all nested under "fastapi_ws_rpc" logger name)
+
+        Args:
+            mode (LoggingMode, optional): The mode to set logging to. Defaults to
+            LoggingModes.UVICORN.
+        """
+        self._mode = mode
+        logging_config = self.config_template.copy()
+        # add logs beside uvicorn
+        if mode == LoggingModes.UVICORN:
+            logging_config["loggers"] = self.UVICORN_LOGGERS.copy()
+            logging_config["loggers"]["fastapi_ws_rpc"]["level"] = level
+            dictConfig(logging_config)
+        elif mode == LoggingModes.SIMPLE:
+            pass
+        elif mode == LoggingModes.LOGURU:
+            pass
+        # no logs
+        else:
+            logging_config["loggers"] = {}
+            dictConfig(logging_config)
+
+
+# Singelton for logging configuration
+logging_config = LoggingConfig()
+
+
+def get_logger(name):
+    """
+    Get a logger object to log with.
+    Called by inner modules for logging.
+    """
+    mode = logging_config.get_mode()
+    # logging through loguru
+    if mode == LoggingModes.LOGURU:
+        from loguru import logger
+    else:
+        # regular python logging
+        logger = logging.getLogger(f"fastapi_ws_rpc.{name}")
+
+    return logger

fastapi_ws_rpc/rpc_channel.py

@@ -0,0 +1,482 @@
+"""
+Definition for an RPC channel protocol on top of a websocket -
+enabling bi-directional request/response interactions
+"""
+import asyncio
+from inspect import _empty, signature
+from typing import Any, Dict, List
+
+from pydantic import ValidationError
+
+from .logger import get_logger
+from .rpc_methods import EXPOSED_BUILT_IN_METHODS, NoResponse, RpcMethodsBase
+from .schemas import RpcMessage, RpcRequest, RpcResponse
+from .utils import gen_uid, pydantic_parse
+
+logger = get_logger("RPC_CHANNEL")
+
+
+class DEFAULT_TIMEOUT:
+    pass
+
+
+class RemoteValueError(ValueError):
+    pass
+
+
+class RpcException(Exception):
+    pass
+
+
+class RpcChannelClosedException(Exception):
+    """
+    Raised when the channel is closed mid-operation
+    """
+
+    pass
+
+
+class UnknownMethodException(RpcException):
+    pass
+
+
+class RpcPromise:
+    """
+    Simple Event and id wrapper/proxy
+    Holds the state of a pending request
+    """
+
+    def __init__(self, request: RpcRequest):
+        self._request = request
+        self._id = request.call_id
+        # event used to wait for the completion of the request
+        # (upon receiving its matching response)
+        self._event = asyncio.Event()
+
+    @property
+    def request(self):
+        return self._request
+
+    @property
+    def call_id(self):
+        return self._id
+
+    def set(self):
+        """
+        Signal completion of request with received response
+        """
+        self._event.set()
+
+    def wait(self):
+        """
+        Wait on the internal event - triggered on response
+        """
+        return self._event.wait()
+
+
+class RpcProxy:
+    """Provides a proxy to a remote method on the other side of the channel."""
+
+    def __init__(self, channel, method_name: str) -> None:
+        self.channel = channel
+        self.method_name = method_name
+
+    def __call__(self, **kwargs: Any) -> Any:
+        """Calls the remote method with the given keyword arguments.
+
+        Parameters
+        ----------
+        kwargs : dict
+            Keyword arguments to pass to the remote method.
+
+        Returns
+        -------
+        Any
+            Return value of the remote method.
+        """
+        logger.debug("Calling RPC method: %s", self.method_name)
+        return self.channel.call(self.method_name, args=kwargs)
+
+
+class RpcCaller:
+    """Calls remote methods on the other side of the channel."""
+
+    def __init__(self, channel: "RpcChannel", methods=None) -> None:
+        self._channel = channel
+
+    def __getattribute__(self, name: str):
+        """Returns an :class:`~fastapi_ws_rpc.rpc_channel.RpcProxy` instance
+        for exposed RPC methods.
+
+        Parameters
+        ----------
+        name : str
+            Name of the requested attribute (or RPC method).
+
+        Returns
+        -------
+        Any
+            Attribute or RPC method.
+        """
+        is_exposed = not name.startswith("_") or name in EXPOSED_BUILT_IN_METHODS
+        if is_exposed:
+            logger.debug("%s was detected to be a remote RPC method.", name)
+            return RpcProxy(self._channel, name)
+        return super().__getattribute__(name)
+
+
+# Callback signatures
+async def OnConnectCallback(channel):
+    pass
+
+
+async def OnDisconnectCallback(channel):
+    pass
+
+
+async def OnErrorCallback(channel, err: Exception):
+    pass
+
+
+class RpcChannel:
+    """
+    A wire agnostic json-rpc channel protocol for both server and client.
+    Enable each side to send RPC-requests (calling exposed methods on other side)
+    and receive rpc-responses with the return value
+
+    provides a .other property for calling remote methods.
+    e.g. answer = channel.other.add(a=1,b=1) will (For example) ask the other
+    side to perform 1+1 and will return an RPC-response of 2
+    """
+
+    def __init__(
+        self,
+        methods: RpcMethodsBase,
+        socket,
+        channel_id=None,
+        default_response_timeout=None,
+        sync_channel_id=False,
+        **kwargs,
+    ):
+        """
+
+        Args:
+            methods (RpcMethodsBase): RPC methods to expose to other side
+            socket: socket object providing simple send/recv methods
+            channel_id (str, optional): uuid for channel. Defaults to None in
+            which case a random UUID is generated.
+            default_response_timeout(float, optional) default timeout for RPC
+            call responses. Defaults to None - i.e. no timeout
+            sync_channel_id(bool, optional) should get the other side of the
+            channel id, helps to identify connections, cost a bit networking time.
+                Defaults to False - i.e. not getting the other side channel id
+        """
+        logger.debug("Initializing RPC channel...")
+        self.methods = methods._copy_()
+        # allow methods to access channel (for recursive calls - e.g. call me as
+        # a response for me calling you)
+        self.methods._set_channel_(self)
+        # Pending requests - id-mapped to async-event
+        self.requests: Dict[str, asyncio.Event] = {}
+        # Received responses
+        self.responses = {}
+        self.socket = socket
+        # timeout
+        self.default_response_timeout = default_response_timeout
+        # Unique channel id
+        self.id = channel_id if channel_id is not None else gen_uid()
+        # flag control should we retrieve the channel id of the other side
+        self._sync_channel_id = sync_channel_id
+        # The channel id of the other side (if sync_channel_id is True)
+        self._other_channel_id = None
+        # asyncio event to check if we got the channel id of the other side
+        self._channel_id_synced = asyncio.Event()
+        #
+        # convenience caller
+        # TODO - pass remote methods object to support validation before call
+        self.other = RpcCaller(self)
+        # core event callback registers
+        self._connect_handlers: List[OnConnectCallback] = []
+        self._disconnect_handlers: List[OnDisconnectCallback] = []
+        self._error_handlers = []
+        # internal event
+        self._closed = asyncio.Event()
+
+        # any other kwarg goes straight to channel context (Accessible to methods)
+        self._context = kwargs or {}
+        logger.debug("RPC channel initialized.")
+
+    @property
+    def context(self) -> Dict[str, Any]:
+        return self._context
+
+    async def get_other_channel_id(self) -> str:
+        """
+        Method to get the channel id of the other side of the channel
+        The _channel_id_synced verify we have it
+        Timeout exception can be raised if the value isn't available
+        """
+        await asyncio.wait_for(
+            self._channel_id_synced.wait(), self.default_response_timeout
+        )
+        return self._other_channel_id
+
+    def get_return_type(self, method):
+        method_signature = signature(method)
+        return (
+            method_signature.return_annotation
+            if method_signature.return_annotation is not _empty
+            else str
+        )
+
+    async def send(self, data):
+        """
+        For internal use. wrap calls to underlying socket
+        """
+        await self.socket.send(data)
+
+    async def receive(self):
+        """
+        For internal use. wrap calls to underlying socket
+        """
+        return await self.socket.recv()
+
+    async def close(self):
+        res = await self.socket.close()
+        # signal closer
+        self._closed.set()
+        return res
+
+    def isClosed(self):
+        return self._closed.is_set()
+
+    async def wait_until_closed(self):
+        """
+        Waits until the close internal event happens.
+        """
+        return await self._closed.wait()
+
+    async def on_message(self, data):
+        """Handle an incoming RPC message."""
+        logger.debug(f"Processing received message: {data}")
+        try:
+            message = pydantic_parse(RpcMessage, data)
+            if message.request is not None:
+                await self.on_request(message.request)
+            if message.response is not None:
+                await self.on_response(message.response)
+        except ValidationError as e:
+            logger.error("Failed to parse message", {"message": data, "error": e})
+            await self.on_error(e)
+        except Exception as e:
+            await self.on_error(e)
+            raise
+
+    def register_connect_handler(self, coros: List[OnConnectCallback] = None):
+        """
+        Register a connection handler callback that will be called (As an async
+        task)) with the channel
+        Args:
+            coros (List[Coroutine]): async callback
+        """
+        if coros is not None:
+            self._connect_handlers.extend(coros)
+
+    def register_disconnect_handler(self, coros: List[OnDisconnectCallback] = None):
+        """
+        Register a disconnect handler callback that will be called (As an async
+        task)) with the channel id
+        Args:
+            coros (List[Coroutine]): async callback
+        """
+        if coros is not None:
+            self._disconnect_handlers.extend(coros)
+
+    def register_error_handler(self, coros: List[OnErrorCallback] = None):
+        """
+        Register an error handler callback that will be called (As an async
+        task)) with the channel and triggered error.
+        Args:
+            coros (List[Coroutine]): async callback
+        """
+        if coros is not None:
+            self._error_handlers.extend(coros)
+
+    async def on_handler_event(self, handlers, *args, **kwargs):
+        await asyncio.gather(*(callback(*args, **kwargs) for callback in handlers))
+
+    async def on_connect(self):
+        """
+        Run get other channel id if sync_channel_id is True
+        Run all callbacks from self._connect_handlers
+        """
+        if self._sync_channel_id:
+            logger.debug("Syncing channel ID...")
+            self._get_other_channel_id_task = asyncio.create_task(
+                self._get_other_channel_id()
+            )
+        await self.on_handler_event(self._connect_handlers, self)
+
+    async def _get_other_channel_id(self):
+        """
+        Perform call to the other side of the channel to get its channel id
+        Each side is generating the channel id by itself so there is no way
+        to identify a connection without this sync
+        """
+        if self._other_channel_id is None:
+            logger.debug("No cached channel ID found, calling _get_channel_id_()...")
+            other_channel_id = await self.other._get_channel_id_()
+            self._other_channel_id = (
+                other_channel_id.result
+                if other_channel_id and other_channel_id.result
+                else None
+            )
+            logger.debug("Got channel ID: %s", self._other_channel_id)
+            if self._other_channel_id is None:
+                raise RemoteValueError()
+            # update asyncio event that we received remote channel id
+            self._channel_id_synced.set()
+            return self._other_channel_id
+        return self._other_channel_id
+
+    async def on_disconnect(self):
+        # disconnect happened - mark the channel as closed
+        self._closed.set()
+        await self.on_handler_event(self._disconnect_handlers, self)
+
+    async def on_error(self, error: Exception):
+        await self.on_handler_event(self._error_handlers, self, error)
+
+    async def on_request(self, message: RpcRequest):
+        """
+        Handle incoming RPC requests - calling relevant exposed method
+        Note: methods prefixed with "_" are protected and ignored.
+
+        Args:
+            message (RpcRequest): the RPC request with the method to call
+        """
+        # TODO add exception support (catch exceptions and pass to other side
+        # as response with errors)
+        logger.debug(f"Handling RPC request on channel {self.id}: {message}")
+        method_name = message.method
+        # Ignore "_" prefixed methods (except the built in "_ping_")
+        if isinstance(method_name, str) and (
+            not method_name.startswith("_") or method_name in EXPOSED_BUILT_IN_METHODS
+        ):
+            method = getattr(self.methods, method_name)
+            if callable(method):
+                result = await method(**message.arguments)
+                if result is not NoResponse:
+                    # get indicated type
+                    result_type = self.get_return_type(method)
+                    # if no type given - try to convert to string
+                    if result_type is str and not isinstance(result, str):
+                        result = str(result)
+                    response = RpcMessage(
+                        request=RpcRequest(
+                            method=method_name,
+                            arguments=message.arguments,
+                            call_id=message.call_id,
+                        ),
+                        response=RpcResponse[result_type](
+                            call_id=message.call_id,
+                            result=result,
+                            result_type=getattr(
+                                result_type,
+                                "__name__",
+                                getattr(result_type, "_name", "unknown-type"),
+                            ),
+                        ),
+                    )
+                    await self.send(response)
+
+    def get_saved_promise(self, call_id):
+        return self.requests[call_id]
+
+    def get_saved_response(self, call_id):
+        return self.responses[call_id]
+
+    def clear_saved_call(self, call_id):
+        del self.requests[call_id]
+        del self.responses[call_id]
+
+    async def on_response(self, response: RpcResponse):
+        """
+        Handle an incoming response to a previous RPC call
+
+        Args:
+            response (RpcResponse): the received response
+        """
+        logger.debug("Handling RPC response - %s", {"response": response})
+        if response.call_id is not None and response.call_id in self.requests:
+            self.responses[response.call_id] = response
+            promise = self.requests[response.call_id]
+            promise.set()
+
+    async def wait_for_response(self, promise, timeout=DEFAULT_TIMEOUT) -> RpcResponse:
+        """
+        Wait on a previously made call
+        Args:
+            promise (RpcPromise): the awaitable-wrapper returned from the RPC
+            request call
+            timeout (float, None, or DEFAULT_TIMEOUT): the timeout to wait on
+            the response, defaults to DEFAULT_TIMEOUT.
+                - DEFAULT_TIMEOUT - use the value passed as
+                'default_response_timeout' in channel init
+                - None - no timeout
+                - a number - seconds to wait before timing out
+        Raises:
+            asyncio.exceptions.TimeoutError - on timeout
+            RpcChannelClosedException - if the channel fails before wait could
+            be completed
+        """
+        if timeout is DEFAULT_TIMEOUT:
+            timeout = self.default_response_timeout
+        # wait for the promise or until the channel is terminated
+        _, pending = await asyncio.wait(
+            [
+                asyncio.ensure_future(promise.wait()),
+                asyncio.ensure_future(self._closed.wait()),
+            ],
+            timeout=timeout,
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+        # Cancel all pending futures and then detect if close was the first done
+        for fut in pending:
+            fut.cancel()
+        response = self.responses.get(promise.call_id, NoResponse)
+        # if the channel was closed before we could finish
+        if response is NoResponse:
+            raise RpcChannelClosedException(
+                f"Channel Closed before RPC response for {promise.call_id} could be received"  # noqa: E501
+            )
+        self.clear_saved_call(promise.call_id)
+        return response
+
+    async def async_call(self, name, args={}, call_id=None) -> RpcPromise:
+        """
+        Call a method and return the event and the sent message (including the
+        chosen call_id)
+        use self.wait_for_response on the event and call_id to get the return
+        value of the call
+        Args:
+            name (str): name of the method to call on the other side
+            args (dict): keyword args to pass tot he other side
+            call_id (string, optional): a UUID to use to track the call () -
+            override only with true UUIDs
+        """
+        call_id = call_id or gen_uid()
+        message = RpcMessage(
+            request=RpcRequest(method=name, arguments=args, call_id=call_id)
+        )
+        logger.debug("Sending the following RPC message: %s", message)
+        await self.send(message)
+        promise = self.requests[message.request.call_id] = RpcPromise(message.request)
+        return promise
+
+    async def call(self, name, args={}, timeout=DEFAULT_TIMEOUT):
+        """
+        Call a method and wait for a response to be received
+        """
+        promise = await self.async_call(name, args)
+        return await self.wait_for_response(promise, timeout=timeout)