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
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
|
# SPDX-FileCopyrightText: 2026 Dennis Fink <me+coding@dennisfink.me>
#
# SPDX-License-Identifier: BSD-3-Clause
from collections.abc import Callable
from functools import wraps
from typing import Any, ParamSpec, TypeVar
from urllib.parse import urlsplit, urlunsplit
from flask import Flask
from huey import Huey as BaseHuey
from huey.api import TaskWrapper
P = ParamSpec("P")
R = TypeVar("R")
class Huey:
"""
Provide Flask integration for a Huey instance.
The extension initializes a Huey backend from the Flask configuration and
wraps tasks so that they execute within an application context.
"""
def __init__(self, app: Flask | None = None) -> None:
"""
Initialize the Huey extension.
:param app: Flask application to initialize immediately, if provided.
"""
self.app: Flask | None = None
self._huey: BaseHuey | None = None
if app is not None:
self.init_app(app)
def init_app(self, app: Flask) -> None:
"""
Initialize Huey for a Flask application.
The backend and its storage options are derived from ``HUEY_URL`` and the
resulting extension is registered with the application.
:param app: Flask application to initialize.
:raises RuntimeError: If the configured Huey backend URL is invalid or uses
an unsupported scheme.
"""
config: dict[str, Any] = {
"name": app.import_name,
"results": True,
"store_none": False,
"utc": True,
"immediate": app.config.get("TESTING", False),
**app.config.get("HUEY", {}),
}
url = app.config.get("HUEY_URL", config.pop("url", "memory://"))
huey_class, storage_kwargs = self.backend_from_url(url)
self.app = app
self._huey = huey_class(**config, **storage_kwargs)
app.extensions["huey"] = self
@property
def huey(self) -> BaseHuey:
"""
Return the initialized Huey instance.
:return: Configured Huey backend instance.
:raises RuntimeError: If the extension has not been initialized.
"""
if self._huey is None:
raise RuntimeError(
"Huey has not been initialized. "
"Call huey.init_app(app) before importing tasks."
)
return self._huey
def task(
self, *task_args: Any, **task_kwargs: Any
) -> Callable[[Callable[P, object]], TaskWrapper]:
"""
Create a Huey task that runs within the Flask application context.
:param task_args: Positional arguments forwarded to Huey's task decorator.
:param task_kwargs: Keyword arguments forwarded to Huey's task decorator.
:return: Decorator that registers the wrapped function as a Huey task.
"""
def decorator(func: Callable[P, R]) -> TaskWrapper:
@wraps(func)
def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
if self.app is None:
raise RuntimeError("Flask app is not available.")
with self.app.app_context():
return func(*args, **kwargs)
return self.huey.task(*task_args, **task_kwargs)(wrapper)
return decorator
def periodic_task(
self, *task_args: Any, **task_kwargs: Any
) -> Callable[[Callable[P, object]], TaskWrapper]:
"""
Create a periodic Huey task that runs within the Flask application context.
:param task_args: Positional arguments forwarded to Huey's periodic task
decorator.
:param task_kwargs: Keyword arguments forwarded to Huey's periodic task
decorator.
:return: Decorator that registers the wrapped function as a periodic Huey
task.
"""
def decorator(func: Callable[P, R]) -> TaskWrapper:
@wraps(func)
def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
if self.app is None:
raise RuntimeError("Flask app is not available.")
with self.app.app_context():
return func(*args, **kwargs)
return self.huey.periodic_task(*task_args, **task_kwargs)(wrapper)
return decorator
def __getattr__(self, name: str) -> Any:
"""
Forward an unknown attribute to the underlying Huey instance.
:param name: Name of the attribute to retrieve.
:return: Attribute from the initialized Huey instance.
:raises RuntimeError: If the extension has not been initialized.
"""
return getattr(self.huey, name)
@staticmethod
def backend_from_url(url: str) -> tuple[type[BaseHuey], dict[str, Any]]:
"""
Determine the Huey backend and storage options from a URL.
:param url: Huey backend URL.
:return: Huey backend class and keyword arguments for its storage backend.
:raises RuntimeError: If the URL is malformed for the selected backend or
uses an unsupported scheme.
"""
parsed = urlsplit(url)
scheme = parsed.scheme.lower()
if scheme.startswith(("redis", "rediss")):
fixed_url = urlunsplit(parsed._replace(scheme=scheme.split("+", 1)[0]))
if scheme.endswith("priority+expire"):
from huey import PriorityRedisExpireHuey
return PriorityRedisExpireHuey, {"url": fixed_url}
elif scheme.endswith("priority"):
from huey import PriorityRedisHuey
return PriorityRedisHuey, {"url": fixed_url}
elif scheme.endswith("expire"):
from huey import RedisExpireHuey
return RedisExpireHuey, {"url": fixed_url}
else:
from huey import RedisHuey
return RedisHuey, {"url": url}
elif scheme == "sqlite":
from huey import SqliteHuey
prefix = "sqlite:///"
if not url.startswith(prefix):
raise RuntimeError(
"SQLite Huey URLs must look like sqlite:///var/huey.db"
)
filename = url.removeprefix(prefix)
if not filename:
raise RuntimeError("SQLite Huey URL must include a database path.")
return SqliteHuey, {"filename": filename}
elif scheme == "file":
from huey import FileHuey
prefix = "file:///"
if not url.startswith(prefix):
raise RuntimeError(
"File Huey URLs must look like file:///var/huey-queue"
)
path = url.removeprefix(prefix)
if not path:
raise RuntimeError("File Huey URL must include a directory path.")
return FileHuey, {"path": path}
elif scheme in {"postgres", "postgresql"}:
from huey import PostgresHuey
return PostgresHuey, {"dsn": url}
elif scheme == "memory":
from huey import MemoryHuey
return MemoryHuey, {}
elif scheme == "blackhole":
from huey import BlackHoleHuey
return BlackHoleHuey, {}
raise RuntimeError(f"Unsupported HUEY_URL scheme: {scheme!r}")
|