Skip to content

Commit 306f77d

Browse files
fregataaclaude
andcommitted
refactor(BA-5071): Apply RBAC Creator pattern to ImageAlias
- Add IMAGE_ALIAS to RBACElementType enum - Replace Creator/execute_creator with RBACEntityCreator/execute_rbac_entity_creator for ImageAlias creation in image domain - Scope association: ImageAlias → IMAGE (parent scope) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 7d1ad28 commit 306f77d

5 files changed

Lines changed: 56 additions & 18 deletions

File tree

src/ai/backend/common/data/permission/types.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -385,6 +385,9 @@ class RBACElementType(enum.StrEnum):
385385
# === Auto-only entities used in permissions ===
386386
NOTIFICATION_RULE = "notification_rule"
387387

388+
# === Auto sub-entities with direct GET APIs ===
389+
IMAGE_ALIAS = "image:alias"
390+
388391
# === Entity-level scopes (for entity-scope permissions) ===
389392
ARTIFACT_REVISION = "artifact_revision"
390393

src/ai/backend/manager/repositories/image/db_source/db_source.py

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from sqlalchemy.orm import selectinload
1313

1414
from ai.backend.common.bgtask.reporter import ProgressReporter
15+
from ai.backend.common.data.permission.types import RBACElementType
1516
from ai.backend.common.docker import ImageRef
1617
from ai.backend.common.exception import UnknownImageReference
1718
from ai.backend.common.types import ImageAlias, ImageID
@@ -26,6 +27,7 @@
2627
RescanImagesResult,
2728
ResourceLimitInput,
2829
)
30+
from ai.backend.manager.data.permission.types import RBACElementRef
2931
from ai.backend.manager.errors.image import (
3032
AliasImageActionDBError,
3133
AliasImageActionValueError,
@@ -45,8 +47,11 @@
4547
)
4648
from ai.backend.manager.models.kernel.row import KernelRow
4749
from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
48-
from ai.backend.manager.repositories.base import BatchQuerier, Creator, execute_batch_querier
49-
from ai.backend.manager.repositories.base.creator import execute_creator
50+
from ai.backend.manager.repositories.base import BatchQuerier, execute_batch_querier
51+
from ai.backend.manager.repositories.base.rbac.entity_creator import (
52+
RBACEntityCreator,
53+
execute_rbac_entity_creator,
54+
)
5055
from ai.backend.manager.repositories.base.updater import Updater, execute_updater
5156
from ai.backend.manager.repositories.image.creators import ImageAliasCreatorSpec
5257

