|
| 1 | +# Copyright (c) Microsoft Corporation. |
| 2 | +# Licensed under the MIT License. |
| 3 | + |
| 4 | +import logging |
| 5 | + |
| 6 | +from durabletask.client import (EntityQuery, OrchestrationStatus, |
| 7 | + TaskHubGrpcClient) |
| 8 | +from durabletask.entities import EntityInstanceId |
| 9 | +from durabletask.internal.helpers import ensure_aware |
| 10 | +from durabletask.scheduled import transitions |
| 11 | +from durabletask.scheduled.exceptions import ScheduleNotFoundError |
| 12 | +from durabletask.scheduled.models import (ScheduleCreationOptions, |
| 13 | + ScheduleDescription, ScheduleQuery, |
| 14 | + ScheduleState, ScheduleUpdateOptions) |
| 15 | +from durabletask.scheduled.orchestrator import ( |
| 16 | + ScheduleOperationRequest, execute_schedule_operation_orchestrator) |
| 17 | +from durabletask.scheduled.schedule_entity import (DELETE_OPERATION, |
| 18 | + ENTITY_NAME) |
| 19 | + |
| 20 | +logger = logging.getLogger("durabletask.scheduled") |
| 21 | + |
| 22 | + |
| 23 | +class ScheduleClient: |
| 24 | + """Client for managing a single schedule instance.""" |
| 25 | + |
| 26 | + def __init__(self, client: TaskHubGrpcClient, schedule_id: str, |
| 27 | + *, operation_timeout: float = 60): |
| 28 | + if not schedule_id: |
| 29 | + raise ValueError("schedule_id cannot be empty.") |
| 30 | + self._client = client |
| 31 | + self._schedule_id = schedule_id |
| 32 | + self._entity_id = EntityInstanceId(ENTITY_NAME, schedule_id) |
| 33 | + self._operation_timeout = operation_timeout |
| 34 | + |
| 35 | + @property |
| 36 | + def schedule_id(self) -> str: |
| 37 | + """Gets the ID of this schedule.""" |
| 38 | + return self._schedule_id |
| 39 | + |
| 40 | + def _run_operation(self, operation_name: str, input: object | None = None) -> None: |
| 41 | + request = ScheduleOperationRequest( |
| 42 | + entity_id=str(self._entity_id), |
| 43 | + operation_name=operation_name, |
| 44 | + input=input, |
| 45 | + ) |
| 46 | + instance_id = self._client.schedule_new_orchestration( |
| 47 | + execute_schedule_operation_orchestrator, input=request) |
| 48 | + state = self._client.wait_for_orchestration_completion( |
| 49 | + instance_id, timeout=self._operation_timeout) |
| 50 | + if state is None or state.runtime_status != OrchestrationStatus.COMPLETED: |
| 51 | + failure = state.failure_details if state else None |
| 52 | + message = failure.message if failure else "unknown error" |
| 53 | + raise RuntimeError( |
| 54 | + f"Failed to '{operation_name}' schedule '{self._schedule_id}': {message}") |
| 55 | + |
| 56 | + def create(self, options: ScheduleCreationOptions) -> None: |
| 57 | + """Create or update this schedule with the given configuration.""" |
| 58 | + self._run_operation(transitions.CREATE_SCHEDULE, options) |
| 59 | + |
| 60 | + def update(self, options: ScheduleUpdateOptions) -> None: |
| 61 | + """Update this schedule's configuration.""" |
| 62 | + self._run_operation(transitions.UPDATE_SCHEDULE, options) |
| 63 | + |
| 64 | + def pause(self) -> None: |
| 65 | + """Pause this schedule.""" |
| 66 | + self._run_operation(transitions.PAUSE_SCHEDULE) |
| 67 | + |
| 68 | + def resume(self) -> None: |
| 69 | + """Resume this schedule.""" |
| 70 | + self._run_operation(transitions.RESUME_SCHEDULE) |
| 71 | + |
| 72 | + def delete(self) -> None: |
| 73 | + """Delete this schedule.""" |
| 74 | + self._run_operation(DELETE_OPERATION) |
| 75 | + |
| 76 | + def describe(self) -> ScheduleDescription: |
| 77 | + """Retrieve the current details of this schedule.""" |
| 78 | + metadata = self._client.get_entity(self._entity_id, include_state=True) |
| 79 | + if metadata is None: |
| 80 | + raise ScheduleNotFoundError(self._schedule_id) |
| 81 | + state = metadata.get_typed_state(ScheduleState) |
| 82 | + if state is None: |
| 83 | + raise ScheduleNotFoundError(self._schedule_id) |
| 84 | + return state.to_description() |
| 85 | + |
| 86 | + |
| 87 | +class ScheduledTaskClient: |
| 88 | + """Client for managing scheduled tasks in a Durable Task application.""" |
| 89 | + |
| 90 | + def __init__(self, client: TaskHubGrpcClient, *, operation_timeout: float = 60): |
| 91 | + self._client = client |
| 92 | + self._operation_timeout = operation_timeout |
| 93 | + |
| 94 | + def get_schedule_client(self, schedule_id: str) -> ScheduleClient: |
| 95 | + """Get a handle to manage a specific schedule.""" |
| 96 | + return ScheduleClient(self._client, schedule_id, |
| 97 | + operation_timeout=self._operation_timeout) |
| 98 | + |
| 99 | + def create_schedule(self, options: ScheduleCreationOptions) -> ScheduleClient: |
| 100 | + """Create a new schedule and return a client for managing it.""" |
| 101 | + schedule_client = self.get_schedule_client(options.schedule_id) |
| 102 | + schedule_client.create(options) |
| 103 | + return schedule_client |
| 104 | + |
| 105 | + def get_schedule(self, schedule_id: str) -> ScheduleDescription | None: |
| 106 | + """Get a schedule description by ID, or None if it does not exist.""" |
| 107 | + try: |
| 108 | + return self.get_schedule_client(schedule_id).describe() |
| 109 | + except ScheduleNotFoundError: |
| 110 | + return None |
| 111 | + |
| 112 | + def list_schedules(self, schedule_query: ScheduleQuery | None = None) -> list[ScheduleDescription]: |
| 113 | + """List schedules matching the given filter criteria. |
| 114 | +
|
| 115 | + > [!NOTE] |
| 116 | + > The ``status`` and ``created_from``/``created_to`` filters are applied |
| 117 | + > client-side after each page of entities is fetched, so an individual |
| 118 | + > page may contain fewer than ``page_size`` matches (or none) even when |
| 119 | + > more matching schedules exist. This mirrors the .NET implementation. |
| 120 | + """ |
| 121 | + prefix = schedule_query.schedule_id_prefix if schedule_query and schedule_query.schedule_id_prefix else "" |
| 122 | + page_size = (schedule_query.page_size if schedule_query and schedule_query.page_size |
| 123 | + else ScheduleQuery.DEFAULT_PAGE_SIZE) |
| 124 | + query = EntityQuery( |
| 125 | + instance_id_starts_with=f"@{ENTITY_NAME}@{prefix}", |
| 126 | + include_state=True, |
| 127 | + page_size=page_size, |
| 128 | + ) |
| 129 | + results: list[ScheduleDescription] = [] |
| 130 | + for metadata in self._client.get_all_entities(query): |
| 131 | + state = metadata.get_typed_state(ScheduleState) |
| 132 | + if state is None or state.schedule_configuration is None: |
| 133 | + continue |
| 134 | + if not self._matches_filter(state, schedule_query): |
| 135 | + continue |
| 136 | + results.append(state.to_description()) |
| 137 | + return results |
| 138 | + |
| 139 | + @staticmethod |
| 140 | + def _matches_filter(state: ScheduleState, schedule_query: ScheduleQuery | None) -> bool: |
| 141 | + if schedule_query is None: |
| 142 | + return True |
| 143 | + if schedule_query.status is not None and state.status != schedule_query.status: |
| 144 | + return False |
| 145 | + # ``ScheduleQuery`` normalizes its bounds to aware UTC; defensively |
| 146 | + # normalize the stored timestamp too (a payload could in principle carry |
| 147 | + # a naive value) so the comparison can never raise on naive-vs-aware. |
| 148 | + # Bounds are exclusive, matching the .NET ScheduledTasks implementation. |
| 149 | + created_at = ensure_aware(state.schedule_created_at) |
| 150 | + if schedule_query.created_from is not None and not (created_at and created_at > schedule_query.created_from): |
| 151 | + return False |
| 152 | + if schedule_query.created_to is not None and not (created_at and created_at < schedule_query.created_to): |
| 153 | + return False |
| 154 | + return True |
0 commit comments