ltq 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

ltq/__init__.py

@@ -1,18 +1,24 @@
 from .app import App
-from .utils import dispatch
+from .broker import Broker
 from .task import Task
 from .worker import Worker
 from .scheduler import Scheduler
 from .logger import get_logger
-from .errors import RejectMessage, RetryMessage
+from .errors import RejectError, RetryError
+from .middleware import Middleware, MaxTries, MaxAge, MaxRate, Sentry
 
 __all__ = [
     "App",
+    "Broker",
     "Worker",
     "Scheduler",
     "Task",
-    "dispatch",
     "get_logger",
-    "RejectMessage",
-    "RetryMessage",
+    "RejectError",
+    "RetryError",
+    "Middleware",
+    "MaxTries",
+    "MaxAge",
+    "MaxRate",
+    "Sentry",
 ]

ltq/app.py

@@ -1,14 +1,40 @@
 import asyncio
+import threading
 
+from .middleware import Middleware
 from .worker import Worker
 
 
 class App:
-    def __init__(self) -> None:
-        self.workers: set[Worker] = set()
+    def __init__(self, middlewares: list[Middleware] | None = None) -> None:
+        self.workers: dict[str, Worker] = dict()
+        self.middlewares: list[Middleware] = middlewares or []
+
+    def register_middleware(self, middleware: Middleware, pos: int = -1) -> None:
+        if pos == -1:
+            self.middlewares.append(middleware)
+        else:
+            self.middlewares.insert(pos, middleware)
 
     def register_worker(self, worker: Worker) -> None:
-        self.workers.add(worker)
+        if worker.name in self.workers:
+            raise RuntimeError(f"Worker '{worker.name}' is already registered")
+        worker.middlewares = list(self.middlewares) + worker.middlewares
+        self.workers[worker.name] = worker
+
+    @staticmethod
+    def _run_worker(worker: Worker) -> None:
+        asyncio.run(worker.run())
 
     async def run(self) -> None:
-        await asyncio.gather(*(w.run() for w in self.workers))
+        threads: list[threading.Thread] = []
+        for worker in self.workers.values():
+            t = threading.Thread(target=self._run_worker, args=(worker,), daemon=True)
+            t.start()
+            threads.append(t)
+
+        try:
+            while any(t.is_alive() for t in threads):
+                await asyncio.sleep(0.2)
+        except asyncio.CancelledError:
+            pass

ltq/broker.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from urllib.parse import urlparse
+import uuid
+from collections import defaultdict
+
+import redis.asyncio as aioredis
+
+from .message import Message
+
+
+class Broker:
+    @staticmethod
+    def from_url(url: str) -> Broker:
+        urlp = urlparse(url)
+        if urlp.scheme == "memory":
+            return MemoryBroker()
+        elif urlp.scheme == "redis":
+            return RedisBroker(url)
+        else:
+            raise RuntimeError(f"Unknown scheme: {urlp.scheme}")
+
+    async def close(self) -> None: ...
+    async def publish(self, message: Message, delay: float = 0) -> None: ...
+    async def consume(self, queue: str) -> Message: ...
+    async def ack(self, message: Message) -> None: ...
+    async def nack(
+        self,
+        message: Message,
+        delay: float = 0,
+        drop: bool = False,
+    ) -> None: ...
+    async def len(self, queue: str) -> int: ...
+    async def clear(self, queue: str) -> None: ...
+
+
+class RedisBroker(Broker):
+    def __init__(self, url: str) -> None:
+        self.url = url
+        self._client = aioredis.from_url(url)
+        self._id = uuid.uuid4().hex[:8]
+
+    async def close(self) -> None:
+        await self._client.aclose()
+
+    async def publish(
+        self,
+        message: Message,
+        delay: float = 0,
+    ) -> None:
+        score = time.time() + delay
+        await self._client.zadd(
+            f"queue:{message.task_name}",
+            {
+                message.to_json(): score,
+            },
+        )  # type: ignore
+
+    async def consume(self, queue: str) -> Message:
+        while True:
+            now = time.time()
+            ready = await self._client.zrangebyscore(
+                f"queue:{queue}", 0, now, start=0, num=1
+            )  # type: ignore
+            if ready:
+                msg = ready[0]
+                await self._client.zadd(f"processing:{queue}:{self._id}", {msg: now,})  # type: ignore
+                await self._client.zrem(f"queue:{queue}", msg)  # type: ignore
+                return Message.from_json(msg)
+            await asyncio.sleep(0.1)
+
+    async def ack(self, message: Message) -> None:
+        key = f"processing:{message.task_name}:{self._id}"
+        await self._client.zrem(key, message.to_json())  # type: ignore
+
+    async def nack(
+        self,
+        message: Message,
+        delay: float = 0,
+        drop: bool = False,
+    ) -> None:
+        key = f"processing:{message.task_name}:{self._id}"
+        await self._client.zrem(key, message.to_json())  # type: ignore
+        if not drop:
+            await self.publish(message, delay=delay)
+
+    async def len(self, queue: str) -> int:
+        return await self._client.zcard(f"queue:{queue}") or 0  # type: ignore
+
+    async def clear(self, queue: str) -> None:
+        await self._client.delete(f"queue:{queue}", f"processing:{queue}:{self._id}")  # type: ignore
+
+
+class MemoryBroker(Broker):
+    def __init__(self) -> None:
+        self._queues: defaultdict[str, dict[str, float]] = defaultdict(dict)
+
+    async def close(self) -> None:
+        pass
+
+    async def publish(
+        self,
+        message: Message,
+        delay: float = 0,
+    ) -> None:
+        self._queues[message.task_name][message.to_json()] = time.time() + delay
+
+    async def consume(self, queue: str) -> Message:
+        while True:
+            now = time.time()
+            for msg, score in list(self._queues[queue].items()):
+                if score <= now:
+                    del self._queues[queue][msg]
+                    return Message.from_json(msg)
+            await asyncio.sleep(0.1)
+
+    async def ack(self, message: Message) -> None:
+        pass
+
+    async def nack(
+        self,
+        message: Message,
+        delay: float = 0,
+        drop: bool = False,
+    ) -> None:
+        if not drop:
+            await self.publish(message, delay=delay)
+
+    async def len(self, queue: str) -> int:
+        return len(self._queues[queue])
+
+    async def clear(self, queue: str) -> None:
+        self._queues.pop(queue, None)

