Skip to content

Commit b596dd2

Browse files
jopemachineclaude
jopemachine
and
claude
committed
refactor(BA-6554): make the fragment write-gate atomic in the repository (review)
Move the allow-list write-gate out of the service and into the fragment db_source so the gate check and the write share one transaction (no check-then-write race). WriteOps is read-capable, so create/update/purge run `w.exists(AppConfigAllowListRow ...)` inside their own write tx; update/purge also fetch the target row in that same tx, removing the service's separate get_by_id (3 txs -> 1). The service now purely delegates and no longer depends on AppConfigAllowListRepository. Gate pass/reject is now covered by real-DB repository tests; the service tests verify delegation. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 53435a5 commit b596dd2

4 files changed

Lines changed: 158 additions & 161 deletions

File tree

‎src/ai/backend/manager/repositories/app_config_fragment/db_source/db_source.py‎

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66

77
import sqlalchemy as sa
88

9+
from ai.backend.common.data.app_config.types import AppConfigScopeType
10+
from ai.backend.common.data.filter_specs import StringMatchSpec
911
from ai.backend.common.exception import BackendAIError
1012
from ai.backend.common.identifier.app_config_fragment import AppConfigFragmentID
1113
from ai.backend.common.metrics.metric import DomainType, LayerType
@@ -16,21 +18,29 @@
1618
AppConfigFragmentData,
1719
AppConfigFragmentSearchResult,
1820
)
19-
from ai.backend.manager.errors.app_config import AppConfigFragmentNotFound
21+
from ai.backend.manager.errors.app_config import (
22+
AppConfigFragmentNotFound,
23+
AppConfigFragmentWriteNotAllowed,
24+
)
25+
from ai.backend.manager.models.app_config_allow_list.conditions import (
26+
AppConfigAllowListConditions,
27+
)
28+
from ai.backend.manager.models.app_config_allow_list.row import AppConfigAllowListRow
2029
from ai.backend.manager.models.app_config_definition.row import AppConfigDefinitionRow
2130
from ai.backend.manager.models.app_config_fragment.row import AppConfigFragmentRow
2231
from ai.backend.manager.repositories.app_config_fragment.creators import (
2332
AppConfigFragmentCreatorSpec,
2433
)
2534
from ai.backend.manager.repositories.base import (
2635
BatchQuerier,
36+
ExistsQuerier,
2737
Purger,
2838
Querier,
2939
SearchScope,
3040
Updater,
3141
)
3242
from ai.backend.manager.repositories.base.creator import NextValuePolicy
33-
from ai.backend.manager.repositories.ops import DBOpsProvider
43+
from ai.backend.manager.repositories.ops import DBOpsProvider, WriteOps
3444

3545
__all__ = ("AppConfigFragmentDBSource",)
3646

@@ -65,6 +75,32 @@ class AppConfigFragmentDBSource:
6575
def __init__(self, ops_provider: DBOpsProvider) -> None:
6676
self._ops = ops_provider
6777

78+
async def _ensure_write_allowed(
79+
self, w: WriteOps, config_name: str, scope_type: AppConfigScopeType
80+
) -> None:
81+
"""Reject the write unless an allow-list row exists for ``(config_name, scope_type)``.
82+
83+
Runs inside the caller's write transaction (``WriteOps`` is read-capable), so the
84+
gate check and the write commit atomically — no check-then-write race. Because an
85+
allow-list row requires a registered ``config_name`` (FK), this also enforces
86+
registration.
87+
"""
88+
allowed = await w.exists(
89+
ExistsQuerier(
90+
row_class=AppConfigAllowListRow,
91+
conditions=[
92+
AppConfigAllowListConditions.by_config_name_equals(
93+
StringMatchSpec(config_name, case_insensitive=False, negated=False)
94+
),
95+
AppConfigAllowListConditions.by_scope_type_equals(scope_type),
96+
],
97+
)
98+
)
99+
if not allowed:
100+
raise AppConfigFragmentWriteNotAllowed(
101+
f"Writing app config {config_name!r} at scope {scope_type.value!r} is not allowed."
102+
)
103+
68104
@app_config_fragment_db_source_resilience.apply()
69105
async def create(self, spec: AppConfigFragmentCreatorSpec) -> AppConfigFragmentData:
70106
policy = NextValuePolicy(
@@ -76,6 +112,7 @@ async def create(self, spec: AppConfigFragmentCreatorSpec) -> AppConfigFragmentD
76112
gap=RANK_GAP,
77113
)
78114
async with self._ops.write_ops() as w:
115+
await self._ensure_write_allowed(w, spec.config_name, spec.scope_type)
79116
created = await w.create_with_next_value(policy, spec)
80117
return created.row.to_data()
81118

