zepben.ewb 1.0.0b1__py3-none-any.whl

zepben/ewb/streaming/grpc/grpc_channel_builder.py

@@ -0,0 +1,185 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ["GrpcChannelBuilder"]
+
+from abc import ABC
+from typing import Optional, List, Tuple, Any, Dict
+
+import grpc
+from grpc._channel import _InactiveRpcError
+from zepben.ewb.auth import ZepbenTokenFetcher
+from zepben.protobuf.cc.cc_pb2_grpc import CustomerConsumerStub
+from zepben.protobuf.connection.connection_requests_pb2 import CheckConnectionRequest
+from zepben.protobuf.dc.dc_pb2_grpc import DiagramConsumerStub
+from zepben.protobuf.nc.nc_pb2_grpc import NetworkConsumerStub
+from zepben.protobuf.ns.network_state_pb2_grpc import QueryNetworkStateServiceStub, UpdateNetworkStateServiceStub
+
+from zepben.ewb.streaming.exceptions import GrpcConnectionException
+from zepben.ewb.streaming.grpc.auth_token_plugin import AuthTokenPlugin
+
+_TWENTY_MEGABYTES = 1024 * 1024 * 20
+
+
+class GrpcChannelBuilder(ABC):
+    """
+    Builder class for gRPC channels to connect to EWB's gRPC service.
+    """
+
+    _stubs: Dict = {
+        "NetworkConsumerClient": NetworkConsumerStub,
+        "DiagramConsumerClient": DiagramConsumerStub,
+        "CustomerConsumerClient": CustomerConsumerStub,
+        "QueryNetworkStateClient": QueryNetworkStateServiceStub,
+        "UpdateNetworkStateClient": UpdateNetworkStateServiceStub
+    }
+
+    def __init__(self):
+        self._socket_address: str = "localhost:50051"
+        self._channel_credentials: Optional[grpc.ChannelCredentials] = None
+
+    def build(self, skip_connection_test: bool = False, timeout_seconds: int = 5, debug: bool = False,
+              options: Optional[List[Tuple[str, Any]]] = None) -> grpc.aio.Channel:
+        """
+        Get the resulting :class:`grpc.aio.Channel` from this builder.
+
+        :param skip_connection_test: Skip confirming a connection can be established to the server. This is not recommended, but provided as a safety
+        mechanism if for any reason the connection test fails unexpectedly even though the connection is fine.
+        :param timeout_seconds: The timeout used for each request made in the connection test.
+        :param debug: Collect and append unhandled RPC errors to the `ConnectionException` raised on an unsuccessful connection test.
+        :param options: An optional list of key-value pairs (channel_arguments in gRPC Core runtime) to configure the channel.
+
+        :return: A gRPC channel resulting from this builder.
+        """
+
+        # Synchronous channel is used purely for confirming channel connection works as we can't do async requests here,
+        # but an aio channel is always returned to the user.
+        if not skip_connection_test:
+            if self._channel_credentials:
+                channel = grpc.secure_channel(self._socket_address, self._channel_credentials)
+            else:
+                channel = grpc.insecure_channel(self._socket_address)
+            self._test_connection(channel, debug=debug, timeout_seconds=timeout_seconds)
+            channel.close()
+
+        if options is None:
+            options = [("grpc.max_receive_message_length", _TWENTY_MEGABYTES), ("grpc.max_send_message_length", _TWENTY_MEGABYTES)]
+        else:
+            has_max_receive_msg_length = False
+            has_max_send_msg_length = False
+            for key, _ in options:
+                if key == "grpc.max_receive_message_length":
+                    has_max_receive_msg_length = True
+                if key == "grpc.max_send_message_length":
+                    has_max_send_msg_length = True
+            if not has_max_send_msg_length:
+                options.append(("grpc.max_send_message_length", _TWENTY_MEGABYTES))
+            if not has_max_receive_msg_length:
+                options.append(("grpc.max_receive_message_length", _TWENTY_MEGABYTES))
+
+        if self._channel_credentials:
+            return grpc.aio.secure_channel(self._socket_address, self._channel_credentials, options=options)
+
+        return grpc.aio.insecure_channel(self._socket_address, options=options)
+
+    def _test_connection(self, channel: grpc.Channel, debug: bool, timeout_seconds: int):
+        debug_errors = []
+        for name, client in zip(self._stubs.keys(), map(lambda stub: stub(channel), list(self._stubs.values()))):
+            try:
+                result = client.checkConnection(CheckConnectionRequest(), timeout=timeout_seconds, wait_for_ready=False)
+                if result:
+                    return
+            except _InactiveRpcError as rpc_error:
+                debug_errors.append(f"Received the following exception with {name}:\n{rpc_error}\n")
+        raise GrpcConnectionException(f"Couldn't establish gRPC connection to any service on {self._socket_address}.\n{''.join(debug_errors) if debug else ''}")
+
+    def for_address(self, host: str, port: int) -> 'GrpcChannelBuilder':
+        """
+        Specify the address for the gRPC channel.
+
+        :param host: The hostname hosting the gRPC service
+        :param port: The port of the gRPC service
+        :return: This builder
+        """
+        self._socket_address = f"{host}:{port}"
+        return self
+
+    def make_secure(
+        self,
+        root_certificates: Optional[str] = None,
+        certificate_chain: Optional[str] = None,
+        private_key: Optional[str] = None
+    ) -> 'GrpcChannelBuilder':
+        """
+        Secures channel with SSL credentials.
+
+        :param root_certificates: The filename of the truststore to use when verifying the RPC service's SSL/TLS certificate
+        :param certificate_chain: The filename of the certificate chain to use for client authentication
+        :param private_key: The filename of the private key to use for client authentication
+        :return: This builder
+        """
+        root_certificates_bytes = None
+        if root_certificates is not None:
+            with open(root_certificates, "rb") as f:
+                root_certificates_bytes = f.read()
+
+        certificate_chain_bytes = None
+        if certificate_chain is not None:
+            with open(certificate_chain, "rb") as f:
+                certificate_chain_bytes = f.read()
+
+        private_key_bytes = None
+        if private_key is not None:
+            with open(private_key, "rb") as f:
+                private_key_bytes = f.read()
+
+        return self.make_secure_with_bytes(root_certificates_bytes, certificate_chain_bytes, private_key_bytes)
+
+    def make_secure_with_bytes(
+        self,
+        root_certificates_bytes: Optional[bytes] = None,
+        certificate_chain_bytes: Optional[bytes] = None,
+        private_key_bytes: Optional[bytes] = None
+    ) -> 'GrpcChannelBuilder':
+        """
+        Secures channel with SSL credentials.
+
+        :param root_certificates_bytes: The bytestring truststore to use when verifying the RPC service's SSL/TLS certificate
+        :param certificate_chain_bytes: The bytestring certificate chain to use for client authentication
+        :param private_key_bytes: The bytestring private key to use for client authentication
+        :return: This builder
+        """
+        self._channel_credentials = grpc.ssl_channel_credentials(root_certificates_bytes, private_key_bytes, certificate_chain_bytes)
+        return self
+
+    def with_token_fetcher(self, token_fetcher: ZepbenTokenFetcher) -> 'GrpcChannelBuilder':
+        """
+        Authenticates calls for the gRPC channel using a :class:`ZepbenTokenFetcher`.
+
+        :param token_fetcher: The :class:`ZepbenTokenFetcher` to use to fetch access tokens.
+        :return: This builder
+        """
+        if self._channel_credentials is None:
+            raise Exception("You must call make_secure before calling with_token_fetcher.")
+        self._channel_credentials = grpc.composite_channel_credentials(
+            self._channel_credentials,
+            grpc.metadata_call_credentials(AuthTokenPlugin(token_fetcher=token_fetcher))
+        )
+        return self
+
+    def with_client_token(self, client_token: str) -> 'GrpcChannelBuilder':
+        """
+        Authenticates calls for the gRPC channel using a client access token string.
+
+        :param client_token: The :class:`str` that has been generated in Evolve application.
+        :return: This builder
+        """
+        if self._channel_credentials is None:
+            raise Exception("You must call make_secure before calling with_client_token.")
+        self._channel_credentials = grpc.composite_channel_credentials(
+            self._channel_credentials,
+            grpc.access_token_call_credentials(client_token)
+        )
+        return self

