@camera.ui/server 0.0.48 → 0.0.50-internal.0

package/dist/server/src/_camera_ui_tools/camera_ui_rpc_client/base/client.py

@@ -0,0 +1,256 @@
+from __future__ import annotations
+
+import inspect
+import logging
+import traceback
+from collections.abc import Awaitable
+from typing import Any, Callable, TypeVar, cast
+
+from ..constants import SUBJECTS
+from ..typings import (
+    ClientConnectionOptions,
+    ClientMessage,
+    ConnectionState,
+    Message,
+    Metadata,
+    ResponseMessage,
+    RPCOptions,
+)
+from ..utils.logger import setup_logger
+from ..utils.routing import EndpointType, MessageRouter
+from ..utils.subscriptions import SubscriptionWrapper
+from .base import RPCBase
+from .connection import RPCConnectionManager
+from .errors import ERROR_CODES, ErrorFactory, NatsError
+from .rpc import RPCConnection
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+logger: logging.Logger = setup_logger("camera:ui:client")
+
+
+class RPCClient(RPCBase):
+    server_metadata: Metadata | None
+
+    _server_status: ConnectionState
+    _registered: bool
+    _is_proxy: bool
+
+    def __init__(
+        self,
+        type_: EndpointType,
+        options: ClientConnectionOptions,
+        is_proxy: bool = False,
+    ):
+        super().__init__(type_, options)
+        self._server_status = ConnectionState.STOPPED
+        self._registered = False
+        self._is_proxy = is_proxy
+        self.server_metadata = None
+
+    async def connect(self) -> None:
+        await self._initialize()
+        await self._setup_subscriptions()
+        logger.debug("Connected successfully")
+
+    async def disconnect(self) -> None:
+        await self._cleanup()
+        self._registered = False
+        logger.debug("Disconnected from NATS")
+
+    async def send_to_server(self, data: Any) -> None:
+        if not self.server_metadata:
+            raise ErrorFactory.client.invalid_state(
+                None,
+                {
+                    "description": "Server metadata is not available",
+                    "code": ERROR_CODES.CLIENT.INVALID_STATE,
+                },
+            )
+
+        subject = (
+            MessageRouter.get_plugin_to_server_subject(
+                self.options["metadata"]["id"], self.server_metadata["id"]
+            )
+            if self.type_ == "plugin"
+            else MessageRouter.get_service_to_server_subject(
+                self.options["metadata"]["id"], self.server_metadata["id"]
+            )
+        )
+        routing = MessageRouter.create_routing_metadata(
+            self.type_, self.options["metadata"]["id"], "server", self.server_metadata["id"]
+        )
+        return await self.publish(subject, data, routing)
+
+    async def request_from_server(self, data: Any, timeout: int | None = None) -> ResponseMessage[Any]:
+        if not self.server_metadata:
+            raise ErrorFactory.client.invalid_state(
+                None,
+                {
+                    "description": "Server metadata is not available",
+                    "code": ERROR_CODES.CLIENT.INVALID_STATE,
+                },
+            )
+
+        subject = (
+            MessageRouter.get_plugin_to_server_subject(
+                self.options["metadata"]["id"], self.server_metadata["id"]
+            )
+            if self.type_ == "plugin"
+            else MessageRouter.get_service_to_server_subject(
+                self.options["metadata"]["id"], self.server_metadata["id"]
+            )
+        )
+        routing = MessageRouter.create_routing_metadata(
+            self.type_, self.options["metadata"]["id"], "server", self.server_metadata["id"]
+        )
+        return await self.request(subject, data, timeout, routing)
+
+    async def on_server_message(self, handler: Callable[[T], R | Awaitable[R]]) -> SubscriptionWrapper:
+        if not self.server_metadata:
+            raise ErrorFactory.client.invalid_state(
+                None,
+                {
+                    "description": "Server metadata is not available",
+                    "code": ERROR_CODES.CLIENT.INVALID_STATE,
+                },
+            )
+
+        async def wrapped_handler(message: Message[T]) -> R:
+            if inspect.iscoroutinefunction(handler) or inspect.isawaitable(handler):
+                return await cast(Awaitable[R], handler(message["data"]))
+            else:
+                return cast(R, handler(message["data"]))
+
+        direct_subject = (
+            MessageRouter.get_server_to_plugin_subject(
+                self.server_metadata["id"], self.options["metadata"]["id"]
+            )
+            if self.type_ == "plugin"
+            else MessageRouter.get_server_to_service_subject(
+                self.server_metadata["id"], self.options["metadata"]["id"]
+            )
+        )
+
+        return await self.handle_request(direct_subject, wrapped_handler)
+
+    def is_server_running(self) -> bool:
+        return self._server_status == ConnectionState.RUNNING
+
+    def is_connected(self) -> bool:
+        return self._connection.is_connected()
+
+    def is_registered(self) -> bool:
+        return self._registered
+
+    async def create_rpc_proxy(
+        self, cls: type[T], namespace: str, rpc_options: RPCOptions | None = None
+    ) -> T:
+        return await super().create_rpc_proxy(cls, namespace, rpc_options)
+
+    async def create_rpc_connection(
+        self, namespace: str, rpc_options: RPCOptions | None = None
+    ) -> RPCConnection:
+        return await super().create_rpc_connection(namespace, rpc_options)
+
+    async def disconnect_rpc_connection(self, namespace: str) -> None:
+        return await super().disconnect_rpc_connection(namespace)
+
+    def get_connection_manager(self) -> RPCConnectionManager:
+        return self._connection
+
+    def get_connection_managers(self) -> list[RPCConnectionManager]:
+        rrpc_connection_managers: list[RPCConnectionManager] = []
+        for rpc_manager in self._rpc_managers.values():
+            rrpc_connection_managers.append(rpc_manager.get_connection_manager())
+
+        return [self._connection, *rrpc_connection_managers]
+
+    async def _initialize(self) -> None:
+        await self._connection.connect()
+        await self._register()
+
+    async def _setup_subscriptions(self) -> None:
+        if not self._is_proxy:
+            await self.handle_request(
+                f"{SUBJECTS.HEALTH.PING}.{self.options['metadata']['id']}", self._handle_ping
+            )
+
+        await self.subscribe(SUBJECTS.SERVER.STATUS, self._handle_server_status)
+
+    async def _handle_ping(self, event: Message[Any]) -> ClientMessage:
+        return {"clientId": self.options["metadata"]["id"]}
+
+    async def _handle_server_status(self, event: Message[ConnectionState]) -> None:
+        logger.debug(f"Server status update: {event}")
+        self._server_status = event["data"]
+
+        if self._server_status in (ConnectionState.STOPPING, ConnectionState.STOPPED):
+            self._registered = False
+        elif self._server_status == ConnectionState.RUNNING and not self.is_registered():
+            if self._is_proxy:
+                self._registered = True
+                self._server_status = ConnectionState.RUNNING
+            else:
+                await self._register()
+
+    async def _cleanup(self) -> None:
+        await super().cleanup()
+        await self._connection.disconnect()
+
+    async def _register(self) -> None:
+        if self.is_registered() or self._is_proxy:
+            return
+
+        SUBJECT = getattr(SUBJECTS, self.type_.upper()).REGISTER
+
+        try:
+            response = await self.request(SUBJECT, self.options["metadata"])
+
+            if "error" in response and response["error"]:
+                raise NatsError.from_json(response["error"])
+
+            if not response["data"]["success"]:
+                raise ErrorFactory.client.registration_failed(
+                    None,
+                    {
+                        "description": "Failed to register client",
+                        "subject": SUBJECT,
+                        "code": ERROR_CODES.CLIENT.REGISTRATION_FAILED,
+                        "metadata": {
+                            "response": response,
+                            "isProxy": self._is_proxy,
+                            "registered": self._registered,
+                            "type": self.type_,
+                            "metadata": self.options["metadata"],
+                        },
+                    },
+                )
+
+            self.server_metadata = response["data"]["serverMetadata"]
+            self._registered = True
+            self._server_status = ConnectionState.RUNNING
+
+            logger.debug("Registration successful")
+
+        except Exception as e:
+            self._registered = False
+
+            print(e)
+            traceback.print_exc()
+
+            raise ErrorFactory.client.registration_failed(
+                e,
+                {
+                    "description": "Failed to register client",
+                    "code": ERROR_CODES.CLIENT.REGISTRATION_FAILED,
+                    "subject": SUBJECT,
+                    "metadata": {
+                        "isProxy": self._is_proxy,
+                        "registered": self._registered,
+                        "type": self.type_,
+                        "metadata": self.options["metadata"],
+                    },
+                },
+            ) from e