@@ -90,6 +127,12 @@ async def get_by_id(self, fragment_id: AppConfigFragmentID) -> AppConfigFragment
90127
@app_config_fragment_db_source_resilience.apply()
91128
async def update(self, updater: Updater[AppConfigFragmentRow]) -> AppConfigFragmentData:
92129
async with self._ops.write_ops() as w:
130+
existing = await w.query(
131+
Querier(row_class=AppConfigFragmentRow, pk_value=updater.pk_value)
132+
)
133+
if existing is None:
134+
raise AppConfigFragmentNotFound(f"App config fragment {updater.pk_value} not found")
135+
await self._ensure_write_allowed(w, existing.row.config_name, existing.row.scope_type)
93136
result = await w.update(updater)
94137
if result is None:
95138
raise AppConfigFragmentNotFound(f"App config fragment {updater.pk_value} not found")
@@ -98,6 +141,12 @@ async def update(self, updater: Updater[AppConfigFragmentRow]) -> AppConfigFragm
98141
@app_config_fragment_db_source_resilience.apply()
99142
async def purge(self, purger: Purger[AppConfigFragmentRow]) -> AppConfigFragmentData:
100143
async with self._ops.write_ops() as w:
144+
existing = await w.query(
145+
Querier(row_class=AppConfigFragmentRow, pk_value=purger.pk_value)
146+
)
147+
if existing is None:
148+
raise AppConfigFragmentNotFound(f"App config fragment {purger.pk_value} not found")
149+
await self._ensure_write_allowed(w, existing.row.config_name, existing.row.scope_type)
101150
result = await w.purge(purger)
102151
if result is None:
103152
raise AppConfigFragmentNotFound(f"App config fragment {purger.pk_value} not found")

‎src/ai/backend/manager/services/app_config_fragment/service.py‎

Lines changed: 8 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,8 @@
11
from __future__ import annotations
22

