Bjorn/db_utils/queue.py

# 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(" ", "_")