forked from taskiq-python/taskiq
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtask.py
More file actions
163 lines (132 loc) · 4.61 KB
/
Copy pathtask.py
File metadata and controls
163 lines (132 loc) · 4.61 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
import asyncio
from abc import ABC, abstractmethod
from time import time
from typing import TYPE_CHECKING, Any, Coroutine, Generic, Optional, Union
from typing_extensions import TypeVar
from taskiq.exceptions import (
ResultGetError,
ResultIsReadyError,
TaskiqResultTimeoutError,
)
if TYPE_CHECKING: # pragma: no cover
from taskiq.abc.result_backend import AsyncResultBackend
from taskiq.depends.progress_tracker import TaskProgress
from taskiq.result import TaskiqResult
_ReturnType = TypeVar("_ReturnType")
class _Task(ABC, Generic[_ReturnType]):
"""TaskiqTask interface."""
@abstractmethod
def is_ready(self) -> Union[bool, Coroutine[Any, Any, bool]]:
"""
Method to check wether result is ready.
:return: True if result is ready.
"""
@abstractmethod
def get_result(
self,
with_logs: bool = False,
) -> Union[
"TaskiqResult[_ReturnType]",
Coroutine[Any, Any, "TaskiqResult[_ReturnType]"],
]:
"""
Get actual execution result.
:param with_logs: wether you want to fetch logs.
:return: TaskiqResult.
"""
@abstractmethod
def wait_result(
self,
check_interval: float = 0.2,
timeout: float = -1.0,
with_logs: bool = False,
) -> Union[
"TaskiqResult[_ReturnType]",
Coroutine[Any, Any, "TaskiqResult[_ReturnType]"],
]:
"""
Wait for result to become ready and get it.
This function constantly checks whether result is ready
and fetches it when it becomes available.
:param check_interval: how often availability is checked.
:param timeout: maximum amount of time it will wait
before raising TaskiqResultTimeoutError.
:param with_logs: whether you need to download logs.
:return: TaskiqResult.
"""
@abstractmethod
def get_progress(
self,
) -> Union[
"Optional[TaskProgress[Any]]",
Coroutine[Any, Any, "Optional[TaskProgress[Any]]"],
]:
"""
Get task progress.
:return: task's progress.
"""
class AsyncTaskiqTask(_Task[_ReturnType]):
"""AsyncTask for AsyncResultBackend."""
def __init__(
self,
task_id: str,
result_backend: "AsyncResultBackend[_ReturnType]",
) -> None:
self.task_id = task_id
self.result_backend = result_backend
async def is_ready(self) -> bool:
"""
Checks if task is completed.
:raises ResultIsReadyError: if we can't get info about task readiness.
:return: True if task is completed.
"""
try:
return await self.result_backend.is_result_ready(self.task_id)
except Exception as exc:
raise ResultIsReadyError from exc
async def get_result(self, with_logs: bool = False) -> "TaskiqResult[_ReturnType]":
"""
Get result of a task from result backend.
:param with_logs: whether you want to fetch logs from worker.
:raises ResultGetError: if we can't get result from ResultBackend.
:return: task's return value.
"""
try:
return await self.result_backend.get_result(
self.task_id,
with_logs=with_logs,
)
except Exception as exc:
raise ResultGetError from exc
async def wait_result(
self,
check_interval: float = 0.2,
timeout: float = -1.0,
with_logs: bool = False,
) -> "TaskiqResult[_ReturnType]":
"""
Waits until result is ready.
This method just checks whether the task is
ready. And if it is it returns the result.
It may throw TaskiqResultTimeoutError if
task didn't became ready in provided
period of time.
:param check_interval: How often checks are performed.
:param timeout: timeout for the result.
:param with_logs: whether you want to fetch logs from worker.
:raises TaskiqResultTimeoutError: if task didn't
become ready in provided period of time.
:return: task's return value.
"""
start_time = time()
while not await self.is_ready():
await asyncio.sleep(check_interval)
if 0 < timeout < time() - start_time:
raise TaskiqResultTimeoutError
return await self.get_result(with_logs=with_logs)
async def get_progress(self) -> "Optional[TaskProgress[Any]]":
"""
Get task progress.
:return: task's progress.
"""
return await self.result_backend.get_progress(self.task_id)