queuectl.py

import os
import click
import multiprocessing
import signal
import sys
from datetime import datetime, timedelta
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.live import Live
from rich.status import Status
from rich import print as rprint
from queue_system import storage, config, worker
from tabulate import tabulate
import webbrowser
import threading
# Initialize rich console
console = Console()

# Global list to track workers


worker_process= []


def handle_signal(sig, frame):
    """Gracefully shut down worker processes on sigint (ctrl+c)"""
    with console.status("[yellow]Shutting down workers...[/yellow]"):
        for p in worker_process:
            p.terminate()
        for p in worker_process:
            p.join()  # wait for process to exit
    console.print("[green]All workers stopped[/green]")
    sys.exit(0)

@click.group()
def cli():
    """
    QueueCTL: A CLI based background job queue system.
    """
    pass

@cli.command()
def init():
    """initialzes the database and config"""
    storage.init_db()

#enqueu commands

@cli.command()
@click.argument('command_str', nargs=-1)
@click.option('--priority', default=0, type=int, help="Job priority (higher runs sooner)")
@click.option('--timeout', default=30, type=int, help="Job timeout, automatically transfers pending task to dlq")
@click.option('--delay', default=0, type=int, help="Delay execution by N seconds.")
def enqueue(command_str, delay, priority, timeout):
    """
    Add a new job to the queue

    Example: queuectl enqueue sleep 5
    """
    if not command_str:
        console.print("[red]Error:[/red] No command provided.")
        return
    
    run_at_time = datetime.now() + timedelta(seconds=delay)


    try:
        command = " ".join(command_str)
        with console.status("[cyan]Adding job to queue...[/cyan]"):
            job_id = storage.create_job(command, run_at_time=run_at_time,priority=priority,timeout=timeout)
        
        if delay > 0:
            console.print(f"[green]Enqueued job[/green] [cyan]{job_id}[/cyan] [yellow](will run in {delay}s)[/yellow]")
        else:
            console.print(f"[green]Enqueued job[/green] [cyan]{job_id}[/cyan]: {command}")
    except Exception as e:
        console.print(f"[red]Error enqueuing job:[/red] {str(e)}")


#worker commands

@cli.command()
@click.option('--count', default=3, type=int, help='Number of workers to start.')
def start_workers(count):
    """starts one or more worker processes"""
    if count < 1:
        console.print("[red]Error:[/red] Must start at least one worker")
        return

    mx_num_workers = os.cpu_count()
    if count > mx_num_workers:
        console.print(f"[red]Error:[/red] Maximum number of workers is {mx_num_workers}")
        return

    # Set up the signal handler for graceful shutdown
    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)

    console.print(Panel(
        f"Starting [cyan]{count}[/cyan] workers\nPress [bold red]Ctrl+C[/bold red] to stop",
        border_style="green"
    ))

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        console=console
    ) as progress:
        task = progress.add_task("Starting workers...", total=count)
        for _ in range(count):
            p = multiprocessing.Process(target=worker.run_worker_process)
            p.start()
            worker_process.append(p)
            progress.advance(task)


    #wait for all process to finish
    # the signall handler will interupt this join

    for p in worker_process:
        p.join()



#status& list commands
@cli.command()
def status():
    """ shows a summary of all job states"""
    summary = storage.get_job_summary()
    if not summary:
        console.print("[yellow]No jobs found[/yellow]")
        return
    
    table = Table(show_header=True, header_style="bold magenta")
    table.add_column("State", style="cyan")
    table.add_column("Count", justify="right", style="blue", width= 50)
    
    for row in summary:
        table.add_row(row['state'], str(row['count']))
    
    console.print(Panel(table, title="Job Queue Status", border_style="blue"))


@cli.command()
@click.option('--state', default='pending', help='job state (pending, processing, completed, dead)')
def list(state):
    """Lists Jobs by state"""
    jobs = storage.get_jobs_by_state(state)
    if not jobs:
        console.print(f"[yellow]No jobs found in '[bold]{state}[/bold]' state.[/yellow]")
        return
    
    table = Table(show_header=True, header_style="bold magenta", width=95)
    table.add_column("ID", style="cyan", width=40)
    table.add_column("Command", style="green", width=32)
    table.add_column("Attempts", justify="right", style="yellow", width=8)
    table.add_column("Updated At", style="blue", width=16)
    
    for job in jobs:
        command = job['command']
        if len(command) > 50:
            command = command[:47] + '...'
        
        table.add_row(
            job['id'],
            command,
            str(job['attempts']),
            job['updated_at']
        )
    
    console.print(Panel(table, title=f"Jobs in '{state}' state", border_style="blue"))

