1 files changed, 159 insertions, 0 deletions
diff --git a/webmentions_ssg/tasks/extension.py b/webmentions_ssg/tasks/extension.py
new file mode 100644
index 0000000..ebdaa76
--- /dev/null
+++ b/webmentions_ssg/tasks/extension.py
@@ -0,0 +1,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}")