taskmanager-engine 0.1.0__py3-none-any.whl

taskmanager/core/task.py

@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+import redis.asyncio as redis
+
+from taskmanager.config import settings
+from taskmanager.core.broker import RedisBroker
+from taskmanager.core.job import Job
+
+
+class TaskRegistry:
+    """Registry holding all user-defined background task functions."""
+
+    def __init__(self):
+        self._tasks: dict[str, Task] = {}
+        self._default_broker: RedisBroker | None = None
+
+    def register(self, task: Task) -> None:
+        self._tasks[task.name] = task
+
+    def get(self, name: str) -> Task | None:
+        return self._tasks.get(name)
+
+    def list_tasks(self) -> list[str]:
+        return list(self._tasks.keys())
+
+    def get_broker(self) -> RedisBroker:
+        if self._default_broker is None:
+            client = redis.from_url(settings.redis_url, decode_responses=True)
+            self._default_broker = RedisBroker(client, prefix=settings.redis_prefix)
+        return self._default_broker
+
+    def set_broker(self, broker: RedisBroker) -> None:
+        self._default_broker = broker
+
+    def task(
+        self,
+        name: str | None = None,
+        queue: str = "default",
+        max_retries: int = 3,
+        retry_backoff: float = 2.0,
+        timeout: float | None = None,
+    ) -> Callable[..., Task]:
+        """Decorator to register a function directly to this registry."""
+        def decorator(func: Callable[..., Any]) -> Task:
+            t = Task(
+                func=func,
+                name=name,
+                queue=queue,
+                max_retries=max_retries,
+                retry_backoff=retry_backoff,
+                timeout=timeout,
+                broker=self._default_broker,
+            )
+            self.register(t)
+            return t
+
+        return decorator
+
+
+registry = TaskRegistry()
+
+
+class Task:
+    """Represents a background task wrapper around a Python callable or coroutine."""
+
+    def __init__(
+        self,
+        func: Callable[..., Any],
+        name: str | None = None,
+        queue: str = "default",
+        max_retries: int = 3,
+        retry_backoff: float = 2.0,
+        timeout: float | None = None,
+        broker: RedisBroker | None = None,
+    ):
+        self.func = func
+        self.name = name or f"{func.__module__}.{func.__name__}"
+        self.queue = queue
+        self.max_retries = max_retries
+        self.retry_backoff = retry_backoff
+        self.timeout = timeout
+        self._broker = broker
+        self.is_async = inspect.iscoroutinefunction(func)
+
+    def get_signature_info(self) -> dict[str, Any]:
+        """Extracts parameters, types, defaults, and generates a sample payload."""
+        try:
+            sig = inspect.signature(self.func)
+            params = []
+            sample_kwargs: dict[str, Any] = {}
+            for param in sig.parameters.values():
+                has_default = param.default is not inspect.Parameter.empty
+                default_val = param.default if has_default else None
+                ann_str = "Any" if param.annotation is inspect.Parameter.empty else str(param.annotation)
+
+                params.append(
+                    {
+                        "name": param.name,
+                        "has_default": has_default,
+                        "default": default_val,
+                        "annotation": ann_str,
+                    }
+                )
+
+                # Heuristics for realistic smart sample values
+                if has_default:
+                    if default_val is not None:
+                        sample_kwargs[param.name] = default_val
+                else:
+                    ann_lower = ann_str.lower()
+                    p_name = param.name.lower()
+                    if "cwd" in p_name:
+                        pass  # Keep optional cwd omitted from sample
+                    elif "int" in ann_lower or "year" in p_name or "month" in p_name or "count" in p_name:
+                        sample_kwargs[param.name] = 2026 if "year" in p_name else (8 if "month" in p_name else 1)
+                    elif "float" in ann_lower:
+                        sample_kwargs[param.name] = 10.0
+                    elif "bool" in ann_lower or "dry_run" in p_name or "force" in p_name:
+                        sample_kwargs[param.name] = False
+                    elif "list" in ann_lower:
+                        sample_kwargs[param.name] = []
+                    elif "dict" in ann_lower:
+                        sample_kwargs[param.name] = {}
+                    elif "email" in p_name:
+                        sample_kwargs[param.name] = "cliente@exemplo.com"
+                    elif "name" in p_name or "nome" in p_name:
+                        sample_kwargs[param.name] = "Carlos Silva"
+                    elif "order" in p_name or "id" in p_name:
+                        sample_kwargs[param.name] = "PED-12345"
+                    elif "command" in p_name:
+                        sample_kwargs[param.name] = "python scripts/backup_database.py --compress"
+                    elif "path" in p_name:
+                        sample_kwargs[param.name] = "scripts/backup_database.py"
+                    else:
+                        sample_kwargs[param.name] = f"valor_{param.name}"
+
+            doc = inspect.getdoc(self.func) or ""
+            return {
+                "parameters": params,
+                "sample_kwargs": sample_kwargs,
+                "docstring": doc,
+            }
+        except Exception:
+            return {
+                "parameters": [],
+                "sample_kwargs": {},
+                "docstring": "",
+            }
+
+    @property
+    def broker(self) -> RedisBroker:
+        return self._broker or registry.get_broker()
+
+    def set_broker(self, broker: RedisBroker) -> None:
+        self._broker = broker
+
+    async def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Direct call execution for testing or inline execution."""
+        if self.is_async:
+            return await self.func(*args, **kwargs)
+        return self.func(*args, **kwargs)
+
+    async def delay(self, *args: Any, **kwargs: Any) -> Job:
+        """Enqueues the task immediately into the default queue."""
+        return await self.apply_async(args=list(args), kwargs=kwargs)
+
+    async def apply_async(
+        self,
+        args: list[Any] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        queue: str | None = None,
+        delay: float | None = None,
+        priority: int = 0,
+        max_retries: int | None = None,
+        retry_backoff: float | None = None,
+        timeout: float | None = None,
+        idempotency_key: str | None = None,
+    ) -> Job:
+        """Enqueues the task with custom execution parameters."""
+        job = Job(
+            task_name=self.name,
+            queue=queue or self.queue,
+            args=args or [],
+            kwargs=kwargs or {},
+            priority=priority,
+            max_retries=self.max_retries if max_retries is None else max_retries,
+            retry_backoff=self.retry_backoff if retry_backoff is None else retry_backoff,
+            timeout=self.timeout if timeout is None else timeout,
+            idempotency_key=idempotency_key,
+        )
+
+        if delay and delay > 0:
+            return await self.broker.schedule_delayed(job, delay_seconds=delay)
+        return await self.broker.enqueue(job)
+
+
+def task(
+    name: str | None = None,
+    queue: str = "default",
+    max_retries: int = 3,
+    retry_backoff: float = 2.0,
+    timeout: float | None = None,
+    broker: RedisBroker | None = None,
+) -> Callable[[Callable[..., Any]], Task]:
+    """Decorator to register a function as a TaskManager background task."""
+
+    def decorator(fn: Callable[..., Any]) -> Task:
+        t = Task(
+            func=fn,
+            name=name or f"{fn.__module__}.{fn.__name__}",
+            queue=queue,
+            max_retries=max_retries,
+            retry_backoff=retry_backoff,
+            timeout=timeout,
+            broker=broker,
+        )
+        registry.register(t)
+        return t
+
+    return decorator

taskmanager/scheduler/cron.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import time
+import uuid
+from datetime import UTC, datetime
+from enum import Enum
+from typing import Any
+
+from croniter import croniter
+from pydantic import BaseModel, Field
+
+
+class ScheduleType(str, Enum):
+    CRON = "cron"
+    INTERVAL = "interval"
+
+
+class Schedule(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    name: str
+    task_name: str
+    queue: str = "default"
+    schedule_type: ScheduleType = ScheduleType.CRON
+    cron_expression: str | None = None  # e.g. "*/5 * * * *"
+    interval_seconds: float | None = None
+    args: list[Any] = Field(default_factory=list)
+    kwargs: dict[str, Any] = Field(default_factory=dict)
+    enabled: bool = True
+    last_run: float | None = None
+    next_run: float | None = None
+    total_runs: int = 0
+    created_at: float = Field(default_factory=time.time)
+
+
+def calculate_next_run(schedule: Schedule, from_timestamp: float | None = None) -> float:
+    """Calculates the next unix timestamp for a given Schedule."""
+    base_time = from_timestamp if from_timestamp is not None else time.time()
+
+    if schedule.schedule_type == ScheduleType.CRON:
+        if not schedule.cron_expression or not croniter.is_valid(schedule.cron_expression):
+            raise ValueError(f"Invalid cron expression: {schedule.cron_expression}")
+        base_dt = datetime.fromtimestamp(base_time, tz=UTC).astimezone()
+        iter_cron = croniter(schedule.cron_expression, base_dt)
+        next_dt = iter_cron.get_next(datetime)
+        return next_dt.timestamp()
+
+    elif schedule.schedule_type == ScheduleType.INTERVAL:
+        if not schedule.interval_seconds or schedule.interval_seconds <= 0:
+            raise ValueError(f"Invalid interval seconds: {schedule.interval_seconds}")
+        return base_time + schedule.interval_seconds
+
+    raise ValueError(f"Unknown schedule type: {schedule.schedule_type}")

taskmanager/scheduler/scheduler.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+
+from taskmanager.config import settings
+from taskmanager.core.broker import RedisBroker
+from taskmanager.core.job import Job
+from taskmanager.scheduler.cron import Schedule, calculate_next_run
+
+logger = logging.getLogger(__name__)
+
+
+class Scheduler:
+    """Dynamic Cron and Interval scheduler utilizing Redis for persistence and distributed locking."""
+
+    def __init__(self, broker: RedisBroker):
+        self.broker = broker
+        self._running = False
+        self._task: asyncio.Task[None] | None = None
+        self._lock_id = f"scheduler-{time.time()}"
+
+    # --- Schedule CRUD Operations ---
+    async def add_schedule(self, schedule: Schedule) -> Schedule:
+        """Registers a new schedule and computes initial next_run."""
+        if schedule.next_run is None:
+            schedule.next_run = calculate_next_run(schedule)
+        key = self.broker._key_schedules()
+        await self.broker.redis.hset(key, schedule.id, schedule.model_dump_json())
+        await self.broker.publish_event("schedule:created", schedule.model_dump())
+        return schedule
+
+    async def get_schedule(self, schedule_id: str) -> Schedule | None:
+        """Fetches a schedule by ID."""
+        key = self.broker._key_schedules()
+        raw = await self.broker.redis.hget(key, schedule_id)
+        if not raw:
+            return None
+        return Schedule.model_validate_json(raw)
+
+    async def list_schedules(self) -> list[Schedule]:
+        """Lists all registered schedules."""
+        key = self.broker._key_schedules()
+        all_raw = await self.broker.redis.hgetall(key)
+        schedules: list[Schedule] = []
+        for raw in all_raw.values():
+            try:
+                schedules.append(Schedule.model_validate_json(raw))
+            except Exception:
+                pass
+        return sorted(schedules, key=lambda s: s.created_at)
+
+    async def update_schedule(self, schedule: Schedule) -> Schedule:
+        """Updates an existing schedule."""
+        if schedule.enabled:
+            schedule.next_run = calculate_next_run(schedule)
+        key = self.broker._key_schedules()
+        await self.broker.redis.hset(key, schedule.id, schedule.model_dump_json())
+        await self.broker.publish_event("schedule:updated", schedule.model_dump())
+        return schedule
+
+    async def delete_schedule(self, schedule_id: str) -> bool:
+        """Deletes a schedule."""
+        key = self.broker._key_schedules()
+        res = await self.broker.redis.hdel(key, schedule_id)
+        if res > 0:
+            await self.broker.publish_event("schedule:deleted", {"schedule_id": schedule_id})
+            return True
+        return False
+
+    async def toggle_schedule(self, schedule_id: str, enabled: bool) -> Schedule | None:
+        """Enables or disables a schedule."""
+        schedule = await self.get_schedule(schedule_id)
+        if not schedule:
+            return None
+        schedule.enabled = enabled
+        if enabled:
+            schedule.next_run = calculate_next_run(schedule)
+        else:
+            schedule.next_run = None
+        return await self.update_schedule(schedule)
+
+    async def trigger_now(self, schedule_id: str) -> Job | None:
+        """Manually triggers a schedule immediately."""
+        schedule = await self.get_schedule(schedule_id)
+        if not schedule:
+            return None
+
+        job = Job(
+            task_name=schedule.task_name,
+            queue=schedule.queue,
+            args=schedule.args,
+            kwargs=schedule.kwargs,
+        )
+        enqueued_job = await self.broker.enqueue(job)
+
+        schedule.last_run = time.time()
+        schedule.total_runs += 1
+        key = self.broker._key_schedules()
+        await self.broker.redis.hset(key, schedule.id, schedule.model_dump_json())
+        await self.broker.publish_event(
+            "schedule:triggered",
+            {
+                "schedule_id": schedule.id,
+                "job_id": enqueued_job.id,
+                "manual": True,
+            },
+        )
+        return enqueued_job
+
+    # --- Scheduler Daemon Loop ---
+    async def start(self) -> None:
+        """Starts the scheduler polling loop."""
+        self._running = True
+        logger.info("Scheduler daemon started.")
+        try:
+            while self._running:
+                await self._tick()
+                await asyncio.sleep(settings.scheduler_poll_interval)
+        except asyncio.CancelledError:
+            pass
+        finally:
+            self._running = False
+            logger.info("Scheduler daemon stopped.")
+
+    async def stop(self) -> None:
+        self._running = False
+
+    async def _tick(self) -> None:
+        """Evaluates schedules and enqueues due jobs using distributed locking."""
+        lock_key = self.broker._key_lock("scheduler")
+        # Attempt to acquire leader lock with 3-second TTL
+        acquired = await self.broker.redis.set(lock_key, self._lock_id, nx=True, ex=3)
+        if not acquired:
+            # Check if this instance already holds the lock
+            current_holder = await self.broker.redis.get(lock_key)
+            if current_holder != self._lock_id:
+                return  # Another scheduler instance is the active leader
+            # Refresh lock
+            await self.broker.redis.expire(lock_key, 3)
+
+        now = time.time()
+        schedules = await self.list_schedules()
+
+        for schedule in schedules:
+            if not schedule.enabled or schedule.next_run is None:
+                continue
+
+            if now >= schedule.next_run:
+                # Enqueue the job
+                job = Job(
+                    task_name=schedule.task_name,
+                    queue=schedule.queue,
+                    args=schedule.args,
+                    kwargs=schedule.kwargs,
+                )
+                await self.broker.enqueue(job)
+
+                # Advance next run
+                schedule.last_run = now
+                schedule.total_runs += 1
+                try:
+                    schedule.next_run = calculate_next_run(schedule, from_timestamp=now)
+                except Exception as err:
+                    logger.error(f"Error calculating next run for schedule {schedule.id}: {err}")
+                    schedule.next_run = now + 60
+
+                key = self.broker._key_schedules()
+                await self.broker.redis.hset(key, schedule.id, schedule.model_dump_json())
+                await self.broker.publish_event(
+                    "schedule:triggered",
+                    {
+                        "schedule_id": schedule.id,
+                        "job_id": job.id,
+                        "manual": False,
+                    },
+                )