#  --- DLQ Commands---


@cli.group()
def dlq():
    """Manage the Dead Letter queue (DLQ)"""

    pass

@dlq.command(name="list")
def dlq_list():
    """List all the jobs in DLQ"""
    jobs = storage.get_jobs_by_state('dead')
    if not jobs:
        console.print("[yellow]DLQ is empty![/yellow]")
        return
    
    table = Table(show_header=True, header_style="bold red", width=95)
    table.add_column("ID", style="cyan", width=38)  
    table.add_column("Command", style="green", width=30)
    table.add_column("Attempts", justify="right", style="yellow", width=8)
    table.add_column("Updated At", style="blue", width=16)
    
    for job in jobs:
        command = job['command']
        if len(command) > 50:
            command = command[:47] + '...'
        
        table.add_row(
            job['id'],
            command,
            str(job['attempts']),
            job['updated_at']
        )
    
    console.print(Panel(table, title="Dead Letter Queue", border_style="red"))


@dlq.command(name="retry")
@click.argument('job_id')
def dlq_retry(job_id):
    """Retry a specific job from DLQ"""
    if storage.retry_dlq_job(job_id):
        print(f"job {job_id} moved from DLQ to 'pending' queue.")

    else:
        print(f"Error: job {job_id} not found in DLQ")




@cli.command(name="config")
@click.argument('key')
@click.argument('value')
def set_config(key, value):
    """set a config value (ex , max_retries, backoff_base)."""
    from queue_system import config as config_module
    config_module.set_config(key, value)


@cli.command()
@click.argument('job_id')
def inspect(job_id):
    """Get detailed info about a specific job"""
    job = storage.get_jobs_by_id(job_id)

    if not job:
        console.print(f"[red]Error:[/red] Job [bold]{job_id}[/bold] not found.")
        return

    # Header
    console.print(Panel(f"[cyan]{job.get('command', '')}[/cyan]", title=f"Job {job.get('id')}", border_style="magenta"))

    # Main info table
    main_table = Table.grid(expand=False)
    main_table.add_column(justify="right", style="bold")
    main_table.add_column()

    main_table.add_row("State", str(job.get('state', 'N/A')))
    main_table.add_row("Priority", str(job.get('priority', 0)))
    main_table.add_row("Attempts", f"{job.get('attempts', 0)} / {job.get('max_retries', 0)}")
    main_table.add_row("Timeout", f"{job.get('timeout', 0)}s")

    console.print(Panel(main_table, title="Main Info", border_style="blue"))

    # Time info
    time_table = Table.grid(expand=False)
    time_table.add_column(justify="right", style="bold")
    time_table.add_column()

    created = job.get('created_at') or "N/A"
    next_run = job.get('next_run_at') or "N/A"
    started = job.get('started_at') or "N/A"
    finished = job.get('finished_at') or "N/A"

    time_table.add_row("Created", created)
    time_table.add_row("Next Run", next_run)
    time_table.add_row("Started", started)
    time_table.add_row("Finished", finished)

    # Duration calculation
    if job.get('started_at') and job.get('finished_at'):
        try:
            start = datetime.fromisoformat(job['started_at'])
            finish = datetime.fromisoformat(job['finished_at'])
            duration = finish - start
            time_table.add_row("Duration", f"{duration.total_seconds():.2f}s")
        except Exception:
            pass

    console.print(Panel(time_table, title="Timestamps", border_style="green"))

    # Output logs
    if job.get('stdout'):
        console.print(Panel(job.get('stdout'), title="STDOUT", border_style="green"))

    if job.get('stderr'):
        console.print(Panel(job.get('stderr'), title="STDERR", border_style="red"))


@cli.command()
@click.option('--port', default=5000, type=int, help='Port to run the dashboard on.')
def dashboard(port):
    """Starts a minimal web dashboard to monitor the queue."""
    print(f"Starting web dashboard on http://127.0.0.1:{port}/")
    
    # Function to run the server
    def run_server():
        try:
            from dashboard import app
            app.run(port=port, debug=False, use_reloader=False)
        except ImportError:
            print("\nError: Could not import dashboard.")
            print("Please run: pip install Flask htmx-flask")
        except Exception as e:
            if "Address already in use" in str(e):
                print(f"Error: Port {port} is already in use.")
            else:
                print(f"\nDashboard server error: {e}")
    threading.Timer(1, lambda: webbrowser.open(f"http://127.0.0.1:{port}/")).start()
    
    run_server()
 


if __name__ == '__main__':
    cli()