1
0

[loop-cycle-50] refactor: replace bare sqlite3.connect() with context managers batch 2 (#157) (#180)

This commit is contained in:
2026-03-15 11:58:43 -04:00
parent bea2749158
commit bcd6d7e321
16 changed files with 512 additions and 510 deletions

View File

@@ -9,8 +9,8 @@ import asyncio
import json
import logging
import sqlite3
from collections.abc import Callable, Coroutine
from contextlib import closing
from collections.abc import Callable, Coroutine, Generator
from contextlib import closing, contextmanager
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
@@ -106,22 +106,23 @@ class EventBus:
conn.executescript(_EVENTS_SCHEMA)
conn.commit()
def _get_persistence_conn(self) -> sqlite3.Connection | None:
@contextmanager
def _get_persistence_conn(self) -> Generator[sqlite3.Connection | None, None, None]:
"""Get a connection to the persistence database."""
if self._persistence_db_path is None:
return None
conn = sqlite3.connect(str(self._persistence_db_path))
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA busy_timeout=5000")
return conn
yield None
return
with closing(sqlite3.connect(str(self._persistence_db_path))) as conn:
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA busy_timeout=5000")
yield conn
def _persist_event(self, event: Event) -> None:
"""Write an event to the persistence database."""
conn = self._get_persistence_conn()
if conn is None:
return
try:
with closing(conn):
with self._get_persistence_conn() as conn:
if conn is None:
return
try:
task_id = event.data.get("task_id", "")
agent_id = event.data.get("agent_id", "")
conn.execute(
@@ -139,8 +140,8 @@ class EventBus:
),
)
conn.commit()
except Exception as exc:
logger.debug("Failed to persist event: %s", exc)
except Exception as exc:
logger.debug("Failed to persist event: %s", exc)
# ── Replay ───────────────────────────────────────────────────────────
@@ -162,12 +163,11 @@ class EventBus:
Returns:
List of Event objects from persistent storage.
"""
conn = self._get_persistence_conn()
if conn is None:
return []
with self._get_persistence_conn() as conn:
if conn is None:
return []
try:
with closing(conn):
try:
conditions = []
params: list = []
@@ -197,9 +197,9 @@ class EventBus:
)
for row in rows
]
except Exception as exc:
logger.debug("Failed to replay events: %s", exc)
return []
except Exception as exc:
logger.debug("Failed to replay events: %s", exc)
return []
# ── Subscribe / Publish ──────────────────────────────────────────────

View File

@@ -11,7 +11,8 @@ model roles (student, teacher, judge/PRM) run on dedicated resources.
import logging
import sqlite3
import threading
from contextlib import closing
from collections.abc import Generator
from contextlib import closing, contextmanager
from dataclasses import dataclass
from datetime import UTC, datetime
from enum import StrEnum
@@ -61,36 +62,37 @@ class CustomModel:
self.registered_at = datetime.now(UTC).isoformat()
def _get_conn() -> sqlite3.Connection:
@contextmanager
def _get_conn() -> Generator[sqlite3.Connection, None, None]:
DB_PATH.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(str(DB_PATH))
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA busy_timeout=5000")
conn.execute("""
CREATE TABLE IF NOT EXISTS custom_models (
name TEXT PRIMARY KEY,
format TEXT NOT NULL,
path TEXT NOT NULL,
role TEXT NOT NULL DEFAULT 'general',
context_window INTEGER NOT NULL DEFAULT 4096,
description TEXT NOT NULL DEFAULT '',
registered_at TEXT NOT NULL,
active INTEGER NOT NULL DEFAULT 1,
default_temperature REAL NOT NULL DEFAULT 0.7,
max_tokens INTEGER NOT NULL DEFAULT 2048
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS agent_model_assignments (
agent_id TEXT PRIMARY KEY,
model_name TEXT NOT NULL,
assigned_at TEXT NOT NULL,
FOREIGN KEY (model_name) REFERENCES custom_models(name)
)
""")
conn.commit()
return conn
with closing(sqlite3.connect(str(DB_PATH))) as conn:
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA busy_timeout=5000")
conn.execute("""
CREATE TABLE IF NOT EXISTS custom_models (
name TEXT PRIMARY KEY,
format TEXT NOT NULL,
path TEXT NOT NULL,
role TEXT NOT NULL DEFAULT 'general',
context_window INTEGER NOT NULL DEFAULT 4096,
description TEXT NOT NULL DEFAULT '',
registered_at TEXT NOT NULL,
active INTEGER NOT NULL DEFAULT 1,
default_temperature REAL NOT NULL DEFAULT 0.7,
max_tokens INTEGER NOT NULL DEFAULT 2048
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS agent_model_assignments (
agent_id TEXT PRIMARY KEY,
model_name TEXT NOT NULL,
assigned_at TEXT NOT NULL,
FOREIGN KEY (model_name) REFERENCES custom_models(name)
)
""")
conn.commit()
yield conn
class ModelRegistry:
@@ -106,7 +108,7 @@ class ModelRegistry:
def _load_from_db(self) -> None:
"""Bootstrap cache from SQLite."""
try:
with closing(_get_conn()) as conn:
with _get_conn() as conn:
for row in conn.execute("SELECT * FROM custom_models WHERE active = 1").fetchall():
self._models[row["name"]] = CustomModel(
name=row["name"],
@@ -130,7 +132,7 @@ class ModelRegistry:
def register(self, model: CustomModel) -> CustomModel:
"""Register a new custom model."""
with self._lock:
with closing(_get_conn()) as conn:
with _get_conn() as conn:
conn.execute(
"""
INSERT OR REPLACE INTO custom_models
@@ -161,7 +163,7 @@ class ModelRegistry:
with self._lock:
if name not in self._models:
return False
with closing(_get_conn()) as conn:
with _get_conn() as conn:
conn.execute("DELETE FROM custom_models WHERE name = ?", (name,))
conn.execute("DELETE FROM agent_model_assignments WHERE model_name = ?", (name,))
conn.commit()
@@ -191,7 +193,7 @@ class ModelRegistry:
return False
with self._lock:
model.active = active
with closing(_get_conn()) as conn:
with _get_conn() as conn:
conn.execute(
"UPDATE custom_models SET active = ? WHERE name = ?",
(int(active), name),
@@ -207,7 +209,7 @@ class ModelRegistry:
return False
with self._lock:
now = datetime.now(UTC).isoformat()
with closing(_get_conn()) as conn:
with _get_conn() as conn:
conn.execute(
"""
INSERT OR REPLACE INTO agent_model_assignments
@@ -226,7 +228,7 @@ class ModelRegistry:
with self._lock:
if agent_id not in self._agent_assignments:
return False
with closing(_get_conn()) as conn:
with _get_conn() as conn:
conn.execute(
"DELETE FROM agent_model_assignments WHERE agent_id = ?",
(agent_id,),