Merge pull request #201 from tmchow/fix/90-store-sql-column-whitelist
fix(store): validate updatable columns in update_run and update_finding
This commit is contained in:
@@ -132,6 +132,32 @@ INSERT OR IGNORE INTO settings (key, value) VALUES ('briefing_format', 'concise'
|
||||
INSERT OR IGNORE INTO settings (key, value) VALUES ('default_schedule', '0 8 * * *');
|
||||
"""
|
||||
|
||||
_UPDATABLE_RUN_COLUMNS = frozenset({
|
||||
"source_mode",
|
||||
"prompt_tokens",
|
||||
"completion_tokens",
|
||||
"token_cost",
|
||||
"duration_seconds",
|
||||
"status",
|
||||
"error_message",
|
||||
"findings_new",
|
||||
"findings_updated",
|
||||
})
|
||||
|
||||
_UPDATABLE_FINDING_COLUMNS = frozenset({
|
||||
"source",
|
||||
"source_url",
|
||||
"source_title",
|
||||
"author",
|
||||
"content",
|
||||
"summary",
|
||||
"engagement_score",
|
||||
"relevance_score",
|
||||
"last_seen",
|
||||
"sighting_count",
|
||||
"dismissed",
|
||||
})
|
||||
|
||||
# Future migrations keyed by version number
|
||||
MIGRATIONS: Dict[int, str] = {}
|
||||
|
||||
@@ -298,6 +324,11 @@ def update_run(run_id: int, **kwargs):
|
||||
"""Update a research run's fields."""
|
||||
conn = _connect()
|
||||
try:
|
||||
invalid_columns = sorted(set(kwargs) - _UPDATABLE_RUN_COLUMNS)
|
||||
if invalid_columns:
|
||||
raise ValueError(
|
||||
f"Invalid run update fields: {', '.join(invalid_columns)}"
|
||||
)
|
||||
sets = ", ".join(f"{k} = ?" for k in kwargs)
|
||||
values = list(kwargs.values()) + [run_id]
|
||||
conn.execute(f"UPDATE research_runs SET {sets} WHERE id = ?", values)
|
||||
@@ -430,6 +461,11 @@ def update_finding(finding_id: int, **kwargs):
|
||||
"""Update a finding's fields."""
|
||||
conn = _connect()
|
||||
try:
|
||||
invalid_columns = sorted(set(kwargs) - _UPDATABLE_FINDING_COLUMNS)
|
||||
if invalid_columns:
|
||||
raise ValueError(
|
||||
f"Invalid finding update fields: {', '.join(invalid_columns)}"
|
||||
)
|
||||
sets = ", ".join(f"{k} = ?" for k in kwargs)
|
||||
values = list(kwargs.values()) + [finding_id]
|
||||
conn.execute(f"UPDATE findings SET {sets} WHERE id = ?", values)
|
||||
|
||||
@@ -384,6 +384,31 @@ def test_store_findings_skips_items_without_url(temp_db):
|
||||
assert counts["new"] == 1
|
||||
|
||||
|
||||
def test_update_validates_allowed_columns(temp_db, sample_report):
|
||||
"""Test update_run/update_finding accept valid keys and reject invalid keys."""
|
||||
topic = store.add_topic("Test Topic")
|
||||
run_id = store.record_run(topic["id"], source_mode="v3")
|
||||
|
||||
# Valid run update key should not raise
|
||||
store.update_run(run_id, status="failed")
|
||||
|
||||
findings = store.findings_from_report(sample_report)
|
||||
store.store_findings(run_id, topic["id"], findings[:1])
|
||||
|
||||
conn = sqlite3.connect(str(temp_db))
|
||||
finding_id = conn.execute("SELECT id FROM findings LIMIT 1").fetchone()[0]
|
||||
conn.close()
|
||||
|
||||
# Valid finding update key should not raise
|
||||
store.update_finding(finding_id, dismissed=1)
|
||||
|
||||
with pytest.raises(ValueError, match="invalid_run_column"):
|
||||
store.update_run(run_id, invalid_run_column="x")
|
||||
|
||||
with pytest.raises(ValueError, match="invalid_finding_column"):
|
||||
store.update_finding(finding_id, invalid_finding_column="x")
|
||||
|
||||
|
||||
# === Tests for topic management ===
|
||||
|
||||
def test_add_topic(temp_db):
|
||||
|
||||
Reference in New Issue
Block a user