|
12 | 12 | from sqlalchemy.orm import selectinload |
13 | 13 |
|
14 | 14 | from ai.backend.common.bgtask.reporter import ProgressReporter |
| 15 | +from ai.backend.common.data.permission.types import RBACElementType |
15 | 16 | from ai.backend.common.docker import ImageRef |
16 | 17 | from ai.backend.common.exception import UnknownImageReference |
17 | 18 | from ai.backend.common.types import ImageAlias, ImageID |
|
26 | 27 | RescanImagesResult, |
27 | 28 | ResourceLimitInput, |
28 | 29 | ) |
| 30 | +from ai.backend.manager.data.permission.types import RBACElementRef |
29 | 31 | from ai.backend.manager.errors.image import ( |
30 | 32 | AliasImageActionDBError, |
31 | 33 | AliasImageActionValueError, |
|
45 | 47 | ) |
46 | 48 | from ai.backend.manager.models.kernel.row import KernelRow |
47 | 49 | 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 | +) |
50 | 55 | from ai.backend.manager.repositories.base.updater import Updater, execute_updater |
51 | 56 | from ai.backend.manager.repositories.image.creators import ImageAliasCreatorSpec |
52 | 57 |
|
@@ -244,10 +249,20 @@ async def insert_image_alias( |
244 | 249 | image_row = await ImageRow.resolve( |
245 | 250 | session, [ImageIdentifier(image_canonical, architecture)] |
246 | 251 | ) |
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) |
249 | 264 | 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 "") |
251 | 266 | return row_id, alias_data |
252 | 267 | except ValueError as e: |
253 | 268 | raise AliasImageActionValueError from e |
@@ -327,16 +342,18 @@ async def clear_image_resource_limits( |
327 | 342 | image_row._resources = {} |
328 | 343 | return image_row.to_dataclass() |
329 | 344 |
|
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: |
331 | 348 | """ |
332 | | - Creates an image alias using the Creator pattern. |
| 349 | + Creates an image alias using the RBACEntityCreator pattern. |
333 | 350 | """ |
334 | | - spec = cast(ImageAliasCreatorSpec, creator.spec) |
335 | 351 | try: |
336 | 352 | async with self._db.begin_session() as session: |
| 353 | + spec = cast(ImageAliasCreatorSpec, creator.spec) |
337 | 354 | # Validate that the image exists |
338 | 355 | 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) |
340 | 357 | return ImageAliasData(id=result.row.id, alias=result.row.alias or "") |
341 | 358 | except ValueError as e: |
342 | 359 | raise AliasImageActionValueError from e |
|
0 commit comments