Add docs and metadata to the services/handlers

Author: igalshilman

python/restate/discovery.py

@@ -21,11 +21,13 @@
 # pylint: disable=C0115
 # pylint: disable=C0103
 # pylint: disable=W0622
+# pylint: disable=R0913,
+# pylint: disable=R0917,
 
 import json
 import typing
 from enum import Enum
-from typing import Optional, Any, List, get_args, get_origin
+from typing import Dict, Optional, Any, List, get_args, get_origin
 
 
 from restate.endpoint import Endpoint as RestateEndpoint
@@ -58,17 +60,21 @@ def __init__(self, contentType: str, setContentTypeIfEmpty: bool, jsonSchema: Op
         self.jsonSchema = jsonSchema
 
 class Handler:
-    def __init__(self, name: str, ty: Optional[ServiceHandlerType] = None, input: Optional[InputPayload] = None, output: Optional[OutputPayload] = None):
+    def __init__(self, name: str, ty: Optional[ServiceHandlerType] = None, input: Optional[InputPayload] = None, output: Optional[OutputPayload] = None, description: Optional[str] = None, metadata: Optional[Dict[str, str]] = None):
         self.name = name
         self.ty = ty
         self.input = input
         self.output = output
+        self.documentation = description
+        self.metadata = metadata
 
 class Service:
-    def __init__(self, name: str, ty: ServiceType, handlers: List[Handler]):
+    def __init__(self, name: str, ty: ServiceType, handlers: List[Handler], description: Optional[str] = None, metadata: Optional[Dict[str, str]] = None):
         self.name = name
         self.ty = ty
         self.handlers = handlers
+        self.documentation = description
+        self.metadata = metadata
 
 class Endpoint:
     def __init__(self, protocolMode: ProtocolMode, minProtocolVersion: int, maxProtocolVersion: int, services: List[Service]):
@@ -182,10 +188,16 @@ def compute_discovery(endpoint: RestateEndpoint, discovered_as : typing.Literal[
                                 contentType=handler.handler_io.content_type,
                                 jsonSchema=json_schema_from_type_hint(handler.handler_io.output_type))
             # add the handler
-            service_handlers.append(Handler(name=handler.name, ty=ty, input=inp, output=out))
-
+            service_handlers.append(Handler(name=handler.name,
+                                            ty=ty,
+                                            input=inp,
+                                            output=out,
+                                            description=handler.description,
+                                            metadata=handler.metadata))
         # add the service
-        services.append(Service(name=service.name, ty=service_type, handlers=service_handlers))
+        description = service.service_tag.description
+        metadata = service.service_tag.metadata
+        services.append(Service(name=service.name, ty=service_type, handlers=service_handlers, description=description, metadata=metadata))
 
     if endpoint.protocol:
         protocol_mode = PROTOCOL_MODES[endpoint.protocol]

python/restate/handler.py

@@ -17,7 +17,7 @@
 
 from dataclasses import dataclass
 from inspect import Signature
-from typing import Any, Callable, Awaitable, Generic, Literal, Optional, TypeVar
+from typing import Any, Callable, Awaitable, Dict, Generic, Literal, Optional, TypeVar
 
 from restate.exceptions import TerminalError
 from restate.serde import JsonSerde, Serde, PydanticJsonSerde
@@ -52,6 +52,8 @@ class ServiceTag:
     """
     kind: Literal["object", "service", "workflow"]
     name: str
+    description: Optional[str] = None
+    metadata: Optional[Dict[str, str]] = None
 
 @dataclass
 class TypeHint(Generic[T]):
@@ -114,6 +116,7 @@ def extract_io_type_hints(handler_io: HandlerIO[I, O], signature: Signature):
         if isinstance(handler_io.output_serde, JsonSerde): # type: ignore
             handler_io.output_serde = PydanticJsonSerde(annotation)
 
+# pylint: disable=R0902
 @dataclass
 class Handler(Generic[I, O]):
     """
@@ -125,6 +128,8 @@ class Handler(Generic[I, O]):
     name: str
     fn: Callable[[Any, I], Awaitable[O]] | Callable[[Any], Awaitable[O]]
     arity: int
+    description: Optional[str] = None
+    metadata: Optional[Dict[str, str]] = None
 
 
 # disable too many arguments warning
@@ -135,7 +140,9 @@ def make_handler(service_tag: ServiceTag,
                  name: str | None,
                  kind: Optional[Literal["exclusive", "shared", "workflow"]],
                  wrapped: Any,
-                 signature: Signature) -> Handler[I, O]:
+                 signature: Signature,
+                 description: Optional[str] = None,
+                 metadata: Optional[Dict[str, str]] = None) -> Handler[I, O]:
     """
     Factory function to create a handler.
     """
@@ -152,12 +159,14 @@ def make_handler(service_tag: ServiceTag,
     arity = len(signature.parameters)
     extract_io_type_hints(handler_io, signature)
 
-    handler = Handler[I, O](service_tag,
-                            handler_io,
-                            kind,
-                            handler_name,
-                            wrapped,
-                            arity)
+    handler = Handler[I, O](service_tag=service_tag,
+                            handler_io=handler_io,
+                            kind=kind,
+                            name=handler_name,
+                            fn=wrapped,
+                            arity=arity,
+                            description=description,
+                            metadata=metadata)
 
     vars(wrapped)[RESTATE_UNIQUE_HANDLER_SYMBOL] = handler
     return handler

python/restate/object.py

@@ -18,7 +18,7 @@
 import typing
 
 from restate.serde import Serde, JsonSerde
-from .handler import HandlerIO, ServiceTag, make_handler
+from restate.handler import Handler, HandlerIO, ServiceTag, make_handler
 
 I = typing.TypeVar('I')
 O = typing.TypeVar('O')
@@ -36,10 +36,16 @@ class VirtualObject:
 
     Args:
         name (str): The name of the object.
+        description (str): The description of the object.
+        metadata (dict): The metadata of the object.
     """
 
-    def __init__(self, name):
-        self.service_tag = ServiceTag("object", name)
+    handlers: typing.Dict[str, Handler[typing.Any, typing.Any]]
+
+    def __init__(self, name,
+                 description: typing.Optional[str] = None,
+                 metadata: typing.Optional[typing.Dict[str, str]]=None):
+        self.service_tag = ServiceTag("object", name, description, metadata)
         self.handlers = {}
 
     @property
@@ -55,7 +61,8 @@ def handler(self,
                 accept: str = "application/json",
                 content_type: str = "application/json",
                 input_serde: Serde[I] = JsonSerde[I](), # type: ignore
-                output_serde: Serde[O] = JsonSerde[O]()) -> typing.Callable: # type: ignore
+                output_serde: Serde[O] = JsonSerde[O](), # type: ignore
+                metadata: typing.Optional[dict] = None) -> typing.Callable:
         """
         Decorator for defining a handler function.
 
@@ -86,7 +93,7 @@ def wrapped(*args, **kwargs):
                 return fn(*args, **kwargs)
 
             signature = inspect.signature(fn, eval_str=True)
-            handler = make_handler(self.service_tag, handler_io, name, kind, wrapped, signature)
+            handler = make_handler(self.service_tag, handler_io, name, kind, wrapped, signature, inspect.getdoc(fn), metadata)
             self.handlers[handler.name] = handler
             return wrapped
 

python/restate/service.py

@@ -19,7 +19,7 @@
 import typing
 
 from restate.serde import Serde, JsonSerde
-from .handler import Handler, HandlerIO, ServiceTag, make_handler
+from restate.handler import Handler, HandlerIO, ServiceTag, make_handler
 
 I = typing.TypeVar('I')
 O = typing.TypeVar('O')
@@ -39,8 +39,10 @@ class Service:
         name (str): The name of the service.
     """
 
-    def __init__(self, name: str) -> None:
-        self.service_tag = ServiceTag("service", name)
+    def __init__(self, name: str,
+                 description: typing.Optional[str] = None,
+                 metadata: typing.Optional[typing.Dict[str, str]] = None) -> None:
+        self.service_tag = ServiceTag("service", name, description, metadata)
         self.handlers: typing.Dict[str, Handler] = {}
 
     @property
@@ -55,7 +57,8 @@ def handler(self,
                 accept: str = "application/json",
                 content_type: str = "application/json",
                 input_serde: Serde[I] = JsonSerde[I](), # type: ignore
-                output_serde: Serde[O] = JsonSerde[O]()) -> typing.Callable: # type: ignore
+                output_serde: Serde[O] = JsonSerde[O](), # type: ignore
+                metadata: typing.Optional[typing.Dict[str, str]] = None) -> typing.Callable:
         """
         Decorator for defining a handler function.
 
@@ -85,7 +88,7 @@ def wrapped(*args, **kwargs):
                 return fn(*args, **kwargs)
 
             signature = inspect.signature(fn, eval_str=True)
-            handler = make_handler(self.service_tag, handler_io, name, None, wrapped, signature)
+            handler = make_handler(self.service_tag, handler_io, name, None, wrapped, signature, inspect.getdoc(fn), metadata)
             self.handlers[handler.name] = handler
             return wrapped
 

python/restate/workflow.py

@@ -19,7 +19,7 @@
 import typing
 
 from restate.serde import Serde, JsonSerde
-from .handler import HandlerIO, ServiceTag, make_handler
+from restate.handler import Handler, HandlerIO, ServiceTag, make_handler
 
 I = typing.TypeVar('I')
 O = typing.TypeVar('O')
@@ -42,8 +42,10 @@ class Workflow:
         name (str): The name of the object.
     """
 
-    def __init__(self, name):
-        self.service_tag = ServiceTag("workflow", name)
+    handlers: typing.Dict[str, Handler[typing.Any, typing.Any]]
+
+    def __init__(self, name, description: typing.Optional[str] = None, metadata: typing.Optional[typing.Dict[str, str]] = None):
+        self.service_tag = ServiceTag("workflow", name, description, metadata)
         self.handlers = {}
 
     @property
@@ -58,33 +60,37 @@ def main(self,
             accept: str = "application/json",
             content_type: str = "application/json",
             input_serde: Serde[I] = JsonSerde[I](), # type: ignore
-            output_serde: Serde[O] = JsonSerde[O]()) -> typing.Callable: # type: ignore
+            output_serde: Serde[O] = JsonSerde[O](), # type: ignore
+            metadata: typing.Optional[typing.Dict[str, str]] = None) -> typing.Callable: # type: ignore
         """Mark this handler as a workflow entry point"""
         return self._add_handler(name,
                             kind="workflow",
                              accept=accept,
                              content_type=content_type,
                              input_serde=input_serde,
-                             output_serde=output_serde)
+                             output_serde=output_serde,
+                             metadata=metadata)
 
     def handler(self,
                 name: typing.Optional[str] = None,
                 accept: str = "application/json",
                 content_type: str = "application/json",
                 input_serde: Serde[I] = JsonSerde[I](), # type: ignore
-                output_serde: Serde[O] = JsonSerde[O]()) -> typing.Callable: # type: ignore
+                output_serde: Serde[O] = JsonSerde[O](), # type: ignore
+                metadata: typing.Optional[typing.Dict[str, str]] = None) -> typing.Callable:
         """
         Decorator for defining a handler function.
         """
-        return self._add_handler(name, "shared", accept, content_type, input_serde, output_serde)
+        return self._add_handler(name, "shared", accept, content_type, input_serde, output_serde, metadata)
 
     def _add_handler(self,
                 name: typing.Optional[str] = None,
                 kind: typing.Literal["workflow", "shared", "exclusive"] = "shared",
                 accept: str = "application/json",
                 content_type: str = "application/json",
                 input_serde: Serde[I] = JsonSerde[I](), # type: ignore
-                output_serde: Serde[O] = JsonSerde[O]()) -> typing.Callable: # type: ignore
+                output_serde: Serde[O] = JsonSerde[O](), # type: ignore
+                metadata: typing.Optional[typing.Dict[str, str]] = None) -> typing.Callable: # type: ignore
         """
         Decorator for defining a handler function.
 
@@ -94,6 +100,7 @@ def _add_handler(self,
             content_type: The content type of the request. Default "application/json".
             serializer: The serializer function to convert the response object to bytes. 
             deserializer: The deserializer function to convert the request bytes to an object.
+            metadata: An optional dictionary of metadata.
 
         Returns:
             Callable: The decorated function.
@@ -115,7 +122,15 @@ def wrapped(*args, **kwargs):
                 return fn(*args, **kwargs)
 
             signature = inspect.signature(fn, eval_str=True)
-            handler = make_handler(self.service_tag, handler_io, name, kind, wrapped, signature)
+            description = inspect.getdoc(fn)
+            handler = make_handler(service_tag=self.service_tag,
+                                   handler_io=handler_io,
+                                   name=name,
+                                   kind=kind,
+                                   wrapped=wrapped,
+                                   signature=signature,
+                                   description=description,
+                                   metadata=metadata)
             self.handlers[handler.name] = handler
             return wrapped