ltq/cli.py

@@ -8,6 +8,7 @@ import argparse
 
 from .logger import setup_logging, get_logger
 from .app import App
+from .broker import Broker
 from .worker import Worker
 
 logger = get_logger()
@@ -37,74 +38,142 @@ def import_from_string(import_str: str):
         sys.exit(1)
 
 
+async def clear_queue(
+    task_name: str,
+    url: str = "redis://localhost:6379",
+) -> None:
+    """Clear a queue for a specific task."""
+    broker = Broker.from_url(url)
+    try:
+        await broker.clear(task_name)
+        logger.info(f"Cleared queue for task: {task_name}")
+    finally:
+        await broker.close()
+
+
+async def get_queue_size(
+    task_name: str,
+    url: str = "redis://localhost:6379",
+) -> int:
+    """Get the size of a queue for a specific task."""
+    broker = Broker.from_url(url)
+    try:
+        return await broker.len(task_name)
+    finally:
+        await broker.close()
+
+
 def main():
     """Run a ltq worker."""
 
     parser = argparse.ArgumentParser(
         prog="ltq",
-        description="Run a ltq worker",
+        description="Run a ltq worker or manage queues",
         formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="Example:\n  ltq example:worker --concurrency 100",
+        epilog="Examples:\n  ltq run examples:worker --concurrency 100\n  ltq clear emails:send_email",
     )
 
