diff --git a/db.py b/db.py index 4434c82..c7a67c7 100644 --- a/db.py +++ b/db.py @@ -5,6 +5,7 @@ """ import sqlite3 import json +from contextlib import contextmanager from datetime import datetime, timezone from pathlib import Path from config import DB_PATH, TASK_STATES, PRIORITIES @@ -18,10 +19,25 @@ def get_conn(): conn.row_factory = sqlite3.Row return conn + +@contextmanager +def get_db(): + """Yield a connection that always commits or rolls back and closes.""" + conn = get_conn() + try: + yield conn + conn.commit() + except Exception: + conn.rollback() + raise + finally: + conn.close() + + def init_db(): """Create all tables if they don't exist.""" - conn = get_conn() - conn.executescript(""" + with get_db() as conn: + conn.executescript(""" -- Sprints: 2-week planning cycles CREATE TABLE IF NOT EXISTS sprints ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -112,204 +128,187 @@ def init_db(): CREATE INDEX IF NOT EXISTS idx_log_task ON agent_log(task_id); CREATE INDEX IF NOT EXISTS idx_benchmarks_metric ON benchmarks(metric); CREATE INDEX IF NOT EXISTS idx_deployments_site ON deployments(site_id); - """) - conn.commit() - conn.close() + """) print(f"[DB] Initialized at {DB_PATH}") # ─── Task Operations ──────────────────────────────────────────── def create_task(sprint_id, title, description, agent_id, priority="P2_MEDIUM", deliverables=None, depends_on=None, estimated_hours=None): - conn = get_conn() - conn.execute(""" - INSERT INTO tasks (sprint_id, title, description, agent_id, priority, - deliverables, depends_on, estimated_hours) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - """, (sprint_id, title, description, agent_id, priority, - json.dumps(deliverables or []), - json.dumps(depends_on or []), - estimated_hours)) - conn.commit() - task_id = conn.execute("SELECT last_insert_rowid()").fetchone()[0] - conn.close() + with get_db() as conn: + conn.execute(""" + INSERT INTO tasks (sprint_id, title, description, agent_id, priority, + deliverables, depends_on, estimated_hours) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, (sprint_id, title, description, agent_id, priority, + json.dumps(deliverables or []), + json.dumps(depends_on or []), + estimated_hours)) + task_id = conn.execute("SELECT last_insert_rowid()").fetchone()[0] return task_id def update_task_state(task_id, new_state): - conn = get_conn() - updates = {"state": new_state} - if new_state == "IN_PROGRESS": - updates["started_at"] = datetime.now(timezone.utc).isoformat() - elif new_state == "DONE": - updates["completed_at"] = datetime.now(timezone.utc).isoformat() - - set_clause = ", ".join(f"{k} = ?" for k in updates) - conn.execute(f"UPDATE tasks SET {set_clause} WHERE id = ?", - list(updates.values()) + [task_id]) - conn.commit() - conn.close() + with get_db() as conn: + if new_state == "IN_PROGRESS": + conn.execute( + "UPDATE tasks SET state = ?, started_at = ? WHERE id = ?", + (new_state, datetime.now(timezone.utc).isoformat(), task_id), + ) + elif new_state == "DONE": + conn.execute( + "UPDATE tasks SET state = ?, completed_at = ? WHERE id = ?", + (new_state, datetime.now(timezone.utc).isoformat(), task_id), + ) + else: + conn.execute( + "UPDATE tasks SET state = ? WHERE id = ?", + (new_state, task_id), + ) def get_tasks_for_agent(agent_id, state=None): - conn = get_conn() - if state: - rows = conn.execute( - "SELECT * FROM tasks WHERE agent_id = ? AND state = ? ORDER BY priority", - (agent_id, state) - ).fetchall() - else: - rows = conn.execute( - "SELECT * FROM tasks WHERE agent_id = ? ORDER BY priority, state", - (agent_id,) - ).fetchall() - conn.close() + with get_db() as conn: + if state: + rows = conn.execute( + "SELECT * FROM tasks WHERE agent_id = ? AND state = ? ORDER BY priority", + (agent_id, state) + ).fetchall() + else: + rows = conn.execute( + "SELECT * FROM tasks WHERE agent_id = ? ORDER BY priority, state", + (agent_id,) + ).fetchall() return [dict(r) for r in rows] def get_blocked_tasks(): """Find tasks whose dependencies aren't DONE yet.""" - conn = get_conn() - tasks = conn.execute( - "SELECT * FROM tasks WHERE state != 'DONE' AND depends_on != '[]'" - ).fetchall() - result = [] - for t in tasks: - deps = json.loads(t["depends_on"]) - for dep_id in deps: - dep = conn.execute("SELECT state FROM tasks WHERE id = ?", (dep_id,)).fetchone() - if dep and dep["state"] != "DONE": - result.append(dict(t)) - break - conn.close() + with get_db() as conn: + tasks = conn.execute( + "SELECT * FROM tasks WHERE state != 'DONE' AND depends_on != '[]'" + ).fetchall() + result = [] + for t in tasks: + deps = json.loads(t["depends_on"]) + for dep_id in deps: + dep = conn.execute("SELECT state FROM tasks WHERE id = ?", (dep_id,)).fetchone() + if dep and dep["state"] != "DONE": + result.append(dict(t)) + break return result # ─── Agent Log Operations ─────────────────────────────────────── def log_action(agent_id, action, task_id=None, details=None, files_modified=None, duration_sec=None, status="SUCCESS"): - conn = get_conn() - conn.execute(""" - INSERT INTO agent_log (agent_id, task_id, action, details, - files_modified, duration_sec, status) - VALUES (?, ?, ?, ?, ?, ?, ?) - """, (agent_id, task_id, action, json.dumps(details) if details else None, - json.dumps(files_modified) if files_modified else None, - duration_sec, status)) - conn.commit() - conn.close() + with get_db() as conn: + conn.execute(""" + INSERT INTO agent_log (agent_id, task_id, action, details, + files_modified, duration_sec, status) + VALUES (?, ?, ?, ?, ?, ?, ?) + """, (agent_id, task_id, action, json.dumps(details) if details else None, + json.dumps(files_modified) if files_modified else None, + duration_sec, status)) # ─── Artifact Operations ──────────────────────────────────────── def register_artifact(task_id, agent_id, file_path, language, lines_of_code, version): - conn = get_conn() - conn.execute(""" - INSERT OR REPLACE INTO artifacts (task_id, agent_id, file_path, language, - lines_of_code, version, updated_at) - VALUES (?, ?, ?, ?, ?, ?, datetime('now')) - """, (task_id, agent_id, file_path, language, lines_of_code, version)) - conn.commit() - conn.close() + with get_db() as conn: + conn.execute(""" + INSERT OR REPLACE INTO artifacts (task_id, agent_id, file_path, language, + lines_of_code, version, updated_at) + VALUES (?, ?, ?, ?, ?, ?, datetime('now')) + """, (task_id, agent_id, file_path, language, lines_of_code, version)) # ─── Benchmark Operations ─────────────────────────────────────── def record_benchmark(build_version, metric, value, unit=None, site_id=None, passed=True): - conn = get_conn() - conn.execute(""" - INSERT INTO benchmarks (build_version, metric, value, unit, site_id, passed) - VALUES (?, ?, ?, ?, ?, ?) - """, (build_version, metric, value, unit, site_id, 1 if passed else 0)) - conn.commit() - conn.close() + with get_db() as conn: + conn.execute(""" + INSERT INTO benchmarks (build_version, metric, value, unit, site_id, passed) + VALUES (?, ?, ?, ?, ?, ?) + """, (build_version, metric, value, unit, site_id, 1 if passed else 0)) def get_latest_benchmarks(build_version=None): - conn = get_conn() - if build_version: - rows = conn.execute( - "SELECT * FROM benchmarks WHERE build_version = ? ORDER BY measured_at DESC", - (build_version,) - ).fetchall() - else: - rows = conn.execute(""" - SELECT * FROM benchmarks WHERE id IN ( - SELECT MAX(id) FROM benchmarks GROUP BY metric - ) ORDER BY metric - """).fetchall() - conn.close() + with get_db() as conn: + if build_version: + rows = conn.execute( + "SELECT * FROM benchmarks WHERE build_version = ? ORDER BY measured_at DESC", + (build_version,) + ).fetchall() + else: + rows = conn.execute(""" + SELECT * FROM benchmarks WHERE id IN ( + SELECT MAX(id) FROM benchmarks GROUP BY metric + ) ORDER BY metric + """).fetchall() return [dict(r) for r in rows] # ─── Deployment Operations ────────────────────────────────────── def record_deployment(site_id, image_version): - conn = get_conn() - conn.execute(""" - INSERT INTO deployments (site_id, image_version) VALUES (?, ?) - """, (site_id, image_version)) - conn.commit() - conn.close() + with get_db() as conn: + conn.execute(""" + INSERT INTO deployments (site_id, image_version) VALUES (?, ?) + """, (site_id, image_version)) def get_fleet_status(): """Current deployment version at each site.""" - conn = get_conn() - rows = conn.execute(""" - SELECT d.* FROM deployments d - INNER JOIN ( - SELECT site_id, MAX(id) as max_id FROM deployments GROUP BY site_id - ) latest ON d.id = latest.max_id - ORDER BY d.site_id - """).fetchall() - conn.close() + with get_db() as conn: + rows = conn.execute(""" + SELECT d.* FROM deployments d + INNER JOIN ( + SELECT site_id, MAX(id) as max_id FROM deployments GROUP BY site_id + ) latest ON d.id = latest.max_id + ORDER BY d.site_id + """).fetchall() return [dict(r) for r in rows] # ─── Sprint Operations ────────────────────────────────────────── def create_sprint(name, goal, start_date, end_date): - conn = get_conn() - conn.execute(""" - INSERT INTO sprints (name, goal, start_date, end_date) VALUES (?, ?, ?, ?) - """, (name, goal, start_date, end_date)) - conn.commit() - sprint_id = conn.execute("SELECT last_insert_rowid()").fetchone()[0] - conn.close() + with get_db() as conn: + conn.execute(""" + INSERT INTO sprints (name, goal, start_date, end_date) VALUES (?, ?, ?, ?) + """, (name, goal, start_date, end_date)) + sprint_id = conn.execute("SELECT last_insert_rowid()").fetchone()[0] return sprint_id def get_sprint_summary(sprint_id): - conn = get_conn() - sprint = dict(conn.execute("SELECT * FROM sprints WHERE id = ?", (sprint_id,)).fetchone()) - tasks = conn.execute("SELECT state, COUNT(*) as cnt FROM tasks WHERE sprint_id = ? GROUP BY state", - (sprint_id,)).fetchall() - sprint["task_counts"] = {r["state"]: r["cnt"] for r in tasks} - conn.close() + with get_db() as conn: + sprint = dict(conn.execute("SELECT * FROM sprints WHERE id = ?", (sprint_id,)).fetchone()) + tasks = conn.execute("SELECT state, COUNT(*) as cnt FROM tasks WHERE sprint_id = ? GROUP BY state", + (sprint_id,)).fetchall() + sprint["task_counts"] = {r["state"]: r["cnt"] for r in tasks} return sprint # ─── Reporting ────────────────────────────────────────────────── def get_agent_stats(): """Summary of each agent's activity.""" - conn = get_conn() - rows = conn.execute(""" - SELECT agent_id, - COUNT(*) as total_actions, - SUM(CASE WHEN status='SUCCESS' THEN 1 ELSE 0 END) as successes, - SUM(CASE WHEN status='FAILED' THEN 1 ELSE 0 END) as failures, - ROUND(SUM(duration_sec), 1) as total_seconds - FROM agent_log - GROUP BY agent_id - ORDER BY agent_id - """).fetchall() - conn.close() + with get_db() as conn: + rows = conn.execute(""" + SELECT agent_id, + COUNT(*) as total_actions, + SUM(CASE WHEN status='SUCCESS' THEN 1 ELSE 0 END) as successes, + SUM(CASE WHEN status='FAILED' THEN 1 ELSE 0 END) as failures, + ROUND(SUM(duration_sec), 1) as total_seconds + FROM agent_log + GROUP BY agent_id + ORDER BY agent_id + """).fetchall() return [dict(r) for r in rows] def dashboard(): """Print full system status.""" - conn = get_conn() - task_states = conn.execute( - "SELECT state, COUNT(*) as cnt FROM tasks GROUP BY state" - ).fetchall() + with get_db() as conn: + task_states = conn.execute( + "SELECT state, COUNT(*) as cnt FROM tasks GROUP BY state" + ).fetchall() total_tasks = sum(r["cnt"] for r in task_states) done_tasks = next((r["cnt"] for r in task_states if r["state"] == "DONE"), 0) agent_stats = get_agent_stats() fleet = get_fleet_status() benchmarks = get_latest_benchmarks() - conn.close() print("\n" + "="*60) print(" SENTINELFORGE DASHBOARD") diff --git a/tests/test_db_lifecycle.py b/tests/test_db_lifecycle.py new file mode 100644 index 0000000..ce652f2 --- /dev/null +++ b/tests/test_db_lifecycle.py @@ -0,0 +1,34 @@ +import sqlite3 + +import pytest + +import db + + +def test_database_context_rolls_back_and_closes_on_error(tmp_path, monkeypatch): + monkeypatch.setattr(db, "DB_PATH", tmp_path / "sentinelforge.db") + db.init_db() + + with pytest.raises(RuntimeError): + with db.get_db() as conn: + conn.execute( + "INSERT INTO deployments (site_id, image_version) VALUES (?, ?)", + ("site-1", "test"), + ) + raise RuntimeError("abort transaction") + + with pytest.raises(sqlite3.ProgrammingError): + conn.execute("SELECT 1") + assert db.get_fleet_status() == [] + + +def test_task_state_update_uses_fixed_sql_and_persists(tmp_path, monkeypatch): + monkeypatch.setattr(db, "DB_PATH", tmp_path / "sentinelforge.db") + db.init_db() + task_id = db.create_task(None, "Test", "Description", agent_id=1) + + db.update_task_state(task_id, "DONE") + + task = db.get_tasks_for_agent(1)[0] + assert task["state"] == "DONE" + assert task["completed_at"] is not None