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 * * *');
|
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
|
# Future migrations keyed by version number
|
||||||
MIGRATIONS: Dict[int, str] = {}
|
MIGRATIONS: Dict[int, str] = {}
|
||||||
|
|
||||||
@@ -298,6 +324,11 @@ def update_run(run_id: int, **kwargs):
|
|||||||
"""Update a research run's fields."""
|
"""Update a research run's fields."""
|
||||||
conn = _connect()
|
conn = _connect()
|
||||||
try:
|
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)
|
sets = ", ".join(f"{k} = ?" for k in kwargs)
|
||||||
values = list(kwargs.values()) + [run_id]
|
values = list(kwargs.values()) + [run_id]
|
||||||
conn.execute(f"UPDATE research_runs SET {sets} WHERE id = ?", values)
|
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."""
|
"""Update a finding's fields."""
|
||||||
conn = _connect()
|
conn = _connect()
|
||||||
try:
|
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)
|
sets = ", ".join(f"{k} = ?" for k in kwargs)
|
||||||
values = list(kwargs.values()) + [finding_id]
|
values = list(kwargs.values()) + [finding_id]
|
||||||
conn.execute(f"UPDATE findings SET {sets} WHERE id = ?", values)
|
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
|
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 ===
|
# === Tests for topic management ===
|
||||||
|
|
||||||
def test_add_topic(temp_db):
|
def test_add_topic(temp_db):
|
||||||
|
|||||||
Reference in New Issue
Block a user