Adds support for gRPC aio server that allows for async handlers

Signed-off-by: Vasilis Tsiolkas <[email protected]>

Author: Vasilis Tsiolkas

ext/dapr-ext-grpc/dapr/ext/grpc/aio/__init__.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+
+"""
+Copyright 2023 The Dapr Authors
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from dapr.clients.grpc._request import InvokeMethodRequest, BindingRequest, JobEvent
+from dapr.clients.grpc._response import InvokeMethodResponse, TopicEventResponse
+from dapr.clients.grpc._jobs import Job, FailurePolicy, DropFailurePolicy, ConstantFailurePolicy
+
+from dapr.ext.grpc.aio.app import App, Rule  # type:ignore
+
+
+__all__ = [
+    'App',
+    'Rule',
+    'InvokeMethodRequest',
+    'InvokeMethodResponse',
+    'BindingRequest',
+    'TopicEventResponse',
+    'Job',
+    'JobEvent',
+    'FailurePolicy',
+    'DropFailurePolicy',
+    'ConstantFailurePolicy',
+]

ext/dapr-ext-grpc/dapr/ext/grpc/aio/_health_servicer.py

@@ -0,0 +1,32 @@
+import grpc
+from typing import Callable, Optional
+
+from dapr.proto import appcallback_service_v1
+from dapr.proto.runtime.v1.appcallback_pb2 import HealthCheckResponse
+
+HealthCheckCallable = Optional[Callable[[], None]]
+
+
+class _AioHealthCheckServicer(appcallback_service_v1.AppCallbackHealthCheckServicer):
+    """The implementation of HealthCheck Server.
+
+    :class:`App` provides useful decorators to register method, topic, input bindings.
+    """
+
+    def __init__(self):
+        self._health_check_cb: Optional[HealthCheckCallable] = None
+
+    def register_health_check(self, cb: HealthCheckCallable) -> None:
+        if not cb:
+            raise ValueError('health check callback must be defined')
+        self._health_check_cb = cb
+
+    async def HealthCheck(self, request, context):
+        """Health check."""
+
+        if not self._health_check_cb:
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)  # type: ignore
+            context.set_details('Method not implemented!')
+            raise NotImplementedError('Method not implemented!')
+        await self._health_check_cb()
+        return HealthCheckResponse()

ext/dapr-ext-grpc/dapr/ext/grpc/aio/_servicer.py

@@ -0,0 +1,280 @@
+# -*- coding: utf-8 -*-
+
+"""
+Copyright 2023 The Dapr Authors
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import grpc.aio
+
+from cloudevents.sdk.event import v1  # type: ignore
+from typing import Awaitable, Callable, Dict, List, Optional, Tuple, Union
+
+from google.protobuf import empty_pb2
+from google.protobuf.message import Message as GrpcMessage
+from google.protobuf.struct_pb2 import Struct
+
+from dapr.proto import appcallback_service_v1, common_v1, appcallback_v1
+from dapr.proto.runtime.v1.appcallback_pb2 import (
+    TopicEventRequest,
+    BindingEventRequest,
+    JobEventRequest,
+)
+from dapr.proto.common.v1.common_pb2 import InvokeRequest
+from dapr.clients.base import DEFAULT_JSON_CONTENT_TYPE
+from dapr.clients.grpc._request import InvokeMethodRequest, BindingRequest, JobEvent
+from dapr.clients.grpc._response import InvokeMethodResponse, TopicEventResponse
+
+InvokeMethodCallable = Callable[
+    [InvokeMethodRequest], Awaitable[Union[str, bytes, InvokeMethodResponse]]
+]
+TopicSubscribeCallable = Callable[[v1.Event], Awaitable[Optional[TopicEventResponse]]]
+BindingCallable = Callable[[BindingRequest], Awaitable[None]]
+JobEventCallable = Callable[[JobEvent], Awaitable[None]]
+
+DELIMITER = ':'
+
+
+class Rule:
+    def __init__(self, match: str, priority: int) -> None:
+        self.match = match
+        self.priority = priority
+
+
+class _RegisteredSubscription:
+    def __init__(
+        self,
+        subscription: appcallback_v1.TopicSubscription,
+        rules: List[Tuple[int, appcallback_v1.TopicRule]],
+    ):
+        self.subscription = subscription
+        self.rules = rules
+
+
+class _AioCallbackServicer(
+    appcallback_service_v1.AppCallbackServicer, appcallback_service_v1.AppCallbackAlphaServicer
+):
+    """The asyncio-native implementation of the AppCallback Server.
+
+    This internal class implements application server and provides helpers to register
+    method, topic, and input bindings. It implements the routing handling logic to route
+    mulitple methods, topics, and bindings.
+
+    :class:`App` provides useful decorators to register method, topic, input bindings.
+    """
+
+    def __init__(self):
+        self._invoke_method_map: Dict[str, InvokeMethodCallable] = {}
+        self._topic_map: Dict[str, TopicSubscribeCallable] = {}
+        self._binding_map: Dict[str, BindingCallable] = {}
+        self._job_event_map: Dict[str, JobEventCallable] = {}
+
+        self._registered_topics_map: Dict[str, _RegisteredSubscription] = {}
+        self._registered_topics: List[appcallback_v1.TopicSubscription] = []
+        self._registered_bindings: List[str] = []
+
+    def register_method(self, method: str, cb: InvokeMethodCallable) -> None:
+        """Registers method for service invocation."""
+        if method in self._invoke_method_map:
+            raise ValueError(f'{method} is already registered')
+        self._invoke_method_map[method] = cb
+
+    def register_topic(
+        self,
+        pubsub_name: str,
+        topic: str,
+        cb: TopicSubscribeCallable,
+        metadata: Optional[Dict[str, str]],
+        dead_letter_topic: Optional[str] = None,
+        rule: Optional[Rule] = None,
+        disable_topic_validation: Optional[bool] = False,
+    ) -> None:
+        """Registers topic subscription for pubsub."""
+        if not disable_topic_validation:
+            topic_key = pubsub_name + DELIMITER + topic
+        else:
+            topic_key = pubsub_name
+        pubsub_topic = topic_key + DELIMITER
+        if rule is not None:
+            path = getattr(cb, '__name__', rule.match)
+            pubsub_topic = pubsub_topic + path
+        if pubsub_topic in self._topic_map:
+            raise ValueError(f'{topic} is already registered with {pubsub_name}')
+        self._topic_map[pubsub_topic] = cb
+
+        registered_topic = self._registered_topics_map.get(topic_key)
+        sub: appcallback_v1.TopicSubscription = appcallback_v1.TopicSubscription()
+        rules: List[Tuple[int, appcallback_v1.TopicRule]] = []
+        if not registered_topic:
+            sub = appcallback_v1.TopicSubscription(
+                pubsub_name=pubsub_name,
+                topic=topic,
+                metadata=metadata,
+                routes=appcallback_v1.TopicRoutes(),
+            )
+            if dead_letter_topic:
+                sub.dead_letter_topic = dead_letter_topic
+            registered_topic = _RegisteredSubscription(sub, rules)
+            self._registered_topics_map[topic_key] = registered_topic
+            self._registered_topics.append(sub)
+
+        sub = registered_topic.subscription
+        rules = registered_topic.rules
+
+        if rule:
+            path = getattr(cb, '__name__', rule.match)
+            rules.append((rule.priority, appcallback_v1.TopicRule(match=rule.match, path=path)))
+            rules.sort(key=lambda x: x[0])
+            rs = [rule for id, rule in rules]
+            del sub.routes.rules[:]
+            sub.routes.rules.extend(rs)
+
+    def register_binding(self, name: str, cb: BindingCallable) -> None:
+        """Registers input bindings."""
+        if name in self._binding_map:
+            raise ValueError(f'{name} is already registered')
+        self._binding_map[name] = cb
+        self._registered_bindings.append(name)
+
+    def register_job_event(self, name: str, cb: JobEventCallable) -> None:
+        """Registers job event handler.
+
+        Args:
+            name (str): The name of the job to handle events for.
+            cb (JobEventCallable): The callback function to handle job events.
+        """
+        if name in self._job_event_map:
+            raise ValueError(f'Job event handler for {name} is already registered')
+        self._job_event_map[name] = cb
+
+    async def OnInvoke(self, request: InvokeRequest, context):
+        """Invokes service method with InvokeRequest."""
+        if request.method not in self._invoke_method_map:
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)  # type: ignore
+            raise NotImplementedError(f'{request.method} method not implemented!')
+
+        req = InvokeMethodRequest(request.data, request.content_type)
+        req.metadata = context.invocation_metadata()
+        resp = await self._invoke_method_map[request.method](req)
+
+        if not resp:
+            return common_v1.InvokeResponse()
+
+        resp_data = InvokeMethodResponse()
+        if isinstance(resp, (bytes, str)):
+            resp_data.set_data(resp)
+            resp_data.content_type = DEFAULT_JSON_CONTENT_TYPE
+        elif isinstance(resp, GrpcMessage):
+            resp_data.set_data(resp)
+        elif isinstance(resp, InvokeMethodResponse):
+            resp_data = resp
+        else:
+            context.set_code(grpc.StatusCode.OUT_OF_RANGE)
+            context.set_details(f'{type(resp)} is the invalid return type.')
+            raise NotImplementedError(f'{request.method} method not implemented!')
+
+        if len(resp_data.get_headers()) > 0:
+            context.send_initial_metadata(resp_data.get_headers())
+
+        content_type = ''
+        if resp_data.content_type:
+            content_type = resp_data.content_type
+
+        return common_v1.InvokeResponse(data=resp_data.proto, content_type=content_type)
+
+    async def ListTopicSubscriptions(self, request, context):
+        """Lists all topics subscribed by this app."""
+        return appcallback_v1.ListTopicSubscriptionsResponse(subscriptions=self._registered_topics)
+
+    async def OnTopicEvent(self, request: TopicEventRequest, context):
+        """Subscribes events from Pubsub."""
+        pubsub_topic = request.pubsub_name + DELIMITER + request.topic + DELIMITER + request.path
+        no_validation_key = request.pubsub_name + DELIMITER + request.path
+
+        if pubsub_topic not in self._topic_map:
+            if no_validation_key in self._topic_map:
+                pubsub_topic = no_validation_key
+            else:
+                context.set_code(grpc.StatusCode.UNIMPLEMENTED)  # type: ignore
+                raise NotImplementedError(f'topic {request.topic} is not implemented!')
+
+        customdata: Struct = request.extensions
+        extensions = dict()
+        for k, v in customdata.items():
+            extensions[k] = v
+        for k, v in context.invocation_metadata():
+            extensions['_metadata_' + k] = v
+
+        event = v1.Event()
+        event.SetEventType(request.type)
+        event.SetEventID(request.id)
+        event.SetSource(request.source)
+        event.SetData(request.data)
+        event.SetContentType(request.data_content_type)
+        event.SetSubject(request.topic)
+        event.SetExtensions(extensions)
+
+        response = await self._topic_map[pubsub_topic](event)
+        if isinstance(response, TopicEventResponse):
+            return appcallback_v1.TopicEventResponse(status=response.status.value)
+        return empty_pb2.Empty()
+
+    async def ListInputBindings(self, request, context):
+        """Lists all input bindings subscribed by this app."""
+        return appcallback_v1.ListInputBindingsResponse(bindings=self._registered_bindings)
+
+    async def OnBindingEvent(self, request: BindingEventRequest, context):
+        """Listens events from the input bindings
+        User application can save the states or send the events to the output
+        bindings optionally by returning BindingEventResponse.
+        """
+        if request.name not in self._binding_map:
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)  # type: ignore
+            raise NotImplementedError(f'{request.name} binding not implemented!')
+
+        req = BindingRequest(request.data, dict(request.metadata))
+        req.metadata = context.invocation_metadata()
+        await self._binding_map[request.name](req)
+
+        # TODO: support output bindings options
+        return appcallback_v1.BindingEventResponse()
+
+    async def OnJobEventAlpha1(self, request: JobEventRequest, context):
+        """Handles job events from Dapr runtime.
+
+        This method is called by Dapr when a scheduled job is triggered.
+        It routes the job event to the appropriate registered handler based on the job name.
+
+        Args:
+            request (JobEventRequest): The job event request from Dapr.
+            context: The gRPC context.
+
+        Returns:
+            appcallback_v1.JobEventResponse: Empty response indicating successful handling.
+        """
+        job_name = request.name
+
+        if job_name not in self._job_event_map:
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)  # type: ignore
+            raise NotImplementedError(f'Job event handler for {job_name} not implemented!')
+
+        # Create a JobEvent object matching Go SDK's common.JobEvent
+        # Extract raw data bytes from the Any proto (matching Go implementation)
+        data_bytes = b''
+        if request.HasField('data') and request.data.value:
+            data_bytes = request.data.value
+
+        job_event = JobEvent(name=request.name, data=data_bytes)
+
+        # Call the registered handler with the JobEvent object
+        await self._job_event_map[job_name](job_event)
+
+        # Return empty response
+        return appcallback_v1.JobEventResponse()