zepben/ewb/streaming/mutations/__init__.py

@@ -0,0 +1,4 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.

zepben/ewb/streaming/mutations/update_network_state_client.py

@@ -0,0 +1,80 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ["UpdateNetworkStateClient"]
+
+from typing import List, Callable, AsyncGenerator, Iterable
+
+from zepben.protobuf.ns.network_state_pb2_grpc import UpdateNetworkStateServiceStub
+from zepben.protobuf.ns.network_state_requests_pb2 import SetCurrentStatesRequest as PBSetCurrentStatesRequest
+
+from zepben.ewb.streaming.data.current_state_event import CurrentStateEvent
+from zepben.ewb.streaming.data.current_state_event_batch import CurrentStateEventBatch
+from zepben.ewb.streaming.data.set_current_states_status import SetCurrentStatesStatus
+from zepben.ewb.streaming.grpc.grpc import GrpcClient
+
+
+class UpdateNetworkStateClient(GrpcClient):
+    """
+    A client class that provides functionality to interact with the gRPC service for updating network states.
+    A gRPC channel or stub must be provided.
+    """
+
+    _stub: UpdateNetworkStateServiceStub = None
+
+    def __init__(self, channel=None, stub: UpdateNetworkStateServiceStub = None, error_handlers: List[Callable[[Exception], bool]] = None, timeout: int = 60):
+        super().__init__(error_handlers=error_handlers, timeout=timeout)
+        if channel is None and stub is None:
+            raise ValueError("Must provide either a channel or a stub")
+        if stub is not None:
+            self._stub = stub
+        else:
+            self._stub = UpdateNetworkStateServiceStub(channel)
+
+    async def set_current_states(self, batch_id: int, batch: Iterable[CurrentStateEvent]) -> SetCurrentStatesStatus:
+        """
+        Sends a single batch of current state events to the gRPC service for processing.
+
+        This method allows for sending a single batch of current state events to the gRPC service using the gRPC stub provided in the constructor.
+
+        Args:
+            batch_id: A unique identifier for the batch of events being processed.
+            batch: A collection of CurrentStateEvent objects representing a single batch of events
+                                             to be processed by the gRPC service.
+
+        Returns:
+            A SetCurrentStatesStatus object representing the status of the batch after being processed by the service.
+        """
+
+        async def request_generator() -> AsyncGenerator[CurrentStateEventBatch, None]:
+            yield CurrentStateEventBatch(batch_id, batch)
+
+        responses = [response async for response in self.set_current_states_in_batches(request_generator())]
+        if len(responses) == 0:
+            raise RuntimeError(f"There are no response for this batch {batch_id}")
+
+        return responses[0]
+
+    async def set_current_states_in_batches(self, batches: AsyncGenerator[CurrentStateEventBatch, None]) -> AsyncGenerator[
+        SetCurrentStatesStatus, None]:
+        """
+        Sends a stream of current state events to the gRPC service for processing.
+
+        This method is responsible for streaming a batch of current state events to the gRPC service using the gRPC stub provided in the constructor.
+
+        Args:
+            batches: A stream of CurrentStateEventBatch objects, where each request contains a
+                     collection of CurrentStateEvent objects to be processed by the gRPC service.
+
+        Returns:
+            A stream of SetCurrentStatesResponse objects representing the status of each batch after being processed by the gRPC service.
+        """
+
+        async def request_generator():
+            async for batch in batches:
+                yield PBSetCurrentStatesRequest(messageId=batch.batch_id, event=[event.to_pb() for event in batch.events])
+
+        async for response in self._stub.setCurrentStates(request_generator()):
+            yield SetCurrentStatesStatus.from_pb(response)