3-
from typing import cast
4-
5-
from ai.backend.common.data.app_config.types import AppConfigScopeType
6-
from ai.backend.common.data.filter_specs import StringMatchSpec
7-
from ai.backend.common.identifier.app_config_fragment import AppConfigFragmentID
8-
from ai.backend.manager.errors.app_config import AppConfigFragmentWriteNotAllowed
9-
from ai.backend.manager.models.app_config_allow_list.conditions import (
10-
AppConfigAllowListConditions,
11-
)
12-
from ai.backend.manager.models.app_config_allow_list.row import AppConfigAllowListRow
13-
from ai.backend.manager.repositories.app_config_allow_list.repository import (
14-
AppConfigAllowListRepository,
15-
)
163
from ai.backend.manager.repositories.app_config_fragment.repository import (
174
AppConfigFragmentRepository,
185
)
19-
from ai.backend.manager.repositories.base import ExistsQuerier
206
from ai.backend.manager.services.app_config_fragment.actions.admin_search import (
217
AdminSearchAppConfigFragmentAction,
228
AdminSearchAppConfigFragmentActionResult,
@@ -46,50 +32,24 @@
4632

4733

4834
class AppConfigFragmentService:
49-
"""Write paths for app config fragments (not admin-only).
35+
"""Write/read paths for app config fragments (not admin-only).
5036
51-
``create``, ``update``, and ``purge`` all pass the allow-list write-gate: a fragment may
52-
only be written or removed when an ``app_config_allow_list`` row exists for its
53-
``(config_name, scope_type)``. Because an allow-list row itself requires a registered
54-
``config_name`` (FK to ``app_config_definitions``), this single check also enforces
55-
registration. An allow-listed user may therefore manage their own ``user``-scope
56-
fragment without admin privileges.
37+
``create`` / ``update`` / ``purge`` are gated by the allow-list write-gate in the
38+
repository: the gate check and the write run in a single transaction, so a fragment is
39+
only written or removed when an ``app_config_allow_list`` row exists for its
40+
``(config_name, scope_type)``. Because an allow-list row requires a registered
41+
``config_name`` (FK), this also enforces registration. An allow-listed user may
42+
therefore manage their own ``user``-scope fragment without admin privileges.
5743
"""
5844

5945
_repository: AppConfigFragmentRepository
60-
_allow_list_repository: AppConfigAllowListRepository
6146

62-
def __init__(
63-
self,
64-
repository: AppConfigFragmentRepository,
65-
allow_list_repository: AppConfigAllowListRepository,
66-
) -> None:
47+
def __init__(self, repository: AppConfigFragmentRepository) -> None:
6748
self._repository = repository
68-
self._allow_list_repository = allow_list_repository
69-
70-
async def _ensure_write_allowed(self, config_name: str, scope_type: AppConfigScopeType) -> None:
71-
allowed = await self._allow_list_repository.exists(
72-
ExistsQuerier(
73-
row_class=AppConfigAllowListRow,
74-
conditions=[
75-
AppConfigAllowListConditions.by_config_name_equals(
76-
StringMatchSpec(config_name, case_insensitive=False, negated=False)
77-
),
78-
AppConfigAllowListConditions.by_scope_type_equals(scope_type),
79-
],
80-
)
81-
)
82-
if not allowed:
83-
raise AppConfigFragmentWriteNotAllowed(
84-
f"Writing app config {config_name!r} at scope {scope_type.value!r} is not allowed."
85-
)
8649

8750
async def create(
8851
self, action: CreateAppConfigFragmentAction
8952
) -> CreateAppConfigFragmentActionResult:
90-
await self._ensure_write_allowed(
91-
action.creator_spec.config_name, action.creator_spec.scope_type
92-
)
9353
data = await self._repository.create(action.creator_spec)
9454
return CreateAppConfigFragmentActionResult(fragment=data)
9555

@@ -125,19 +85,11 @@ async def scoped_search(
12585
async def update(
12686
self, action: UpdateAppConfigFragmentAction
12787
) -> UpdateAppConfigFragmentActionResult:
128-
existing = await self._repository.get_by_id(
129-
cast(AppConfigFragmentID, action.updater.pk_value)
130-
)
131-
await self._ensure_write_allowed(existing.config_name, existing.scope_type)
13288
data = await self._repository.update(action.updater)
13389
return UpdateAppConfigFragmentActionResult(fragment=data)
13490

13591
async def purge(
13692
self, action: PurgeAppConfigFragmentAction
13793
) -> PurgeAppConfigFragmentActionResult:
138-
existing = await self._repository.get_by_id(
139-
cast(AppConfigFragmentID, action.purger.pk_value)
140-
)
141-
await self._ensure_write_allowed(existing.config_name, existing.scope_type)
14294
data = await self._repository.purge(action.purger)
14395
return PurgeAppConfigFragmentActionResult(fragment=data)

‎tests/unit/manager/repositories/app_config_fragment/test_repository.py‎

Lines changed: 86 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,12 @@
1515
from ai.backend.manager.data.app_config_fragment.types import (
1616
AppConfigFragmentData,
1717
)
18-
from ai.backend.manager.errors.app_config import AppConfigFragmentNotFound
18+
from ai.backend.manager.errors.app_config import (
19+
AppConfigFragmentNotFound,
20+
AppConfigFragmentWriteNotAllowed,
21+
)
1922
from ai.backend.manager.errors.repository import UniqueConstraintViolationError
23+
from ai.backend.manager.models.app_config_allow_list.row import AppConfigAllowListRow
2024
from ai.backend.manager.models.app_config_definition.row import AppConfigDefinitionRow
2125
from ai.backend.manager.models.app_config_fragment.conditions import AppConfigFragmentConditions
2226
from ai.backend.manager.models.app_config_fragment.orders import AppConfigFragmentOrders
@@ -51,8 +55,11 @@
5155
async def database(
5256
database_connection: ExtendedAsyncSAEngine,
5357
) -> AsyncGenerator[ExtendedAsyncSAEngine, None]:
54-
# FK order: app_config_definitions (parent) before app_config_fragments (child).
55-
async with with_tables(database_connection, [AppConfigDefinitionRow, AppConfigFragmentRow]):
58+
# FK order: app_config_definitions (parent) before the allow-list and fragments (children).
59+
async with with_tables(
60+
database_connection,
61+
[AppConfigDefinitionRow, AppConfigAllowListRow, AppConfigFragmentRow],
62+
):
5663
yield database_connection
5764

5865

@@ -63,15 +70,53 @@ def repository(database: ExtendedAsyncSAEngine) -> AppConfigFragmentRepository:
6370

6471
@pytest.fixture
6572
async def theme_registered(database: ExtendedAsyncSAEngine) -> None:
66-
"""Situation: ``theme`` is registered (FK parent) but no fragments exist yet."""
73+
"""Situation: ``theme`` is registered and allow-listed at every scope; no fragments yet."""
74+
async with database.begin_session() as db_sess:
75+
db_sess.add(AppConfigDefinitionRow(config_name="theme"))
76+
await db_sess.flush()
77+
db_sess.add_all([
78+
AppConfigAllowListRow(config_name="theme", scope_type=AppConfigScopeType.PUBLIC),
79+
AppConfigAllowListRow(config_name="theme", scope_type=AppConfigScopeType.DOMAIN),
80+
AppConfigAllowListRow(config_name="theme", scope_type=AppConfigScopeType.USER),
81+
])
82+
await db_sess.flush()
83+
84+
85+
@pytest.fixture
86+
async def theme_defined_not_allow_listed(database: ExtendedAsyncSAEngine) -> None:
87+
"""Situation: ``theme`` is registered but has no allow-list row (writes are gated)."""
6788
async with database.begin_session() as db_sess:
6889
db_sess.add(AppConfigDefinitionRow(config_name="theme"))
6990
await db_sess.flush()
7091

7192

7293
@pytest.fixture
7394
async def domain_scoped_fragment(database: ExtendedAsyncSAEngine) -> AppConfigFragmentData:
74-
"""Situation: a single pre-existing ``theme``/domain fragment."""
95+
"""Situation: a pre-existing ``theme``/domain fragment, allow-listed at the domain scope."""
96+
async with database.begin_session() as db_sess:
97+
db_sess.add(AppConfigDefinitionRow(config_name="theme"))
98+
await db_sess.flush()
99+
db_sess.add(
100+
AppConfigAllowListRow(config_name="theme", scope_type=AppConfigScopeType.DOMAIN)
101+
)
102+
row = AppConfigFragmentRow(
103+
config_name="theme",
104+
scope_type=AppConfigScopeType.DOMAIN,
105+
scope_id=_DOMAIN_ID,
106+
rank=100,
107+
config={"k": "v"},
108+
)
109+
db_sess.add(row)
110+
await db_sess.flush()
111+
return row.to_data()
112+
113+
114+
@pytest.fixture
115+
async def fragment_not_allow_listed(database: ExtendedAsyncSAEngine) -> AppConfigFragmentData:
116+
"""Situation: a ``theme``/domain fragment whose ``(config_name, scope)`` is NOT allow-listed.
117+
118+
Models a row that survives after its allow-list entry was removed — writes must be gated.
119+
"""
75120
async with database.begin_session() as db_sess:
76121
db_sess.add(AppConfigDefinitionRow(config_name="theme"))
77122
await db_sess.flush()
@@ -171,6 +216,19 @@ async def test_get_by_id_missing_raises(self, repository: AppConfigFragmentRepos
171216
with pytest.raises(AppConfigFragmentNotFound):
172217
await repository.get_by_id(AppConfigFragmentID(uuid.uuid4()))
173218

219+
async def test_create_rejected_when_not_allow_listed(
220+
self, repository: AppConfigFragmentRepository, theme_defined_not_allow_listed: None
221+
) -> None:
222+
with pytest.raises(AppConfigFragmentWriteNotAllowed):
223+
await repository.create(
224+
AppConfigFragmentCreatorSpec(
225+
config_name="theme",
226+
scope_type=AppConfigScopeType.PUBLIC,
227+
scope_id="public",
228+
config={"theme": "dark"},
229+
)
230+
)
231+
174232
async def test_unique_constraint_violation(
175233
self,
176234
repository: AppConfigFragmentRepository,
@@ -242,6 +300,19 @@ async def test_update_missing_raises(self, repository: AppConfigFragmentReposito
242300
)
243301
)
244302

303+
async def test_update_rejected_when_not_allow_listed(
304+
self,
305+
repository: AppConfigFragmentRepository,
306+
fragment_not_allow_listed: AppConfigFragmentData,
307+
) -> None:
308+
with pytest.raises(AppConfigFragmentWriteNotAllowed):
309+
await repository.update(
310+
Updater(
311+
spec=AppConfigFragmentUpdaterSpec(config=OptionalState.update({"b": 2})),
312+
pk_value=fragment_not_allow_listed.id,
313+
)
314+
)
315+
245316

246317
class TestPurge:
247318
async def test_purge_removes_row(
@@ -262,6 +333,16 @@ async def test_purge_missing_raises(self, repository: AppConfigFragmentRepositor
262333
Purger(row_class=AppConfigFragmentRow, pk_value=AppConfigFragmentID(uuid.uuid4()))
263334
)
264335

336+
async def test_purge_rejected_when_not_allow_listed(
337+
self,
338+
repository: AppConfigFragmentRepository,
339+
fragment_not_allow_listed: AppConfigFragmentData,
340+
) -> None:
341+
with pytest.raises(AppConfigFragmentWriteNotAllowed):
342+
await repository.purge(
343+
Purger(row_class=AppConfigFragmentRow, pk_value=fragment_not_allow_listed.id)
344+
)
345+
265346

266347
class TestSearch:
267348
async def test_search_returns_all_and_paginates(

0 commit comments

Comments
 (0)