cli.py

#!/usr/bin/env python3
"""
Command-line interface for the task queue.

Examples:
    python cli.py submit add --args '[2, 3]'
    python cli.py submit flaky_task --kwargs '{"fail_probability": 0.5}' --priority 5
    python cli.py submit always_fails --max-retries 2
    python cli.py status <task_id>
    python cli.py list --status dead
    python cli.py stats
    python cli.py requeue <task_id>

    # Phase 2: task chaining
    python cli.py submit process_report --depends-on <task_id_of_fetch_step>

    # Phase 2: recurring schedules
    python cli.py schedule-add add --interval 30 --args '[2, 3]'
    python cli.py schedule-list
    python cli.py schedule-disable <schedule_id>
    python cli.py schedule-enable <schedule_id>
"""

import argparse
import json
import sys

# Registers the example task functions so the CLI can enqueue them by name.
import examples.sample_tasks  # noqa: F401
from task_queue.broker import Broker
from task_queue.schedule import Schedule
from task_queue.task import Task, TaskStatus


def cmd_submit(args, broker: Broker):
    depends_on = tuple(x.strip() for x in args.depends_on.split(",") if x.strip()) if args.depends_on else ()
    task = Task(
        func_name=args.func_name,
        args=tuple(json.loads(args.args)),
        kwargs=json.loads(args.kwargs),
        priority=args.priority,
        max_retries=args.max_retries,
        depends_on=depends_on,
    )
    task_id = broker.enqueue(task)
    note = f" (waiting on {len(depends_on)} dependenc{'y' if len(depends_on) == 1 else 'ies'})" if depends_on else ""
    print(f"Enqueued task {task_id} ({args.func_name}){note}")


def cmd_status(args, broker: Broker):
    task = broker.get_task(args.task_id)
    if task is None:
        print(f"No task found with id {args.task_id}", file=sys.stderr)
        sys.exit(1)
    print(json.dumps(task.to_dict(), indent=2, default=str))


def cmd_list(args, broker: Broker):
    status = TaskStatus(args.status) if args.status else None
    tasks = broker.list_tasks(status=status, limit=args.limit)
    for t in tasks:
        deps = f"  deps={len(t.depends_on)}" if t.depends_on else ""
        print(f"{t.id}  {t.status.value:8s}  attempts={t.attempts}/{t.max_retries}  {t.func_name}{t.args}{deps}")
    if not tasks:
        print("(no tasks)")


def cmd_stats(args, broker: Broker):
    print(json.dumps(broker.stats(), indent=2))


def cmd_requeue(args, broker: Broker):
    broker.requeue_dead_task(args.task_id)
    print(f"Requeued task {args.task_id} for a fresh set of attempts")


def cmd_schedule_add(args, broker: Broker):
    schedule = Schedule(
        func_name=args.func_name,
        interval_seconds=args.interval,
        args=tuple(json.loads(args.args)),
        kwargs=json.loads(args.kwargs),
        priority=args.priority,
        max_retries=args.max_retries,
    )
    schedule_id = broker.add_schedule(schedule)
    print(f"Added schedule {schedule_id}: runs {args.func_name} every {args.interval}s (starting now)")


def cmd_schedule_list(args, broker: Broker):
    schedules = broker.list_schedules()
    for s in schedules:
        state = "enabled " if s.enabled else "disabled"
        print(f"{s.id}  {state}  every {s.interval_seconds:>5}s  next={s.next_run_at}  {s.func_name}{s.args}")
    if not schedules:
        print("(no schedules)")


def cmd_schedule_enable(args, broker: Broker):
    broker.set_schedule_enabled(args.schedule_id, True)
    print(f"Enabled schedule {args.schedule_id}")


def cmd_schedule_disable(args, broker: Broker):
    broker.set_schedule_enabled(args.schedule_id, False)
    print(f"Disabled schedule {args.schedule_id}")


def build_parser():
    parser = argparse.ArgumentParser(description="Distributed Task Queue CLI")
    sub = parser.add_subparsers(dest="command", required=True)

    p_submit = sub.add_parser("submit", help="Enqueue a new task")
    p_submit.add_argument("func_name", help="Name the task was registered under, e.g. 'add'")
    p_submit.add_argument("--args", default="[]", help="JSON array of positional args")
    p_submit.add_argument("--kwargs", default="{}", help="JSON object of keyword args")
    p_submit.add_argument("--priority", type=int, default=0, help="Higher runs sooner")
    p_submit.add_argument("--max-retries", type=int, default=3)
    p_submit.add_argument(
        "--depends-on", default=None,
        help="Comma-separated task ids that must SUCCEED before this task can run",
    )
    p_submit.set_defaults(func=cmd_submit)

    p_status = sub.add_parser("status", help="Show a single task's full record")
    p_status.add_argument("task_id")
    p_status.set_defaults(func=cmd_status)

    p_list = sub.add_parser("list", help="List tasks, optionally filtered by status")
    p_list.add_argument("--status", choices=[s.value for s in TaskStatus], default=None)
    p_list.add_argument("--limit", type=int, default=50)
    p_list.set_defaults(func=cmd_list)

    p_stats = sub.add_parser("stats", help="Show task counts by status")
    p_stats.set_defaults(func=cmd_stats)

    p_requeue = sub.add_parser("requeue", help="Give a dead-lettered task a fresh set of attempts")
    p_requeue.add_argument("task_id")
    p_requeue.set_defaults(func=cmd_requeue)

    p_sched_add = sub.add_parser("schedule-add", help="Create a recurring schedule (runs immediately, then every --interval seconds)")
    p_sched_add.add_argument("func_name")
    p_sched_add.add_argument("--interval", type=int, required=True, help="Seconds between runs")
    p_sched_add.add_argument("--args", default="[]")
    p_sched_add.add_argument("--kwargs", default="{}")
    p_sched_add.add_argument("--priority", type=int, default=0)
    p_sched_add.add_argument("--max-retries", type=int, default=3)
    p_sched_add.set_defaults(func=cmd_schedule_add)

    p_sched_list = sub.add_parser("schedule-list", help="List all recurring schedules")
    p_sched_list.set_defaults(func=cmd_schedule_list)

    p_sched_enable = sub.add_parser("schedule-enable", help="Re-enable a paused schedule")
    p_sched_enable.add_argument("schedule_id")
    p_sched_enable.set_defaults(func=cmd_schedule_enable)

    p_sched_disable = sub.add_parser("schedule-disable", help="Pause a schedule without deleting it")
    p_sched_disable.add_argument("schedule_id")
    p_sched_disable.set_defaults(func=cmd_schedule_disable)

    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()
    broker = Broker()
    args.func(args, broker)


if __name__ == "__main__":
    main()