fix(db): transaction context manager + fixed parameterized updates (CORR-10)

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")

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