mirror of
https://github.com/infinition/Bjorn.git
synced 2026-03-09 22:22:00 +00:00
BREAKING CHANGE: Complete refactor of architecture to prepare BJORN V2 release, APIs, assets, and UI, webapp, logics, attacks, a lot of new features...
This commit is contained in:
Fabien POLLY
410
db_utils/queue.py
Normal file
410
db_utils/queue.py
Normal file
@@ -0,0 +1,410 @@
|
||||
# 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(" ", "_")
|
||||
Reference in New Issue
Block a user