-    parser.add_argument(
+    subparsers = parser.add_subparsers(dest="command", help="Command to run")
+
+    # Run command
+    run_parser = subparsers.add_parser("run", help="Run a worker or app")
+    run_parser.add_argument(
         "worker", nargs="?", help="Worker import string (module:attribute)"
     )
-    parser.add_argument(
+    run_parser.add_argument(
         "--app", dest="app", help="App import string (module:attribute)"
     )
-    parser.add_argument("--concurrency", type=int, help="Override worker concurrency")
-    parser.add_argument("--poll-sleep", type=float, help="Override worker poll sleep")
-    parser.add_argument(
+    run_parser.add_argument(
+        "--concurrency", type=int, help="Override worker concurrency"
+    )
+    run_parser.add_argument(
         "--log-level",
         type=str,
         default="INFO",
         choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
         help="Set logging level (default: INFO)",
     )
-    args = parser.parse_args()
 
-    if not args.worker and not args.app:
-        parser.error("either worker or --app is required")
-    if args.worker and args.app:
-        parser.error("cannot specify both worker and --app")
+    # Clear command
+    clear_parser = subparsers.add_parser("clear", help="Clear a task queue")
+    clear_parser.add_argument("task_name", help="Task name (namespace:function)")
+    clear_parser.add_argument(
+        "--redis-url",
+        default="redis://localhost:6379",
+        help="Redis URL (default: redis://localhost:6379)",
+    )
+    clear_parser.add_argument(
+        "--log-level",
+        type=str,
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        help="Set logging level (default: INFO)",
+    )
+
+    # Size command
+    size_parser = subparsers.add_parser("size", help="Get queue size for a task")
+    size_parser.add_argument("task_name", help="Task name (namespace:function)")
+    size_parser.add_argument(
+        "--redis-url",
+        default="redis://localhost:6379",
+        help="Redis URL (default: redis://localhost:6379)",
+    )
+
+    args = parser.parse_args()
 
     # Setup colored logging for CLI
-    setup_logging(level=args.log_level)
-    if args.log_level:
+    setup_logging(level=getattr(args, "log_level", "INFO"))
+    if hasattr(args, "log_level") and args.log_level:
         logger.setLevel(args.log_level)
 
-    if args.app:
-        app: App = import_from_string(args.app)
+    # Handle clear command
+    if args.command == "clear":
+        asyncio.run(clear_queue(args.task_name, args.redis_url))
+        return
+
+    # Handle size command
+    if args.command == "size":
+        size = asyncio.run(get_queue_size(args.task_name, args.redis_url))
+        print(f"{args.task_name}: {size}")
+        return
+
+    # Handle run command
+    if args.command == "run":
+        if not args.worker and not args.app:
+            run_parser.error("either worker or --app is required")
+        if args.worker and args.app:
+            run_parser.error("cannot specify both worker and --app")
+
+        if args.app:
+            app: App = import_from_string(args.app)
+
+            for w in app.workers.values():
+                if args.concurrency:
+                    w.concurrency = args.concurrency
+
+            logger.info("Starting ltq app")
+
+            try:
+                asyncio.run(app.run())
+            except KeyboardInterrupt:
+                logger.info("Shutting down...")
+        else:
+            worker: Worker = import_from_string(args.worker)
 
-        for w in app.workers:
             if args.concurrency:
-                w.concurrency = args.concurrency
-            if args.poll_sleep:
-                w.poll_sleep = args.poll_sleep
-
-        logger.info("Starting ltq app")
-        logger.info("App: %s (%d workers)", args.app, len(app.workers))
-
-        try:
-            asyncio.run(app.run())
-        except KeyboardInterrupt:
-            logger.info("Shutting down...")
-    else:
-        worker: Worker = import_from_string(args.worker)
-
-        if args.concurrency:
-            worker.concurrency = args.concurrency
-        if args.poll_sleep:
-            worker.poll_sleep = args.poll_sleep
-
-        logger.info("Starting ltq worker")
-        logger.info("Worker: %s", args.worker)
-
-        try:
-            asyncio.run(worker.run())
-        except KeyboardInterrupt:
-            logger.info("Shutting down...")
+                worker.concurrency = args.concurrency
+
+            logger.info("Starting ltq worker")
+
+            try:
+                asyncio.run(worker.run())
+            except KeyboardInterrupt:
+                logger.info("Shutting down...")
+        return
+
+    # No command specified
+    parser.print_help()
+    sys.exit(1)
 
 
 if __name__ == "__main__":

ltq/errors.py

@@ -1,10 +1,10 @@
-class RejectMessage(Exception):
+class RejectError(Exception):
     """Signal that a message should be dropped."""
 
     pass
 
 
-class RetryMessage(Exception):
+class RetryError(Exception):
     """Signal that a message should be retried after a delay."""
 
     def __init__(self, delay: float = 0.0, message: str = "") -> None:

ltq/logger.py

@@ -1,5 +1,6 @@
 import logging
 import sys
+import threading
 
 
 class ColoredFormatter(logging.Formatter):
@@ -25,14 +26,18 @@ class ColoredFormatter(logging.Formatter):
     def format(self, record: logging.LogRecord) -> str:
         """Format log record with colored severity level."""
         color = self.COLORS.get(record.levelno, self.RESET)
-        log_time = self.formatTime(record, "%H:%M:%S")
-        timestamp = f"{self.GRAY}{log_time}{self.RESET}"
+        log_time = self.formatTime(record, "%Y-%m-%d %H:%M:%S")
+        log_time_ms = f"{log_time}.{int(record.msecs):03d}"
+        timestamp = f"{self.GRAY}{log_time_ms}{self.RESET}"
 
-        levelname = f"{color}{record.levelname:<8}{self.RESET}"
-        workername = f"{self.CYAN}{record.name:<8}{self.RESET}"
+        levelname = f"{color}{record.levelname:<4}{self.RESET}"
+
+        name = record.name.removeprefix("ltq.")
+        workername = f"{self.CYAN}{name:<6}{self.RESET}"
+        thread_id = f"{self.GRAY}[{threading.current_thread().ident}]{self.RESET}"
 
         message = record.getMessage()
-        log_line = f"{timestamp} {levelname} {workername} {message}"
+        log_line = f"{timestamp} {thread_id} {levelname} {workername} {message}"
 
         if record.exc_info and not record.exc_text:
             record.exc_text = self.formatException(record.exc_info)
@@ -60,4 +65,6 @@ def setup_logging(level: int | str = logging.INFO) -> None:
 
 
 def get_logger(name: str = "ltq") -> logging.Logger:
-    return logging.getLogger(name)
+    if name == "ltq":
+        return logging.getLogger("ltq")
+    return logging.getLogger(f"ltq.{name}")

ltq/message.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 import json
+import time
 import uuid
 from dataclasses import dataclass, field
-from typing import Any, TYPE_CHECKING
+from typing import Any
 
-if TYPE_CHECKING:
-    from .task import Task
+def _default_ctx() -> dict[str, Any]:
+    return {"created_at": time.time()}
 
 
 @dataclass
@@ -14,9 +15,8 @@ class Message:
     args: tuple[Any, ...]
     kwargs: dict[str, Any]
     task_name: str
-    task: Task | None = None  # only set when Message created with Task.message
-    ctx: dict[str, Any] = field(default_factory=dict)
-    id: str = field(default_factory=lambda: uuid.uuid4().hex)
+    ctx: dict[str, Any] = field(default_factory=_default_ctx)
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
 
     def to_json(self) -> str:
         return json.dumps(