""" Task Queue Service for bulk VM/Container operations Processes tasks from the queue in background """ from datetime import datetime, timezone from typing import List, Dict, Any, Optional from loguru import logger from sqlalchemy.orm import Session def utcnow() -> datetime: """Get current UTC time as timezone-aware datetime""" return datetime.now(timezone.utc) try: from backend.app.db import SessionLocal from backend.app.models import TaskQueue, ProxmoxServer, User from backend.app.proxmox import ProxmoxClient from backend.app.websocket_manager import broadcast_task_update except ImportError: from app.db import SessionLocal from app.models import TaskQueue, ProxmoxServer, User from app.proxmox import ProxmoxClient from app.websocket_manager import broadcast_task_update def _ws_broadcast(task: TaskQueue, event: str = "task_update") -> None: """Fire-and-forget WS broadcast for a TaskQueue update.""" try: data = task.to_dict() data["kind"] = "bulk" broadcast_task_update(task.user_id, event, data) except Exception as e: logger.debug(f"[TASK QUEUE] WS broadcast skipped: {e}") class TaskQueueService: """Service for managing and processing bulk operation tasks""" # Valid task types TASK_TYPES = ['bulk_start', 'bulk_stop', 'bulk_restart', 'bulk_delete', 'bulk_shutdown'] @staticmethod def create_task( db: Session, task_type: str, user_id: int, items: List[Dict[str, Any]] ) -> TaskQueue: """ Create a new bulk operation task Args: db: Database session task_type: Type of operation (bulk_start, bulk_stop, etc.) user_id: ID of user who initiated items: List of VMs/containers to process Each item: {server_id, vmid, vm_type, name, node} """ if task_type not in TaskQueueService.TASK_TYPES: raise ValueError(f"Invalid task type: {task_type}") if not items: raise ValueError("No items provided for bulk operation") task = TaskQueue( task_type=task_type, status='pending', user_id=user_id, total_items=len(items), completed_items=0, failed_items=0, task_data=items, results=[] ) db.add(task) db.commit() db.refresh(task) logger.info(f"[TASK QUEUE] Created task #{task.id}: {task_type} with {len(items)} items") return task @staticmethod def get_pending_tasks(db: Session) -> List[TaskQueue]: """Get all pending tasks ordered by creation time""" return db.query(TaskQueue).filter( TaskQueue.status == 'pending' ).order_by(TaskQueue.created_at.asc()).all() @staticmethod def get_running_tasks(db: Session) -> List[TaskQueue]: """Get all currently running tasks""" return db.query(TaskQueue).filter( TaskQueue.status == 'running' ).all() @staticmethod def get_user_tasks(db: Session, user_id: int, limit: int = 20) -> List[TaskQueue]: """Get recent tasks for a user""" return db.query(TaskQueue).filter( TaskQueue.user_id == user_id ).order_by(TaskQueue.created_at.desc()).limit(limit).all() @staticmethod def get_active_tasks(db: Session) -> List[TaskQueue]: """Get all active (pending or running) tasks""" return db.query(TaskQueue).filter( TaskQueue.status.in_(['pending', 'running']) ).order_by(TaskQueue.created_at.asc()).all() @staticmethod def cancel_task(db: Session, task_id: int, user_id: int) -> bool: """Cancel a pending task""" task = db.query(TaskQueue).filter( TaskQueue.id == task_id, TaskQueue.user_id == user_id, TaskQueue.status == 'pending' ).first() if not task: return False task.status = 'cancelled' task.completed_at = utcnow() db.commit() _ws_broadcast(task, "task_update") logger.info(f"[TASK QUEUE] Task #{task_id} cancelled by user {user_id}") return True class TaskQueueProcessor: """Processor that executes tasks from the queue""" def __init__(self): self._proxmox_clients: Dict[int, ProxmoxClient] = {} def _get_proxmox_client(self, db: Session, server_id: int) -> Optional[ProxmoxClient]: """Get or create Proxmox client for server""" if server_id in self._proxmox_clients: client = self._proxmox_clients[server_id] if client.is_connected(): return client server = db.query(ProxmoxServer).filter(ProxmoxServer.id == server_id).first() if not server: return None # Require valid credentials before building (from_server would otherwise # construct a client with no usable auth). has_password = server.use_password and server.password has_token = server.api_token_name and server.api_token_value if not (has_password or has_token): return None try: client = ProxmoxClient.from_server(server) self._proxmox_clients[server_id] = client return client except Exception as e: logger.error(f"[TASK QUEUE] Failed to connect to server {server_id}: {e}") return None def _execute_vm_action( self, client: ProxmoxClient, action: str, vmid: int, vm_type: str, node: str ) -> tuple[bool, str]: """Execute action on VM/container""" try: if vm_type == 'lxc': if action == 'start': client.start_container(node, vmid) elif action == 'stop': client.stop_container(node, vmid) elif action == 'restart': client.restart_container(node, vmid) elif action == 'shutdown': # Shutdown = graceful stop client.stop_container(node, vmid) elif action == 'delete': client.delete_container(node, vmid) else: return False, f"Unknown action: {action}" else: # qemu if action == 'start': client.start_vm(node, vmid) elif action == 'stop': client.stop_vm(node, vmid) elif action == 'restart': client.restart_vm(node, vmid) elif action == 'shutdown': # Shutdown = graceful stop client.stop_vm(node, vmid) elif action == 'delete': client.delete_vm(node, vmid) else: return False, f"Unknown action: {action}" return True, "OK" except Exception as e: return False, str(e) def process_task(self, task: TaskQueue) -> None: """Process a single task""" db = SessionLocal() try: # Mark as running task = db.query(TaskQueue).filter(TaskQueue.id == task.id).first() if not task or task.status != 'pending': return task.status = 'running' task.started_at = utcnow() db.commit() _ws_broadcast(task, "task_update") logger.info(f"[TASK QUEUE] Starting task #{task.id}: {task.task_type}") # Determine action from task type action_map = { 'bulk_start': 'start', 'bulk_stop': 'stop', 'bulk_restart': 'restart', 'bulk_shutdown': 'shutdown', 'bulk_delete': 'delete', } action = action_map.get(task.task_type) if not action: task.status = 'failed' task.error_message = f"Unknown task type: {task.task_type}" task.completed_at = utcnow() db.commit() return results = [] items = task.task_data or [] for item in items: server_id = item.get('server_id') vmid = item.get('vmid') vm_type = item.get('vm_type', 'qemu') node = item.get('node', '') name = item.get('name', f'{vm_type.upper()}-{vmid}') client = self._get_proxmox_client(db, server_id) if not client: results.append({ 'server_id': server_id, 'vmid': vmid, 'name': name, 'success': False, 'message': 'Failed to connect to Proxmox server' }) task.failed_items += 1 else: success, message = self._execute_vm_action(client, action, vmid, vm_type, node) results.append({ 'server_id': server_id, 'vmid': vmid, 'name': name, 'success': success, 'message': message }) if success: task.completed_items += 1 else: task.failed_items += 1 # Update progress task.results = results db.commit() _ws_broadcast(task, "task_update") # Mark as completed task.status = 'completed' task.completed_at = utcnow() db.commit() _ws_broadcast(task, "task_update") logger.info(f"[TASK QUEUE] Completed task #{task.id}: {task.completed_items} success, {task.failed_items} failed") except Exception as e: logger.error(f"[TASK QUEUE] Task #{task.id} failed with error: {e}") try: task.status = 'failed' task.error_message = str(e) task.completed_at = utcnow() db.commit() _ws_broadcast(task, "task_update") except Exception: pass finally: db.close() def process_pending_tasks(self) -> int: """Process all pending tasks. Returns number of tasks processed.""" db = SessionLocal() try: # Get pending tasks pending_tasks = TaskQueueService.get_pending_tasks(db) if not pending_tasks: return 0 # Check if there's already a running task running_tasks = TaskQueueService.get_running_tasks(db) if running_tasks: logger.debug(f"[TASK QUEUE] {len(running_tasks)} task(s) already running, skipping") return 0 # Process first pending task task = pending_tasks[0] db.close() # Close before processing (will create new session) self.process_task(task) return 1 except Exception as e: logger.error(f"[TASK QUEUE] Error processing pending tasks: {e}") return 0 finally: try: db.close() except Exception: pass # Global processor instance task_processor = TaskQueueProcessor() def process_task_queue(): """Function to be called by scheduler to process queue""" return task_processor.process_pending_tasks()