function_based_entity.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

"""End-to-end sample that demonstrates how to configure an orchestrator
that calls an activity function in a sequence and prints the outputs."""
import os

from azure.identity import DefaultAzureCredential

from durabletask import client, entities, task
from durabletask.azuremanaged.client import DurableTaskSchedulerClient
from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker


def counter(ctx: entities.EntityContext, input: int) -> int | None:
    if ctx.operation == "set":
        ctx.set_state(input)
    elif ctx.operation == "add":
        current_state = ctx.get_state(int, 0)
        new_state = current_state + (1 if input is None else input)
        ctx.set_state(new_state)
        return new_state
    elif ctx.operation == "get":
        return ctx.get_state(int, 0)
    else:
        raise ValueError(f"Unknown operation '{ctx.operation}'")


def counter_orchestrator(ctx: task.OrchestrationContext, _):
    """Orchestrator function that demonstrates the behavior of the counter entity"""

    entity_id = entities.EntityInstanceId("counter", "myCounter")

    # Initialize the entity with state 0
    ctx.signal_entity(entity_id, "set", 0)
    # Increment the counter by 1
    yield ctx.call_entity(entity_id, "add", 1)
    # Return the entity's current value (should be 1)
    return (yield ctx.call_entity(entity_id, "get"))


# Use environment variables if provided, otherwise use default emulator values
taskhub_name = os.getenv("TASKHUB", "default")
endpoint = os.getenv("ENDPOINT", "http://localhost:8080")

print(f"Using taskhub: {taskhub_name}")
print(f"Using endpoint: {endpoint}")

# Set credential to None for emulator, or DefaultAzureCredential for Azure
secure_channel = endpoint.startswith("https://")
credential = DefaultAzureCredential() if secure_channel else None
with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=secure_channel,
                                taskhub=taskhub_name, token_credential=credential) as w:
    w.add_orchestrator(counter_orchestrator)
    w.add_entity(counter)
    w.start()

    # Construct the client and run the orchestrations
    c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=secure_channel,
                                   taskhub=taskhub_name, token_credential=credential)
    instance_id = c.schedule_new_orchestration(counter_orchestrator)
    state = c.wait_for_orchestration_completion(instance_id, timeout=60)
    if state and state.runtime_status == client.OrchestrationStatus.COMPLETED:
        print(f'Orchestration completed! Result: {state.serialized_output}')
    elif state:
        print(f'Orchestration failed: {state.failure_details}')