fix(campaigns): use RETURNING in create/update_campaign to avoid extra round-trip and TOCTOU race

create_campaign/update_campaign previously committed then opened a
second connection via get_campaign() just to re-read the row they had
just written, adding an extra round-trip and a race window where a
concurrent delete between commit and re-fetch was misreported as
"disappeared immediately after insert/update". Both now RETURNING the
row from the same statement/cursor that wrote it, using a shared
_RETURNING_CLAUSE derived from _SELECT_COLUMNS so the two column lists
cannot drift apart. update_campaign now raises LookupError for an
unknown id (0 rows affected) instead of misdiagnosing it as a race.

Also: delete_campaign uses the (cur.rowcount or 0) > 0 idiom already
established in postgres_repo.py, the module docstring documents that
this repository's SQL paths rely on integration tests rather than
unit tests, and a new DB-free test asserts the RETURNING and SELECT
column lists stay identical.
This commit is contained in:
AFFAANh
2026-08-01 13:36:36 +05:30
parent 1df49aee37
commit 4098af5a55
2 changed files with 57 additions and 13 deletions
+27 -12
View File
@@ -1,5 +1,13 @@
# -*- coding: utf-8 -*-
"""PostgreSQL repository for campaigns and their events."""
"""PostgreSQL repository for campaigns and their events.
Note: CampaignRepository's SQL paths (list/get/create/update/delete,
add_event/list_events) are not covered by unit tests in this module —
they require a live PostgreSQL connection and are exercised by
integration tests instead. Only the pure row-mapping helper
(`campaign_from_row`) and the RETURNING/SELECT column contract are
unit-tested here.
"""
from __future__ import annotations
import uuid
@@ -15,6 +23,13 @@ _SELECT_COLUMNS = """
created_at, updated_at
"""
# Reused verbatim by create_campaign/update_campaign so the row a write
# returns always has exactly the columns _SELECT_COLUMNS reads back on a
# plain SELECT. Keeping this a single f-string derived from
# _SELECT_COLUMNS (rather than a second hand-written column list) makes
# the two impossible to drift apart.
_RETURNING_CLAUSE = f"RETURNING {_SELECT_COLUMNS}"
def campaign_from_row(row: dict[str, Any]) -> CampaignSpec:
"""Build a CampaignSpec from a database row."""
@@ -98,7 +113,8 @@ INSERT INTO maskanx_campaigns (
) VALUES (
%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
)
""",
"""
+ _RETURNING_CLAUSE,
(
spec.id,
spec.company_id,
@@ -115,16 +131,14 @@ INSERT INTO maskanx_campaigns (
jsonb(spec.schedule),
),
)
row = await cur.fetchone()
await conn.commit()
except Exception:
await conn.rollback()
raise
finally:
await conn.close()
created = await self.get_campaign(spec.id)
if created is None:
raise RuntimeError(f"Campaign {spec.id} disappeared immediately after insert")
return created
return campaign_from_row(row)
async def update_campaign(self, spec: CampaignSpec) -> CampaignSpec:
conn = await connect_database()
@@ -138,7 +152,8 @@ UPDATE maskanx_campaigns SET
channels = %s, schedule = %s, sync_status = %s, sync_error = %s,
approved_by = %s, updated_at = NOW()
WHERE id = %s
""",
"""
+ _RETURNING_CLAUSE,
(
spec.name,
spec.status,
@@ -156,16 +171,16 @@ WHERE id = %s
spec.id,
),
)
row = await cur.fetchone()
if row is None:
raise LookupError(f"Campaign {spec.id} does not exist")
await conn.commit()
except Exception:
await conn.rollback()
raise
finally:
await conn.close()
updated = await self.get_campaign(spec.id)
if updated is None:
raise RuntimeError(f"Campaign {spec.id} disappeared immediately after update")
return updated
return campaign_from_row(row)
async def delete_campaign(self, campaign_id: str) -> bool:
conn = await connect_database()
@@ -175,7 +190,7 @@ WHERE id = %s
"DELETE FROM maskanx_campaigns WHERE id = %s",
(campaign_id,),
)
deleted = cur.rowcount > 0
deleted = (cur.rowcount or 0) > 0
await conn.commit()
except Exception:
await conn.rollback()
+30 -1
View File
@@ -1,8 +1,10 @@
# -*- coding: utf-8 -*-
"""Row mapping for the campaign repository."""
import inspect
from datetime import datetime, timezone
from adclaw.campaigns.repo import campaign_from_row
from adclaw.campaigns import repo as repo_module
from adclaw.campaigns.repo import CampaignRepository, campaign_from_row
def _row(**overrides):
@@ -51,3 +53,30 @@ def test_campaign_from_row_defaults_null_json_to_empty():
assert spec.advanced == {}
assert spec.channels == []
assert spec.schedule == {}
def _column_names(column_block):
return [c.strip() for c in column_block.strip().split(",") if c.strip()]
def test_returning_columns_match_select_columns():
"""create_campaign/update_campaign RETURNING must match _SELECT_COLUMNS
exactly, so the row a write returns maps to a CampaignSpec the same way
a plain SELECT would, without a second query. Guards against the two
column lists silently drifting apart.
"""
select_columns = _column_names(repo_module._SELECT_COLUMNS)
returning_columns = _column_names(
repo_module._RETURNING_CLAUSE.split("RETURNING", 1)[1],
)
assert returning_columns == select_columns
# Both methods must build their SQL from the shared _RETURNING_CLAUSE
# constant rather than a hand-duplicated column list that could
# silently diverge from _SELECT_COLUMNS.
for method in (CampaignRepository.create_campaign, CampaignRepository.update_campaign):
source = inspect.getsource(method)
assert "_RETURNING_CLAUSE" in source, (
f"{method.__name__} must reuse _RETURNING_CLAUSE, not a "
"hand-written RETURNING column list"
)