zepben/ewb/streaming/mutations/update_network_state_service.py

@@ -0,0 +1,61 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ["UpdateNetworkStateService"]
+
+from typing import Tuple, Callable, AsyncGenerator, Any
+
+from zepben.protobuf.ns.network_state_pb2_grpc import UpdateNetworkStateServiceServicer
+from zepben.protobuf.ns.network_state_requests_pb2 import SetCurrentStatesRequest as PBSetCurrentStatesRequest
+from zepben.protobuf.ns.network_state_responses_pb2 import SetCurrentStatesResponse as PBSetCurrentStatesResponse
+
+from zepben.ewb.streaming.data.current_state_event import CurrentStateEvent
+from zepben.ewb.streaming.data.set_current_states_status import SetCurrentStatesStatus
+
+
+class UpdateNetworkStateService(UpdateNetworkStateServiceServicer):
+    """
+    A service class that provides a simplified interface for updating network state events
+    via gRPC without exposing the underlying complexity of gRPC mechanisms.
+
+    This class serves as a wrapper around the gRPC-generated service implementation,
+    allowing users to update the network state using a more convenient function type.
+
+    Attributes:
+        on_set_current_states: A function that takes CurrentStateEvent objects asynchronously and returns a
+                               SetCurrentStatesStatus asynchronously that reflects the success or failure of the batch
+                               update process.
+    """
+
+    def __init__(self, on_set_current_states: Callable[
+        [AsyncGenerator[Tuple[int, Tuple[CurrentStateEvent, Any]], None]], AsyncGenerator[SetCurrentStatesStatus, None]]):
+        self.on_set_current_states = on_set_current_states
+
+    async def setCurrentStates(self, request_iterator: AsyncGenerator[PBSetCurrentStatesRequest, None], context) -> AsyncGenerator[
+        PBSetCurrentStatesResponse, None]:
+        """
+        Handles streaming requests for setting current state events and responds with the result of the operation.
+
+        This method is a bidirectional streaming gRPC implementation that processes incoming
+        SetCurrentStatesRequest objects from the client, applies the provided state events using the callback function
+        passed in the constructor, and sends back a SetCurrentStatesResponse indicating the outcome of the update operation.
+
+        It allows clients to stream multiple state events for asynchronous processing. As each event is processed,
+        the service responds with the status of the operation in real-time, without waiting for all events to be received.
+
+        Args:
+            request_iterator: The stream of protobuf requests to process.
+            context: The gRPC context.
+
+        Returns:
+            A stream of protobuf SetCurrentStatesResponse sent back.
+        """
+
+        async def request_generator() -> AsyncGenerator[Tuple[int, Tuple[CurrentStateEvent, Any]], None]:
+            async for request in request_iterator:
+                yield request.messageId, tuple([CurrentStateEvent.from_pb(event) for event in request.event])
+
+        async for status in self.on_set_current_states(request_generator()):
+            yield status.to_pb()

zepben/ewb/testing/__init__.py

@@ -0,0 +1,4 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.