@@ -244,10 +249,20 @@ async def insert_image_alias(
244249
image_row = await ImageRow.resolve(
245250
session, [ImageIdentifier(image_canonical, architecture)]
246251
)
247-
image_alias = ImageAliasRow(alias=alias, image_id=image_row.id)
248-
image_row.aliases.append(image_alias)
252+
rbac_creator = RBACEntityCreator(
253+
spec=ImageAliasCreatorSpec(
254+
alias=alias,
255+
image_id=image_row.id,
256+
),
257+
element_type=RBACElementType.IMAGE_ALIAS,
258+
scope_ref=RBACElementRef(
259+
element_type=RBACElementType.IMAGE,
260+
element_id=str(image_row.id),
261+
),
262+
)
263+
result = await execute_rbac_entity_creator(session, rbac_creator)
249264
row_id = image_row.id
250-
alias_data = ImageAliasData(id=image_alias.id, alias=image_alias.alias or "")
265+
alias_data = ImageAliasData(id=result.row.id, alias=result.row.alias or "")
251266
return row_id, alias_data
252267
except ValueError as e:
253268
raise AliasImageActionValueError from e
@@ -327,16 +342,18 @@ async def clear_image_resource_limits(
327342
image_row._resources = {}
328343
return image_row.to_dataclass()
329344

330-
async def insert_image_alias_by_id(self, creator: Creator[ImageAliasRow]) -> ImageAliasData:
345+
async def insert_image_alias_by_id(
346+
self, creator: RBACEntityCreator[ImageAliasRow]
347+
) -> ImageAliasData:
331348
"""
332-
Creates an image alias using the Creator pattern.
349+
Creates an image alias using the RBACEntityCreator pattern.
333350
"""
334-
spec = cast(ImageAliasCreatorSpec, creator.spec)
335351
try:
336352
async with self._db.begin_session() as session:
353+
spec = cast(ImageAliasCreatorSpec, creator.spec)
337354
# Validate that the image exists
338355
await self._get_image_by_id(session, spec.image_id)
339-
result = await execute_creator(session, creator)
356+
result = await execute_rbac_entity_creator(session, creator)
340357
return ImageAliasData(id=result.row.id, alias=result.row.alias or "")
341358
except ValueError as e:
342359
raise AliasImageActionValueError from e

src/ai/backend/manager/repositories/image/repository.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@
3232
)
3333
from ai.backend.manager.models.image.row import ImageAliasRow
3434
from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
35-
from ai.backend.manager.repositories.base import BatchQuerier, Creator
35+
from ai.backend.manager.repositories.base import BatchQuerier
36+
from ai.backend.manager.repositories.base.rbac.entity_creator import RBACEntityCreator
3637
from ai.backend.manager.repositories.base.updater import Updater
3738
from ai.backend.manager.repositories.image.db_source.db_source import ImageDBSource
3839
from ai.backend.manager.repositories.image.stateful_source.stateful_source import (
@@ -281,9 +282,11 @@ async def clear_image_custom_resource_limit(
281282
return await self._db_source.clear_image_resource_limits(image_canonical, architecture)
282283

283284
@image_repository_resilience.apply()
284-
async def add_image_alias_by_id(self, creator: Creator[ImageAliasRow]) -> ImageAliasData:
285+
async def add_image_alias_by_id(
286+
self, creator: RBACEntityCreator[ImageAliasRow]
287+
) -> ImageAliasData:
285288
"""
286-
Creates an image alias using the Creator pattern.
289+
Creates an image alias using the RBACEntityCreator pattern.
287290
"""
288291
return await self._db_source.insert_image_alias_by_id(creator)
289292

src/ai/backend/manager/services/image/service.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,21 +2,23 @@
22
from uuid import UUID
33

44
from ai.backend.common.contexts.user import current_user
5+
from ai.backend.common.data.permission.types import RBACElementType
56
from ai.backend.common.docker import ImageRef
67
from ai.backend.common.dto.manager.rpc_request import PurgeImagesReq
78
from ai.backend.common.exception import UnknownImageReference
89
from ai.backend.common.types import AgentId, ImageAlias, ImageID
910
from ai.backend.logging.utils import BraceStyleAdapter
1011
from ai.backend.manager.config.provider import ManagerConfigProvider
1112
from ai.backend.manager.data.image.types import ImageWithAgentInstallStatus
13+
from ai.backend.manager.data.permission.types import RBACElementRef
1214
from ai.backend.manager.errors.image import ImageAccessForbiddenError, ImageNotFound
1315
from ai.backend.manager.models.image import (
1416
ImageIdentifier,
1517
ImageRow,
1618
)
1719
from ai.backend.manager.models.user import UserRole
1820
from ai.backend.manager.registry import AgentRegistry
19-
from ai.backend.manager.repositories.base import Creator
21+
from ai.backend.manager.repositories.base.rbac.entity_creator import RBACEntityCreator
2022
from ai.backend.manager.repositories.base.updater import Updater
2123
from ai.backend.manager.repositories.image.creators import ImageAliasCreatorSpec
2224
from ai.backend.manager.repositories.image.repository import ImageRepository
@@ -405,13 +407,18 @@ async def alias_image_by_id(self, action: AliasImageByIdAction) -> AliasImageByI
405407
"""
406408
Creates an alias for an image by its ID.
407409
"""
408-
creator = Creator(
410+
rbac_creator = RBACEntityCreator(
409411
spec=ImageAliasCreatorSpec(
410412
alias=action.alias,
411413
image_id=action.image_id,
412-
)
414+
),
415+
element_type=RBACElementType.IMAGE_ALIAS,
416+
scope_ref=RBACElementRef(
417+
element_type=RBACElementType.IMAGE,
418+
element_id=str(action.image_id),
419+
),
413420
)
414-
image_alias = await self._image_repository.add_image_alias_by_id(creator)
421+
image_alias = await self._image_repository.add_image_alias_by_id(rbac_creator)
415422
return AliasImageByIdActionResult(
416423
image_id=action.image_id,
417424
image_alias=image_alias,

tests/unit/manager/services/image/test_image_service.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
from ai.backend.common.container_registry import ContainerRegistryType
1818
from ai.backend.common.contexts.user import with_user
19+
from ai.backend.common.data.permission.types import RBACElementType
1920
from ai.backend.common.data.user.types import UserData
2021
from ai.backend.common.dto.agent.response import PurgeImageResp, PurgeImagesResp
2122
from ai.backend.common.exception import UnknownImageReference
@@ -31,14 +32,16 @@
3132
RescanImagesResult,
3233
ResourceLimitInput,
3334
)
35+
from ai.backend.manager.data.permission.types import RBACElementRef
3436
from ai.backend.manager.errors.image import (
3537
ImageAccessForbiddenError,
3638
ImageAliasNotFound,
3739
ImageNotFound,
3840
)
3941
from ai.backend.manager.models.image import ImageStatus, ImageType
4042
from ai.backend.manager.models.user import UserRole
41-
from ai.backend.manager.repositories.base import BatchQuerier, Creator, OffsetPagination
43+
from ai.backend.manager.repositories.base import BatchQuerier, OffsetPagination
44+
from ai.backend.manager.repositories.base.rbac.entity_creator import RBACEntityCreator
4245
from ai.backend.manager.repositories.image.creators import ImageAliasCreatorSpec
4346
from ai.backend.manager.repositories.image.repository import ImageRepository
4447
from ai.backend.manager.repositories.image.updaters import ImageUpdaterSpec
@@ -992,10 +995,15 @@ async def test_alias_image_by_id_success(
992995
assert result.image_alias == image_alias_data
993996
mock_image_repository.add_image_alias_by_id.assert_called_once()
994997
creator_arg = mock_image_repository.add_image_alias_by_id.call_args[0][0]
995-
assert isinstance(creator_arg, Creator)
998+
assert isinstance(creator_arg, RBACEntityCreator)
996999
assert isinstance(creator_arg.spec, ImageAliasCreatorSpec)
9971000
assert creator_arg.spec.alias == "python"
9981001
assert creator_arg.spec.image_id == image_id
1002+
assert creator_arg.element_type == RBACElementType.IMAGE_ALIAS
1003+
assert creator_arg.scope_ref == RBACElementRef(
1004+
element_type=RBACElementType.IMAGE,
1005+
element_id=str(image_id),
1006+
)
9991007

10001008

10011009
class TestClearImageCustomResourceLimitById(ImageServiceBaseFixtures):

0 commit comments

Comments
 (0)