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
|
from functools import wraps
from typing import Any, Callable
from urllib.parse import urlsplit, urlunsplit
from flask import Flask
class Huey:
def __init__(self, app: Flask | None = None):
self.app: Flask | None = None
self._huey = None
if app is not None:
self.init_app(app)
def init_app(self, app: Flask):
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):
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):
def decorator(func: Callable):
@wraps(func)
def wrapper(*args: Any, **kwargs: Any):
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):
def decorator(func: Callable):
@wraps(func)
def wrapper(*args: Any, **kwargs: Any):
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):
"""
Forward unknown attributes to the real Huey instance.
This lets you still use things like:
huey.enqueue(...)
huey.scheduled()
huey.pending()
"""
return getattr(self.huey, name)
@staticmethod
def backend_from_url(url: str) -> tuple[Any, dict[str, str]]:
parsed = urlsplit(url)
scheme = parsed.scheme.lower()
if scheme.startswith("redis") or scheme.startswith("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}")
|
