# db_utils/queue.py # Action queue management operations import json import sqlite3 from typing import Any, Dict, Iterable, List, Optional import logging from logger import Logger logger = Logger(name="db_utils.queue", level=logging.DEBUG) class QueueOps: """Action queue scheduling and execution tracking operations""" def __init__(self, base): self.base = base def create_tables(self): """Create action queue table and indexes""" self.base.execute(""" CREATE TABLE IF NOT EXISTS action_queue ( id INTEGER PRIMARY KEY AUTOINCREMENT, action_name TEXT NOT NULL, mac_address TEXT NOT NULL, ip TEXT NOT NULL, port INTEGER, hostname TEXT, service TEXT, priority INTEGER DEFAULT 50, status TEXT DEFAULT 'pending', retry_count INTEGER DEFAULT 0, max_retries INTEGER DEFAULT 3, created_at TEXT DEFAULT CURRENT_TIMESTAMP, scheduled_for TEXT, started_at TEXT, completed_at TEXT, expires_at TEXT, trigger_source TEXT, dependencies TEXT, conditions TEXT, result_summary TEXT, error_message TEXT, tags TEXT, metadata TEXT, FOREIGN KEY (mac_address) REFERENCES hosts(mac_address) ); """) # Optimized indexes for queue operations self.base.execute("CREATE INDEX IF NOT EXISTS idx_queue_pending ON action_queue(status) WHERE status='pending';") self.base.execute("CREATE INDEX IF NOT EXISTS idx_queue_scheduled ON action_queue(scheduled_for) WHERE status='scheduled';") self.base.execute("CREATE INDEX IF NOT EXISTS idx_queue_mac_action ON action_queue(mac_address, action_name);") self.base.execute("CREATE INDEX IF NOT EXISTS idx_queue_key_status ON action_queue(action_name, mac_address, port, status);") self.base.execute("CREATE INDEX IF NOT EXISTS idx_queue_key_time ON action_queue(action_name, mac_address, port, completed_at);") # Unique constraint for a single upcoming schedule per action/target self.base.execute(""" CREATE UNIQUE INDEX IF NOT EXISTS uq_next_scheduled ON action_queue(action_name, COALESCE(mac_address,''), COALESCE(service,''), COALESCE(port,-1)) WHERE status='scheduled'; """) logger.debug("Action queue table created/verified") # ========================================================================= # QUEUE RETRIEVAL OPERATIONS # ========================================================================= def get_next_queued_action(self) -> Optional[Dict[str, Any]]: """ Fetch the next action to execute from the queue. Priority is dynamically boosted: +1 per 5 minutes since creation, capped at +100. """ rows = self.base.query(""" SELECT *, MIN(100, priority + CAST((strftime('%s','now') - strftime('%s',created_at))/300 AS INTEGER)) AS priority_effective FROM action_queue WHERE status = 'pending' AND (scheduled_for IS NULL OR scheduled_for <= datetime('now')) ORDER BY priority_effective DESC, COALESCE(scheduled_for, created_at) ASC LIMIT 1 """) return rows[0] if rows else None def list_action_queue(self, statuses: Optional[Iterable[str]] = None) -> List[Dict[str, Any]]: """List queue entries with a computed `priority_effective` column for pending items""" order_sql = """ CASE status WHEN 'running' THEN 1 WHEN 'pending' THEN 2 WHEN 'scheduled' THEN 3 WHEN 'failed' THEN 4 WHEN 'success' THEN 5 WHEN 'expired' THEN 6 WHEN 'cancelled' THEN 7 ELSE 99 END ASC, priority_effective DESC, COALESCE(scheduled_for, created_at) ASC """ select_sql = """ SELECT *, MIN(100, priority + CAST((strftime('%s','now') - strftime('%s',created_at))/300 AS INTEGER)) AS priority_effective FROM action_queue """ if statuses: in_clause = ",".join("?" for _ in statuses) return self.base.query(f""" {select_sql} WHERE status IN ({in_clause}) ORDER BY {order_sql} """, tuple(statuses)) return self.base.query(f""" {select_sql} ORDER BY {order_sql} """) def get_upcoming_actions_summary(self) -> List[Dict[str, Any]]: """Summary: next run per action_name from the schedule""" return self.base.query(""" SELECT action_name, MIN(scheduled_for) AS next_run_at FROM action_queue WHERE status='scheduled' AND scheduled_for IS NOT NULL GROUP BY action_name ORDER BY next_run_at ASC """) # ========================================================================= # QUEUE UPDATE OPERATIONS # ========================================================================= def update_queue_status(self, queue_id: int, status: str, error_msg: str = None, result: str = None): """Update queue entry status with retry management on failure/expiry""" self.base.invalidate_stats_cache() if status == 'running': self.base.execute( "UPDATE action_queue SET status=?, started_at=CURRENT_TIMESTAMP WHERE id=?", (status, queue_id) ) elif status in ('failed', 'expired'): self.base.execute(""" UPDATE action_queue SET status=?, completed_at=CURRENT_TIMESTAMP, error_message=?, result_summary=COALESCE(?, result_summary), retry_count = MIN(retry_count + 1, max_retries) WHERE id=? """, (status, error_msg, result, queue_id)) elif status in ('success', 'cancelled'): self.base.execute(""" UPDATE action_queue SET status=?, completed_at=CURRENT_TIMESTAMP, error_message=?, result_summary=COALESCE(?, result_summary) WHERE id=? """, (status, error_msg, result, queue_id)) # When execution succeeds, supersede old failed/expired attempts if status == 'success': row = self.base.query_one(""" SELECT action_name, mac_address, port, COALESCE(completed_at, started_at, created_at) AS ts FROM action_queue WHERE id=? LIMIT 1 """, (queue_id,)) if row: try: self.supersede_old_attempts(row['action_name'], row['mac_address'], row['port'], row['ts']) except Exception: pass def promote_due_scheduled_to_pending(self) -> int: """Promote scheduled actions that are due (returns number of rows affected)""" self.base.invalidate_stats_cache() return self.base.execute(""" UPDATE action_queue SET status='pending' WHERE status='scheduled' AND scheduled_for <= CURRENT_TIMESTAMP """) # ========================================================================= # QUEUE INSERTION OPERATIONS # ========================================================================= def ensure_scheduled_occurrence( self, action_name: str, next_run_at: str, mac: Optional[str] = "", ip: Optional[str] = "", *, port: Optional[int] = None, hostname: Optional[str] = None, service: Optional[str] = None, priority: int = 40, trigger: str = "scheduler", tags: Optional[Iterable[str]] = None, metadata: Optional[Dict[str, Any]] = None, max_retries: Optional[int] = None, ) -> bool: """ Ensure a single upcoming 'scheduled' row exists for the given action/target. Returns True if inserted, False if already present (enforced by unique partial index). """ js_tags = json.dumps(list(tags)) if tags is not None and not isinstance(tags, str) else (tags if isinstance(tags, str) else None) js_meta = json.dumps(metadata, ensure_ascii=False) if metadata else None try: self.base.execute(""" INSERT INTO action_queue( action_name, mac_address, ip, port, hostname, service, priority, status, scheduled_for, trigger_source, tags, metadata, max_retries ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?) """, ( action_name, mac or "", ip or "", port, hostname, service, int(priority), "scheduled", next_run_at, trigger, js_tags, js_meta, max_retries )) self.base.invalidate_stats_cache() return True except sqlite3.IntegrityError: return False def queue_action(self, action_name: str, mac: str, ip: str, port: int = None, priority: int = 50, trigger: str = None, metadata: Dict = None) -> None: """Quick enqueue of a 'pending' action""" meta_json = json.dumps(metadata, ensure_ascii=False) if metadata else None self.base.execute(""" INSERT INTO action_queue (action_name, mac_address, ip, port, priority, trigger_source, metadata) VALUES (?,?,?,?,?,?,?) """, (action_name, mac, ip, port, priority, trigger, meta_json)) def queue_action_at( self, action_name: str, mac: Optional[str] = "", ip: Optional[str] = "", *, port: Optional[int] = None, hostname: Optional[str] = None, service: Optional[str] = None, priority: int = 50, status: str = "pending", scheduled_for: Optional[str] = None, trigger: Optional[str] = "scheduler", tags: Optional[Iterable[str]] = None, metadata: Optional[Dict[str, Any]] = None, max_retries: Optional[int] = None, ) -> None: """Generic enqueue that can publish 'pending' or 'scheduled' items with a date""" js_tags = json.dumps(list(tags)) if tags is not None and not isinstance(tags, str) else (tags if isinstance(tags, str) else None) js_meta = json.dumps(metadata, ensure_ascii=False) if metadata else None self.base.execute(""" INSERT INTO action_queue( action_name, mac_address, ip, port, hostname, service, priority, status, scheduled_for, trigger_source, tags, metadata, max_retries ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?) """, ( action_name, mac or "", ip or "", port, hostname, service, int(priority), status, scheduled_for, trigger, js_tags, js_meta, max_retries )) # ========================================================================= # HISTORY AND STATUS OPERATIONS # ========================================================================= def supersede_old_attempts(self, action_name: str, mac_address: str, port: Optional[int] = None, ref_ts: Optional[str] = None) -> int: """ Mark as 'superseded' all old attempts (failed|expired) for the triplet (action, mac, port) earlier than or equal to ref_ts (if provided). Returns affected row count. """ params: List[Any] = [action_name, mac_address, port] time_clause = "" if ref_ts: time_clause = " AND datetime(COALESCE(completed_at, started_at, created_at)) <= datetime(?)" params.append(ref_ts) return self.base.execute(f""" UPDATE action_queue SET status='superseded', error_message = COALESCE(error_message, 'superseded by newer success'), completed_at = COALESCE(completed_at, CURRENT_TIMESTAMP) WHERE action_name = ? AND mac_address = ? AND COALESCE(port,0) = COALESCE(?,0) AND status IN ('failed','expired') {time_clause} """, tuple(params)) def list_attempt_history(self, action_name: str, mac_address: str, port: Optional[int] = None, limit: int = 20) -> List[Dict[str, Any]]: """ Return history of attempts for (action, mac, port), most recent first. """ return self.base.query(""" SELECT action_name, mac_address, port, status, retry_count, max_retries, COALESCE(completed_at, started_at, scheduled_for, created_at) AS ts FROM action_queue WHERE action_name=? AND mac_address=? AND COALESCE(port,0)=COALESCE(?,0) ORDER BY datetime(ts) DESC LIMIT ? """, (action_name, mac_address, port, int(limit))) def get_action_status_from_queue( self, action_name: str, mac_address: Optional[str] = None ) -> Optional[Dict[str, Any]]: """ Return the latest status row for an action (optionally filtered by MAC). """ if mac_address: rows = self.base.query(""" SELECT status, created_at, started_at, completed_at, error_message, result_summary, retry_count, max_retries, mac_address, port, hostname, service, priority FROM action_queue WHERE mac_address=? AND action_name=? ORDER BY datetime(COALESCE(completed_at, started_at, scheduled_for, created_at)) DESC LIMIT 1 """, (mac_address, action_name)) else: rows = self.base.query(""" SELECT status, created_at, started_at, completed_at, error_message, result_summary, retry_count, max_retries, mac_address, port, hostname, service, priority FROM action_queue WHERE action_name=? ORDER BY datetime(COALESCE(completed_at, started_at, scheduled_for, created_at)) DESC LIMIT 1 """, (action_name,)) return rows[0] if rows else None def get_last_action_status_from_queue(self, mac_address: str, action_name: str) -> Optional[Dict[str, str]]: """ Return {'status': 'success|failed|running|pending', 'raw': 'status_YYYYMMDD_HHMMSS'} based only on action_queue. """ rows = self.base.query( """ SELECT status, COALESCE(completed_at, started_at, scheduled_for, created_at) AS ts FROM action_queue WHERE mac_address=? AND action_name=? ORDER BY datetime(COALESCE(completed_at, started_at, scheduled_for, created_at)) DESC LIMIT 1 """, (mac_address, action_name) ) if not rows: return None status = rows[0]["status"] ts = self._format_ts_for_raw(rows[0]["ts"]) return {"status": status, "raw": f"{status}_{ts}"} def get_last_action_statuses_for_mac(self, mac_address: str) -> Dict[str, Dict[str, str]]: """ Map action_name -> {'status':..., 'raw':...} from the latest queue rows for a MAC. """ rows = self.base.query( """ SELECT action_name, status, COALESCE(completed_at, started_at, scheduled_for, created_at) AS ts FROM ( SELECT action_name, status, completed_at, started_at, scheduled_for, created_at, ROW_NUMBER() OVER ( PARTITION BY action_name ORDER BY datetime(COALESCE(completed_at, started_at, scheduled_for, created_at)) DESC ) AS rn FROM action_queue WHERE mac_address=? ) WHERE rn=1 """, (mac_address,) ) out: Dict[str, Dict[str, str]] = {} for r in rows: ts = self._format_ts_for_raw(r["ts"]) st = r["status"] out[r["action_name"]] = {"status": st, "raw": f"{st}_{ts}"} return out # ========================================================================= # HELPER METHODS # ========================================================================= def _format_ts_for_raw(self, ts_db: Optional[str]) -> str: """ Convert SQLite 'YYYY-MM-DD HH:MM:SS' to 'YYYYMMDD_HHMMSS'. Fallback to current UTC when no timestamp is available. """ from datetime import datetime as _dt ts = (ts_db or "").strip() if not ts: return _dt.utcnow().strftime("%Y%m%d_%H%M%S") return ts.replace("-", "").replace(":", "").replace(" ", "_")