Early-stage Nexus sample prototype

hello_nexus/__init__.py

@@ -0,0 +1,3 @@
+from rich.traceback import install
+
+install(show_locals=True)

hello_nexus/caller/app.py

@@ -0,0 +1,58 @@
+import asyncio
+import sys
+from typing import Any, Type
+
+from temporalio.client import Client
+from temporalio.worker import UnsandboxedWorkflowRunner, Worker
+
+from hello_nexus.caller.workflows import (
+    Echo2CallerWorkflow,
+    Echo3CallerWorkflow,
+    EchoCallerWorkflow,
+    Hello2CallerWorkflow,
+    HelloCallerWorkflow,
+)
+
+interrupt_event = asyncio.Event()
+
+
+async def execute_workflow(workflow_cls: Type[Any], input: Any) -> None:
+    client = await Client.connect("localhost:7233", namespace="my-caller-namespace-python")
+    task_queue = "my-caller-task-queue"
+
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[workflow_cls],
+        workflow_runner=UnsandboxedWorkflowRunner(),
+    ):
+        print("🟠 Caller worker started")
+        result = await client.execute_workflow(
+            workflow_cls.run,
+            input,
+            id="my-caller-workflow-id",
+            task_queue=task_queue,
+        )
+        print("🟢 workflow result:", result)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python -m nexus.caller.app [echo|hello]")
+        sys.exit(1)
+
+    [wf_name] = sys.argv[1:]
+    fn = {
+        "echo": lambda: execute_workflow(EchoCallerWorkflow, "hello"),
+        "echo2": lambda: execute_workflow(Echo2CallerWorkflow, "hello"),
+        "echo3": lambda: execute_workflow(Echo3CallerWorkflow, "hello"),
+        "hello": lambda: execute_workflow(HelloCallerWorkflow, "world"),
+        "hello2": lambda: execute_workflow(Hello2CallerWorkflow, "world"),
+    }[wf_name]
+
+    loop = asyncio.new_event_loop()
+    try:
+        loop.run_until_complete(fn())
+    except KeyboardInterrupt:
+        interrupt_event.set()
+        loop.run_until_complete(loop.shutdown_asyncgens())

hello_nexus/caller/workflows.py

@@ -0,0 +1,105 @@
+from datetime import timedelta
+
+import xray
+from temporalio import workflow
+from temporalio.exceptions import FailureError
+from temporalio.workflow import NexusClient
+
+from hello_nexus.service.interface import (
+    EchoInput,
+    EchoOutput,
+    HelloInput,
+    HelloOutput,
+    MyNexusService,
+)
+
+
+class CallerWorkflowBase:
+    def __init__(self):
+        self.nexus_client = NexusClient(
+            MyNexusService,  # or string name "my-nexus-service",
+            "my-nexus-endpoint-name-python",
+            schedule_to_close_timeout=timedelta(seconds=30),
+        )
+
+
+@workflow.defn
+class EchoCallerWorkflow(CallerWorkflowBase):
+    @xray.start_as_current_workflow_method_span()
+    @workflow.run
+    async def run(self, message: str) -> EchoOutput:
+        op_output = await self.nexus_client.execute_operation(
+            MyNexusService.echo,
+            EchoInput(message),
+        )
+        return op_output
+
+
+@workflow.defn
+class Echo2CallerWorkflow(CallerWorkflowBase):
+    @xray.start_as_current_workflow_method_span()
+    @workflow.run
+    async def run(self, message: str) -> EchoOutput:
+        op_output = await self.nexus_client.execute_operation(
+            MyNexusService.echo2,
+            EchoInput(message),
+        )
+        return op_output
+
+
+@workflow.defn
+class Echo3CallerWorkflow(CallerWorkflowBase):
+    @xray.start_as_current_workflow_method_span()
+    @workflow.run
+    async def run(self, message: str) -> EchoOutput:
+        op_output = await self.nexus_client.execute_operation(
+            MyNexusService.echo3,
+            EchoInput(message),
+        )
+        return op_output
+
+
+@workflow.defn
+class HelloCallerWorkflow(CallerWorkflowBase):
+    @xray.start_as_current_workflow_method_span()
+    @workflow.run
+    async def run(self, name: str) -> HelloOutput:
+        handle = await self.nexus_client.start_operation(
+            MyNexusService.hello,
+            HelloInput(name),
+        )
+        assert handle.cancel()
+        try:
+            await handle
+        except FailureError:
+            handle = await self.nexus_client.start_operation(
+                MyNexusService.hello,
+                HelloInput(name),
+            )
+            result = await handle
+            return result
+        raise AssertionError("Expected Nexus operation to be cancelled")
+
+
+@workflow.defn
+class Hello2CallerWorkflow(CallerWorkflowBase):
+    @xray.start_as_current_workflow_method_span()
+    @workflow.run
+    async def run(self, name: str) -> HelloOutput:
+        handle = await self.nexus_client.start_operation(
+            MyNexusService.hello2,
+            HelloInput(name),
+        )
+        return await handle
+
+
+@workflow.defn
+class Hello3CallerWorkflow(CallerWorkflowBase):
+    @xray.start_as_current_workflow_method_span()
+    @workflow.run
+    async def run(self, name: str) -> HelloOutput:
+        handle = await self.nexus_client.start_operation(
+            MyNexusService.hello3,
+            HelloInput(name),
+        )
+        return await handle

