|
21 | 21 | from pyncette.model import ContinuationToken |
22 | 22 | from pyncette.model import ExecutionMode |
23 | 23 | from pyncette.model import Lease |
| 24 | +from pyncette.model import ListTasksResponse |
24 | 25 | from pyncette.model import PollResponse |
25 | 26 | from pyncette.model import QueryResponse |
26 | 27 | from pyncette.model import ResultType |
| 28 | +from pyncette.model import TaskState |
27 | 29 | from pyncette.repository import Repository |
28 | 30 | from pyncette.task import Task |
29 | 31 |
|
@@ -378,6 +380,97 @@ async def unlock_task(self, utc_now: datetime.datetime, task: Task, lease: Lease |
378 | 380 |
|
379 | 381 | raise PyncetteException("Unable to acquire the lock on the task due to contention") |
380 | 382 |
|
| 383 | + async def get_task_state( |
| 384 | + self, |
| 385 | + utc_now: datetime.datetime, |
| 386 | + task: Task, |
| 387 | + ) -> TaskState | None: |
| 388 | + response = await self._table.get_item( |
| 389 | + Key={ |
| 390 | + "partition_id": self._get_partition_id(task), |
| 391 | + "task_id": task.canonical_name, |
| 392 | + } |
| 393 | + ) |
| 394 | + |
| 395 | + if "Item" not in response: |
| 396 | + return None |
| 397 | + |
| 398 | + item = response["Item"] |
| 399 | + |
| 400 | + # For dynamic tasks, re-instantiate from spec to ensure we have fresh parameters |
| 401 | + if task.parent_task is not None: |
| 402 | + task_spec = json.loads(item["task_spec"]) if item.get("task_spec") else None |
| 403 | + if not task_spec: |
| 404 | + raise PyncetteException(f"Task {task.canonical_name} has no task_spec stored") |
| 405 | + instantiated_task = task.parent_task.instantiate_from_spec(task_spec) |
| 406 | + else: |
| 407 | + # Static task - use as-is (no task_spec stored) |
| 408 | + instantiated_task = task |
| 409 | + |
| 410 | + execute_after = item.get("execute_after") |
| 411 | + assert execute_after is not None, "execute_after should not be None for existing tasks" |
| 412 | + scheduled_at = datetime.datetime.fromisoformat(execute_after) |
| 413 | + |
| 414 | + return TaskState( |
| 415 | + task=instantiated_task, |
| 416 | + scheduled_at=scheduled_at, |
| 417 | + locked_until=datetime.datetime.fromisoformat(item["locked_until"]) if item.get("locked_until") else None, |
| 418 | + locked_by=item.get("locked_by"), |
| 419 | + ) |
| 420 | + |
| 421 | + async def list_task_states( |
| 422 | + self, |
| 423 | + utc_now: datetime.datetime, |
| 424 | + parent_task: Task, |
| 425 | + limit: int | None = None, |
| 426 | + continuation_token: ContinuationToken | None = None, |
| 427 | + ) -> ListTasksResponse: |
| 428 | + if limit is None: |
| 429 | + limit = self._batch_size |
| 430 | + |
| 431 | + # Query DynamoDB on partition_id (which encodes parent_name) |
| 432 | + prefix = (self._partition_prefix or "").replace(":", "::") |
| 433 | + partition_id = f"{prefix}:{parent_task.canonical_name}" |
| 434 | + |
| 435 | + query_kwargs: dict[str, Any] = { |
| 436 | + "KeyConditionExpression": Key("partition_id").eq(partition_id), |
| 437 | + "Limit": limit, |
| 438 | + } |
| 439 | + |
| 440 | + if continuation_token is not None: |
| 441 | + query_kwargs["ExclusiveStartKey"] = continuation_token |
| 442 | + |
| 443 | + response = await self._table.query(**query_kwargs) |
| 444 | + |
| 445 | + tasks = [] |
| 446 | + for item in response.get("Items", []): |
| 447 | + task_spec = json.loads(item["task_spec"]) if item.get("task_spec") else None |
| 448 | + if not task_spec: |
| 449 | + logger.warning(f"Task {item['task_id']} has no task_spec, skipping") |
| 450 | + continue |
| 451 | + |
| 452 | + instantiated_task = parent_task.instantiate_from_spec(task_spec) |
| 453 | + |
| 454 | + execute_after = item.get("execute_after") |
| 455 | + assert execute_after is not None, "execute_after should not be None for existing tasks" |
| 456 | + scheduled_at = datetime.datetime.fromisoformat(execute_after) |
| 457 | + |
| 458 | + tasks.append( |
| 459 | + TaskState( |
| 460 | + task=instantiated_task, |
| 461 | + scheduled_at=scheduled_at, |
| 462 | + locked_until=datetime.datetime.fromisoformat(item["locked_until"]) if item.get("locked_until") else None, |
| 463 | + locked_by=item.get("locked_by"), |
| 464 | + ) |
| 465 | + ) |
| 466 | + |
| 467 | + next_token = ContinuationToken(response["LastEvaluatedKey"]) if "LastEvaluatedKey" in response else None |
| 468 | + |
| 469 | + return ListTasksResponse( |
| 470 | + tasks=tasks, |
| 471 | + continuation_token=next_token, |
| 472 | + ) |
| 473 | + |
381 | 474 | async def extend_lease(self, utc_now: datetime.datetime, task: Task, lease: Lease) -> Lease | None: |
382 | 475 | assert isinstance(lease, _TaskRecord) |
383 | 476 | record: _TaskRecord | None = cast(_TaskRecord, copy.copy(lease)) |
|
0 commit comments