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