hello_nexus/clean

@@ -0,0 +1,14 @@
+temporal-delete-all my-target-namespace-python
+temporal-delete-all my-caller-namespace-python
+
+temporal operator namespace create --namespace my-target-namespace-python
+temporal operator namespace create --namespace my-caller-namespace-python
+
+sleep 1
+
+temporal operator nexus endpoint create \
+  --name my-nexus-endpoint-name-python \
+  --target-namespace my-target-namespace-python \
+  --target-task-queue my-target-task-queue-python \
+  --description-file ./hello_nexus/service/description.md
+

hello_nexus/handler/activities.py

@@ -0,0 +1,14 @@
+import asyncio
+
+from temporalio import activity
+
+from hello_nexus.service.interface import (
+    HelloInput,
+    HelloOutput,
+)
+
+
+@activity.defn
+async def hello_activity(input: HelloInput) -> HelloOutput:
+    await asyncio.sleep(1)
+    return HelloOutput(message=f"hello {input.name}")

hello_nexus/handler/dbclient.py

@@ -0,0 +1,3 @@
+class MyDBClient:
+    def execute(self, query: str) -> str:
+        return "<query result>"

hello_nexus/handler/nexus_service.py

@@ -0,0 +1,175 @@
+"""
+Notes:
+
+Sync operations:
+---------------
+Implementations are free to make arbitrary network calls, or perform CPU-bound
+computations such as this one. Total execution duration must not exceed 10s. To
+perform Temporal client calls such as signaling/querying/listing workflows, use
+self.client.
+
+
+Workflow operations:
+---------------------
+The task queue defaults to the task queue being used by the Nexus worker.
+"""
+
+from __future__ import annotations
+
+import nexusrpc.handler
+import temporalio.common
+import temporalio.nexus.handler
+
+from hello_nexus.handler.dbclient import MyDBClient
+from hello_nexus.handler.workflows import HelloWorkflow
+from hello_nexus.service import interface
+from hello_nexus.service.interface import (
+    EchoInput,
+    EchoOutput,
+    HelloInput,
+    HelloOutput,
+)
+
+
+# Inheriting from the protocol here is optional. Users who do it will get the
+# operation definition itself type-checked in situ against the interface (*).
+# Call-sites using instances of the operation are always type-checked.
+#
+# (*) However, in VSCode/Pyright this is done only when type-checking is set to
+# 'strict'.
+class EchoOperation(nexusrpc.handler.Operation[EchoInput, EchoOutput]):
+    def __init__(self, service: MyNexusService):
+        self.service = service
+
+    async def start(
+        self, input: EchoInput, options: nexusrpc.handler.StartOperationOptions
+    ) -> EchoOutput:
+        return EchoOutput(message=f"Echo {input.message}!")
+
+    async def cancel(
+        self, token: str, options: nexusrpc.handler.CancelOperationOptions
+    ) -> None:
+        raise NotImplementedError
+
+    async def fetch_info(
+        self, token: str, options: nexusrpc.handler.FetchOperationInfoOptions
+    ) -> nexusrpc.handler.OperationInfo:
+        raise NotImplementedError
+
+    async def fetch_result(
+        self, token: str, options: nexusrpc.handler.FetchOperationResultOptions
+    ) -> EchoOutput:
+        raise NotImplementedError
+
+
+# Inheriting from the protocol here is optional. Users who do it will get the
+# operation definition itself type-checked in situ against the interface (*).
+# Call-sites using instances of the operation are always type-checked.
+#
+# (*) However, in VSCode/Pyright this is done only when type-checking is set to
+# 'strict'.
+class HelloOperation:  # (nexusrpc.handler.Operation[HelloInput, HelloOutput]):
+    def __init__(self, service: "MyNexusService"):
+        self.service = service
+
+    async def start(
+        self, input: HelloInput, options: nexusrpc.handler.StartOperationOptions
+    ) -> temporalio.nexus.handler.StartWorkflowOperationResult[HelloOutput]:
+        self.service.db_client.execute("<some query>")
+        workflow_id = "default-workflow-id"
+        return await temporalio.nexus.handler.start_workflow(
+            HelloWorkflow.run,
+            input,
+            id=workflow_id,
+            options=options,
+        )
+
+    async def cancel(
+        self, token: str, options: nexusrpc.handler.CancelOperationOptions
+    ) -> None:
+        return await temporalio.nexus.handler.cancel_workflow(token, options)
+
+    async def fetch_info(
+        self, token: str, options: nexusrpc.handler.FetchOperationInfoOptions
+    ) -> nexusrpc.handler.OperationInfo:
+        return await temporalio.nexus.handler.fetch_workflow_info(token, options)
+
+    async def fetch_result(
+        self, token: str, options: nexusrpc.handler.FetchOperationResultOptions
+    ) -> HelloOutput:
+        return await temporalio.nexus.handler.fetch_workflow_result(token, options)
+
+
+class EchoOperation3(nexusrpc.handler.AbstractOperation[EchoInput, EchoOutput]):
+    async def start(
+        self, input: EchoInput, options: nexusrpc.handler.StartOperationOptions
+    ) -> EchoOutput:
+        return EchoOutput(message=f"Echo {input.message}! [from base class variant]")
+
+
+@nexusrpc.handler.service(interface=interface.MyNexusService)
+class MyNexusService:
+    def __init__(self, db_client: MyDBClient):
+        # An example of something that might be held by the service instance.
+        self.db_client = db_client
+
+    # --------------------------------------------------------------------------
+    # Operations defined by explicitly implementing the Operation interface.
+    #
+
+    @nexusrpc.handler.operation
+    def echo(self) -> nexusrpc.handler.Operation[EchoInput, EchoOutput]:
+        return EchoOperation(self)
+
+    @nexusrpc.handler.operation
+    def hello(self) -> nexusrpc.handler.Operation[HelloInput, HelloOutput]:
+        return HelloOperation(self)
+
+    @nexusrpc.handler.operation
+    def echo3(self) -> nexusrpc.handler.Operation[EchoInput, EchoOutput]:
+        return EchoOperation3()
+
+    # --------------------------------------------------------------------------
+    # Operations defined by providing the start method only, using the
+    # "shorthand" decorators.
+    #
+    # Note that a start method defined this way has access to the service
+    # instance, but not to the operation instance (users who need the latter
+    # should implement the Operation interface directly).
+
+    @nexusrpc.handler.sync_operation
+    async def echo2(
+        self, input: EchoInput, _: nexusrpc.handler.StartOperationOptions
+    ) -> EchoOutput:
+        return EchoOutput(message=f"Echo {input.message} [via shorthand]!")
+
+    # --------------------------------------------------------------------------
+    # Operations defined by providing the start method only, using the
+    # "shorthand" decorators.
+    #
+    # Note that a start method defined this way has access to the service
+    # instance, but not to the operation instance (users who need the latter
+    # should implement the Operation interface directly).
+
+    @temporalio.nexus.handler.workflow_operation
+    async def hello2(
+        self, input: HelloInput, options: nexusrpc.handler.StartOperationOptions
+    ) -> temporalio.nexus.handler.StartWorkflowOperationResult[HelloOutput]:
+        self.db_client.execute("<some query>")
+        workflow_id = "default-workflow-id"
+        input.name += " [via shorthand]"
+        return await temporalio.nexus.handler.start_workflow(
+            HelloWorkflow.run,
+            input,
+            id=workflow_id,
+            options=options,
+        )
+
+
+if __name__ == "__main__":
+    # Check run-time type annotations resulting from the decorators.
+    service = MyNexusService(MyDBClient())
+    print("echo:", temporalio.common._type_hints_from_func(service.echo2().start))
+    print(
+        "hello:", temporalio.common._type_hints_from_func(service.hello2().fetch_result)
+    )